import os
import xml.etree.ElementTree as ET
from pathlib import Path

import matplotlib.pyplot as plt
import matplotlib.tri as tri
import numpy as np
import pyvista as pv
import vtuIO
from ogs6py.ogs import OGS

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(Path(out_dir) / 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 f"{out_dir}/" + file_attribute


def get_top_uy(pvd_file_name):
    top_point = (48.0e-3, 60.0e-3, 0)
    file_name = get_last_vtu_file_name(pvd_file_name)
    mesh = pv.read(file_name)
    p_id = mesh.find_closest_point(top_point)
    u = mesh.point_data["displacement"][p_id]
    return u[1]

def run_single_test(mesh_name, output_prefix, use_bbar="false"):
    model = OGS(INPUT_FILE="CooksMembrane.prj", PROJECT_FILE=f"{out_dir}/modified.prj")
    model.replace_text(mesh_name, xpath="./mesh")
    model.replace_text(use_bbar, xpath="./processes/process/use_b_bar")
    model.replace_text(output_prefix, xpath="./time_loop/output/prefix")
    model.replace_text(
        "BiCGSTAB", xpath="./linear_solvers/linear_solver/eigen/solver_type"
    )
    model.replace_text("ILUT", xpath="./linear_solvers/linear_solver/eigen/precon_type")
    vtu_file_name = output_prefix + "_ts_1_t_1.000000.vtu"
    model.replace_text(vtu_file_name, xpath="./test_definition/vtkdiff[1]/file")
    model.replace_text(vtu_file_name, xpath="./test_definition/vtkdiff[2]/file")
    model.replace_text(vtu_file_name, xpath="./test_definition/vtkdiff[3]/file")

    model.write_input()

    # Run OGS
    model.run_model(logfile=f"{out_dir}/out.txt", args=f"-o {out_dir} -m .")

    # Get uy at the top
    return get_top_uy(output_prefix + ".pvd")

mesh_names = [
    "mesh.vtu",
    "mesh_n10.vtu",
    "mesh_n15.vtu",
    "mesh_n20.vtu",
    "mesh_n25.vtu",
    "mesh_n30.vtu",
]
output_prefices_non_bbar = [
    "cooks_membrane_sd_edge_div_4_non_bbar",
    "cooks_membrane_sd_refined_mesh_10_non_bbar",
    "cooks_membrane_sd_refined_mesh_15_non_bbar",
    "cooks_membrane_sd_refined_mesh_20_non_bbar",
    "cooks_membrane_sd_refined_mesh_25_non_bbar",
    "cooks_membrane_sd_refined_mesh_30_non_bbar",
]

uys_at_top_non_bbar = []
for mesh_name, output_prefix in zip(mesh_names, output_prefices_non_bbar):
    uy_at_top = run_single_test(mesh_name, output_prefix)
    uys_at_top_non_bbar.append(uy_at_top)

expected_uys_at_top_non_bbar = np.array(
    [
        0.002164586784123102,
        0.0022603329644579383,
        0.002375295856067169,
        0.002519725590136146,
        0.0026515294133790837,
        0.002868289617025223,
    ]
)
np.testing.assert_allclose(
    actual=uys_at_top_non_bbar, desired=expected_uys_at_top_non_bbar, atol=1e-10
)

output_prefices = [
    "cooks_membrane_sd_edge_div_4",
    "cooks_membrane_sd_refined_mesh_10",
    "cooks_membrane_sd_refined_mesh_15",
    "cooks_membrane_sd_refined_mesh_20",
    "cooks_membrane_sd_refined_mesh_25",
    "cooks_membrane_sd_refined_mesh_30",
]

uys_at_top_bbar = []
for mesh_name, output_prefix in zip(mesh_names, output_prefices):
    uy_at_top = run_single_test(mesh_name, output_prefix, "true")
    uys_at_top_bbar.append(uy_at_top)

expected_uys_at_top_bbar = np.array(
    [
        0.0069574713856979,
        0.007772616910217863,
        0.007897597955618913,
        0.007951479575082158,
        0.007976349858390623,
        0.007999718483861992,
    ]
)

np.testing.assert_allclose(
    actual=uys_at_top_bbar, desired=expected_uys_at_top_bbar, atol=2e-4
)

ne = [4, 10, 15, 20, 25, 30]


def plot_data(ne, u_y_bbar, uy_non_bbar, file_name=""):
    # Plotting
    plt.rcParams["figure.figsize"] = [5, 5]

    if len(u_y_bbar) != 0:
        plt.plot(
            ne, np.array(u_y_bbar) * 1e3, marker="o", linestyle="dashed", label="B bar"
        )
    if len(uy_non_bbar) != 0:
        plt.plot(
            ne,
            np.array(uy_non_bbar) * 1e3,
            marker="x",
            linestyle="dashed",
            label="non B bar",
        )

    plt.xlabel("Number of elements per side")
    plt.ylabel("Top right corner displacement /mm")
    plt.legend()

    plt.tight_layout()
    if file_name != "":
        plt.savefig(file_name)
    plt.show()

nedges = ["4", "10", "15", "20", "25", "30"]


def contour_plot(pvd_file_name, title):
    file_name = get_last_vtu_file_name(pvd_file_name)
    m_plot = vtuIO.VTUIO(file_name, dim=2)
    triang = tri.Triangulation(m_plot.points[:, 0], m_plot.points[:, 1])
    triang = tri.Triangulation(m_plot.points[:, 0], m_plot.points[:, 1])
    s_plot = m_plot.get_point_field("sigma")
    s_trace = s_plot[:, 0] + s_plot[:, 1] + s_plot[:, 2]
    u_plot = m_plot.get_point_field("displacement")

    fig, ax = plt.subplots(ncols=2, figsize=(8, 3))
    ax[0].set_title(title, loc="left", y=1.12)
    plt.subplots_adjust(wspace=0.5)

    contour_stress = ax[0].tricontourf(triang, s_trace, cmap="viridis")
    contour_displacement = ax[1].tricontourf(triang, u_plot[:, 1], cmap="gist_rainbow")
    fig.colorbar(contour_stress, ax=ax[0], label="Stress trace / MPa")
    fig.colorbar(contour_displacement, ax=ax[1], label="Dispplacement / m")
    fig.tight_layout()
    plt.savefig(pvd_file_name + ".png")
    plt.show()