# Standard library imports
import os
import xml.etree.ElementTree as ET
from pathlib import Path

# Local imports
import ogstools as ot

# Third-party imports
import pyvista as pv

out_dir = Path(os.environ.get("OGS_TESTRUNNER_OUT_DIR", "_out"))

if not out_dir.exists():
    out_dir.mkdir(parents=True)

def get_last_vtu_file_name(pvd_file_name):
    tree = ET.parse(pvd_file_name)
    root = tree.getroot()
    # Get the last DataSet tag
    last_dataset = root.findall(".//DataSet")[-1]

    # Get the 'file' attribute of the last DataSet tag
    file_attribute = last_dataset.attrib["file"]
    return Path(out_dir, file_attribute)

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

    def __init__(
        self, project_file, output_prefix, mesh_path, out_dir=out_dir, use_fbar=True
    ):
        self.prj = ot.Project(
            input_file=project_file, output_file=Path(out_dir, "modified.prj")
        )

        self.prj.replace_text(output_prefix, xpath="./time_loop/output/prefix")
        if not use_fbar:
            self.prj.replace_text("none", xpath="./processes/process/f_bar")

        self.use_fbar = use_fbar
        self.out_dir = out_dir
        self.output_prefix = output_prefix
        self.pvd_file_name = Path(self.out_dir, self.output_prefix + ".pvd")
        self.meth_path = mesh_path
        self.resulted_mesh = pv.UnstructuredGrid()

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

    def run(self):
        self.prj.write_input()

        self.prj.run_model(
            logfile=Path(self.out_dir, "out.txt"),
            args=f"-o {self.out_dir} -m {self.meth_path}",
        )
        self.resulted_mesh = pv.read(get_last_vtu_file_name(self.pvd_file_name))

    def get_pvd_name(self):
        return self.pvd_file_name

    def plot3D(self, variable_name, camera_position, component_id=-1):
        point_data = self.resulted_mesh.point_data[variable_name]
        plot_data = point_data[:, component_id] if component_id > 0 else point_data
        deformed_mesh = self.resulted_mesh.copy()
        dim = self.resulted_mesh.point_data["displacement"].shape[1]
        deformed_mesh.points[:, :dim] += self.resulted_mesh.point_data["displacement"]

        pl = pv.Plotter()
        pl.add_mesh(self.resulted_mesh, color="white", style="wireframe", line_width=1)
        pl.add_mesh(
            deformed_mesh,
            scalars=plot_data,
            label="Vertical displacement",
            show_edges=False,
            cmap="jet",
        )
        pl.show_axes()
        pl.camera_position = camera_position
        pl.reset_camera()

        pl.show()

    def check_vertical_displacement_at_point(
        self, point, expected_value, component_id=2
    ):
        p_id = self.resulted_mesh.find_closest_point(point)
        u = self.resulted_mesh.point_data["displacement"][p_id]
        u_z_max = u[component_id]
        print(u_z_max)
        try:
            rel_error = abs((u_z_max - expected_value) / u_z_max)
            print(f"The relative error is {rel_error}.")
            assert rel_error < 1e-4, (
                f"The caculated vertical_displacement at {point} is not"
                f" {expected_value}"
            )
            return u_z_max
        except Exception as e:
            print(f"An error occurred: {e}")

    def run_benchmark(
        self,
        point_having_u_compoment_max,
        expected_u_compoment_max,
        camera_position,
        benchmark_name,
        component_id,
    ):
        self.run()

        u_compoment_max = self.check_vertical_displacement_at_point(
            point_having_u_compoment_max, expected_u_compoment_max, component_id
        )
        f_bar_info = "With the F-bar" if self.use_fbar else "Without the F-bar"
        print(
            f"{f_bar_info}, the obtain maximum displacement component"
            f" of {benchmark_name} is {u_compoment_max} m."
        )

        self.plot3D("displacement", camera_position, component_id)

