from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt


def main():
    timesteps = [500,1000,5000,10000]

    diamond_dir = Path(__file__).resolve().parent.parent
    line_probe_dir = diamond_dir / "line_probe"

    plt.figure(figsize=(7.5, 4.5))

    for nt in timesteps:
        f = line_probe_dir / f"000.Cart.{nt:09d}.line"
        data = np.loadtxt(f, comments="#")
        x = data[:, 2]
        p = data[:, 4]
        plt.plot(x, p, lw=1.8, label=f"{nt}")

    plt.xlabel("x")
    plt.ylabel("Pressure (Pa)")
    plt.grid(True, alpha=0.25)
    plt.legend(title="timestep", fontsize=9)

    out_png = Path(__file__).resolve().parent / "line_probe_pressure.png"
    plt.tight_layout()
    plt.savefig(out_png, dpi=200)
    print(f"Wrote: {out_png}")


if __name__ == "__main__":
    main()