import os
from pathlib import Path

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)

import xml.etree.ElementTree as ET

import pyvista as pv

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)

print(uys_at_top_non_bbar)

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)

print(uys_at_top_bbar)

import matplotlib.pyplot as plt
import numpy as np

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()

plot_data(ne, uys_at_top_bbar, uys_at_top_non_bbar, "b_bar_linear.png")

import matplotlib.tri as tri
import vtuIO

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()

for nedge, output_prefix in zip(nedges, output_prefices_non_bbar):
    contour_plot(output_prefix + ".pvd", "Number of elements per side: " + nedge)

for nedge, output_prefix in zip(nedges, output_prefices):
    contour_plot(output_prefix + ".pvd", "Number of elements per side: " + nedge)