try:
    project_file = Path("RubberIndentation", "RubberIndentation.prj")
    output_prefix = "RubberIndentation"
    ogs_model = SingleOGSModel(
        project_file,
        output_prefix,
        mesh_path=output_prefix,
        out_dir=out_dir,
        use_fbar=False,
    )
    # ogs_model.reset_time_step_size(0.3, 4)

    camera_position = "xy"
    point_have_u_compoment_max = [0.0, 0.01, 0.0]
    expected_u_compoment_max = -0.00010428732840299523
    ogs_model.run_benchmark(
        point_have_u_compoment_max,
        expected_u_compoment_max,
        camera_position,
        "Rubber indentation",
        component_id=1,
    )
except Exception:
    pass

try:
    project_file = Path("RubberIndentation", "RubberIndentation.prj")
    output_prefix = "RubberIndentation"
    ogs_model = SingleOGSModel(
        project_file,
        output_prefix,
        mesh_path=output_prefix,
        out_dir=out_dir,
        use_fbar=True,
    )

    camera_position = "xy"
    point_have_u_compoment_max = [0.0, 0.01, 0.0]
    expected_u_compoment_max = -0.00026707547768988984
    ogs_model.run_benchmark(
        point_have_u_compoment_max,
        expected_u_compoment_max,
        camera_position,
        "Rubber indentation",
        component_id=1,
    )
except Exception:
    pass

try:
    project_file = Path("Indentation3D", "Indentation3D.prj")
    output_prefix = "Indentation3D"
    ogs_model = SingleOGSModel(
        project_file,
        output_prefix,
        mesh_path=output_prefix,
        out_dir=out_dir,
        use_fbar=False,
    )

    camera_position = (-0.5, -0.5, 0.4)
    point_have_u_compoment_max = [0.0, 0.0, 0.01]
    expected_u_compoment_max = -1.6534484343296282e-05
    ogs_model.run_benchmark(
        point_have_u_compoment_max,
        expected_u_compoment_max,
        camera_position,
        "3D indentation",
        component_id=2,
    )
except Exception:
    pass

try:
    project_file = Path("Indentation3D", "Indentation3D.prj")
    output_prefix = "Indentation3D"
    ogs_model = SingleOGSModel(
        project_file,
        output_prefix,
        mesh_path=output_prefix,
        out_dir=out_dir,
        use_fbar=True,
    )

    camera_position = (-0.5, -0.5, 0.4)
    point_have_u_compoment_max = [0.0, 0.0, 0.01]
    expected_u_compoment_max = -0.00010953948059235105
    ogs_model.run_benchmark(
        point_have_u_compoment_max,
        expected_u_compoment_max,
        camera_position,
        "3D indentation",
        component_id=2,
    )
except Exception:
    pass

try:
    project_file = Path("ThickCylindricalShell", "ThickCylindricalShell.prj")
    output_prefix = "ThickCylindricalShell"
    ogs_model = SingleOGSModel(
        project_file,
        output_prefix,
        mesh_path=output_prefix,
        out_dir=out_dir,
        use_fbar=False,
    )
    ogs_model.reset_time_step_size(0.3, 4)

    camera_position = (0.5, -0.5, 0.4)
    point_have_u_compoment_max = [0.0, 0.01, 0.0075]
    expected_u_compoment_max = -0.010194869353071633
    ogs_model.run_benchmark(
        point_have_u_compoment_max,
        expected_u_compoment_max,
        camera_position,
        "Thick cylindrical shell",
        component_id=1,
    )
except Exception:
    pass

try:
    project_file = Path("ThickCylindricalShell", "ThickCylindricalShell.prj")
    output_prefix = "ThickCylindricalShell"
    ogs_model = SingleOGSModel(
        project_file,
        output_prefix,
        mesh_path=output_prefix,
        out_dir=out_dir,
        use_fbar=True,
    )

    camera_position = (0.5, -0.5, 0.4)
    point_have_u_compoment_max = [0.0, 0.01, 0.0075]
    expected_u_compoment_max = -0.014236570288857237
    ogs_model.run_benchmark(
        point_have_u_compoment_max,
        expected_u_compoment_max,
        camera_position,
        "Thick cylindrical shell",
        component_id=1,
    )
except Exception:
    pass