import os
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import ogstools as ot
import vtuIO

# Import analytical solution from a CSV file
exact = np.loadtxt("data/ref_solution_saturation.csv", delimiter=",")
# Zeroth column is location, first column is saturation

out_dir = Path(os.environ.get("OGS_TESTRUNNER_OUT_DIR", "_out"))
if not out_dir.exists():
    out_dir.mkdir(parents=True)

# run OGS
model = ot.Project(input_file="mcWhorter_h2.prj", output_file="mcWhorter_h2.prj")
model.run_model(logfile=f"{out_dir}/out.txt", args=f"-o {out_dir}")

# read OGS results from PVD file
pvdfile = vtuIO.PVDIO(
    f"{out_dir}/result_McWhorter_H2.pvd", dim=2, interpolation_backend="vtk"
)

# The sampling routine requires a line of points in space
x_axis = [(i, 0, 0) for i in exact[:, 0]]

# Only the last timestep is written (together with the initial condition).
# Thus the second element of the time vector (index = 1) is the to be sampled
time = pvdfile.timesteps[1]

# The numerical solution is sampled at the same supporting points as the analytical solution
sL_num = pvdfile.read_set_data(time, "saturation", pointsetarray=x_axis)

# Absolute and relative errors
err_abs = exact[:, 1] - sL_num
err_rel = err_abs / exact[:, 1]

plt.rcParams["figure.figsize"] = (14, 4)
fig1, (ax1, ax2) = plt.subplots(1, 2)

fig1.suptitle(r"Liquid saturation and errors at t=" + str(time) + " seconds")

# Saturation vs. time
ax1.set_xlabel("$x$ / m", fontsize=12)
ax1.set_ylabel(r"$s_\mathrm{L}$", fontsize=12)

ax1.plot(exact[:, 0], sL_num, "b", label=r"$s_\mathrm{L}$ numerical")
ax1.plot(exact[:, 0], exact[:, 1], "g", label=r"$s_\mathrm{L}$ exact")

lns2 = ax2.plot(
    exact[:, 0], err_abs, "b", linewidth=2, linestyle="-", label=r"absolute error"
)

ax2.set_xlabel("$x$ / m", fontsize=12)
ax2.set_ylabel(r"$\epsilon_\mathrm{abs}$", fontsize=12)

ax3 = ax2.twinx()
lns3 = ax3.plot(
    exact[:, 0], err_rel, "g", linewidth=2, linestyle="-", label=r"relative error"
)
ax3.set_ylabel(r"$\epsilon_\mathrm{rel}$", fontsize=12)

ax1.legend()

lns = lns2 + lns3
labs = [label.get_label() for label in lns]
ax2.legend(lns, labs, loc=0)


fig1.savefig(f"{out_dir}/mcWhorter.pdf")