import os
from pathlib import Path

import numpy as np
import ogstools as ot
import pyvista as pv

pv.set_jupyter_backend("static" if "CI" in os.environ else "client")

out_dir = Path(os.environ.get("OGS_TESTRUNNER_OUT_DIR", "_out"))
out_dir.mkdir(parents=True, exist_ok=True)

class SingleOGSModel:
    """An OGS run model"""

    def __init__(self, project_name: str, use_fbar: bool):
        self.prj = ot.Project(
            input_file=Path(project_name) / f"{project_name}.prj",
            output_file=out_dir / "modified.prj",
        )
        self.prj.replace_text(project_name, "./time_loop/output/prefix")
        if not use_fbar:
            self.prj.replace_text("none", "./processes/process/f_bar")
        self.project_name = project_name
        self.use_fbar = use_fbar

    def reset_time_step_size(self, dt: float, n_steps: int):
        xpath = "./time_loop/processes/process[1]/time_stepping/timesteps/pair[1]/"
        self.prj.replace_text(str(dt), xpath + "delta_t")
        self.prj.replace_text(n_steps, xpath + "repeat")

    def run(self) -> ot.MeshSeries:
        self.prj.write_input()
        self.prj.run_model(
            logfile=out_dir / "out.txt", args=f"-o {out_dir} -m {self.project_name}"
        )
        return ot.MeshSeries(out_dir / (self.project_name + ".pvd"))

    def run_benchmark(self, u_max_ref: list[float]):
        results = self.run()
        u_max_pt = np.argmax(ot.variables.displacement.magnitude.transform(results[-1]))
        u_max_num = results[-1]["displacement"][u_max_pt]
        np.testing.assert_allclose(u_max_num, u_max_ref, rtol=1e-4)
        f_bar_info = "With F-bar" if self.use_fbar else "Without F-bar"
        print(
            f"{f_bar_info}, the obtained maximum displacement"
            f" of {self.project_name} is {u_max_num} m."
        )
        self.plot3D(results[-1])

    def plot3D(self, mesh: pv.UnstructuredGrid):
        dim = np.max([cell.dimension for cell in mesh.cell])
        if dim == 2:
            zeros = np.zeros((mesh.n_points, 1))
            mesh["displacement"] = np.hstack((mesh["displacement"], zeros))
        deformed_mesh = mesh.warp_by_vector("displacement")
        pl = pv.Plotter()
        pl.camera_position = "xy" if dim == 2 else (-0.5, -0.5, 0.4)
        pl.add_mesh(mesh, color="lightgrey", style="wireframe", line_width=1)
        pl.add_mesh(deformed_mesh, scalars="displacement", show_edges=True, cmap="jet")
        pl.show_axes()
        pl.reset_camera()
        pl.show()

ogs_model = SingleOGSModel("ThickCylindricalShell", use_fbar=False)
ogs_model.reset_time_step_size(0.3, 4)
ogs_model.run_benchmark(u_max_ref=[0, -0.01192143, 0.00029815])