from importlib.resources import files
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import argparse
import os

parser = argparse.ArgumentParser()
parser.add_argument("--filename", type=str, default="src/project1/poisson_solution.csv")
parser.add_argument("--reference", type=str, default=None)
args = parser.parse_args()

FILENAME = args.filename
REFERENCE = args.reference

def get_rc_params():
    colors = ["FF220C", "70D6FF", "8AAA79", "666370", "1C1F33"]
    rcParams = plt.rcParams
    # Use LaTeX for rendering
    # Setup fonts
    rcParams["text.usetex"] = True
    rcParams["font.family"] = "serif"
    rcParams["font.size"] = 12
    rcParams["axes.labelsize"] = 12
    rcParams["axes.titlesize"] = 12
    rcParams["legend.fontsize"] = 10
    rcParams["xtick.labelsize"] = 10
    rcParams["ytick.labelsize"] = 10
    # Figure size and resolution
    rcParams["figure.figsize"] = (6, 4)
    rcParams["figure.dpi"] = 300
    # Use colors from the palette
    rcParams["axes.prop_cycle"] = plt.cycler(color=[f"#{color}" for color in colors])
    # Grid
    rcParams["axes.grid"] = True
    rcParams["grid.alpha"] = 0.5
    rcParams["grid.linestyle"] = "--"
    # Point ticks to the inside of the axes
    rcParams["xtick.direction"] = "in"
    rcParams["ytick.direction"] = "in"
    rcParams["xtick.top"] = True
    rcParams["ytick.right"] = True
    return rcParams


def singular_plot():
    plt.rcParams.update(get_rc_params())
    df = pd.read_csv(FILENAME)
    fig, ax = plt.subplots(figsize=(6,3))
    ax.plot(df["x"], df[df.columns[1]])
    ax.set_xlabel(r"$x$")
    ax.set_ylabel(rf"${df.columns[1]}$")

    fig.tight_layout()

    output_directory = os.path.abspath(os.path.join(__file__, "../../../../projects/project1/include"))
    output_file = os.path.basename(FILENAME).replace(".csv", ".pdf")
    fig.savefig(os.path.join(output_directory, output_file))

def reference_plot():
    plt.rcParams.update(get_rc_params())
    df = pd.read_csv(FILENAME).sort_values("x")

    fig, (ax1, ax2) = plt.subplots(nrows=2, sharex=True, height_ratios=[3, 1], figsize=(4, 8))
    ax1.plot(df["x"], df["v(x)"], label="Numerical Solution")
    ax1.plot(df["x"], df["u(x)"], label="Reference Solution", linestyle="--")
    ax1.set_xlabel(r"$x$")
    ax1.set_ylabel(r"$u(x)$")
    ax1.legend()

    ratio = (df["v(x)"] - df["u(x)"]) / df["u(x)"]

    ax2.plot(df["x"], ratio.abs())
    ax2.set_xlabel(r"$x$")
    ax2.set_ylabel(r"$|\frac{v(x) - u(x)}{u(x)}|$")
    ax2.axhline(y=1, color="gray", linestyle="--")
    ax2.set_yscale("log")

    fig.tight_layout()

    output_directory = os.path.abspath(os.path.join(__file__, "../../../../projects/project1/include"))
    output_file = os.path.basename(FILENAME).replace(".csv", "_ref.pdf")
    fig.savefig(os.path.join(output_directory, output_file))

def multiple_reference_plot():
    plt.rcParams.update(get_rc_params())
    # Filename is wildcard -> Find files
    import glob
    files = glob.glob(FILENAME)
    files.sort()
    fig, (ax1, ax2) = plt.subplots(nrows=2, sharex=True, height_ratios=[3, 1], figsize=(6,5))
   

    for file in files:
        df = pd.read_csv(file).sort_values("x")
        label = f"Num. (N = {os.path.basename(file).replace(".csv", "").split("_")[-1]})"
        ax1.plot(df["x"], df["v(x)"], label=label)

        ratio = (df["v(x)"] - df["u(x)"]) / df["u(x)"]
        ax2.plot(df["x"], ratio.abs(), label=label)

    ax1.plot(df["x"], df["u(x)"], label="Reference Solution", linestyle="--", color="gray")
    ax1.set_xlabel(r"$x$")
    ax1.set_ylabel(r"$u(x)$")
    ax1.legend()
    ax2.axhline(y=1, color="gray", linestyle="--")
    ax2.set_xlabel(r"$x$")
    ax2.set_ylabel(r"$|\frac{v(x) - u(x)}{u(x)}|$")
    ax2.set_yscale("log")

    fig.tight_layout()

    output_directory = os.path.abspath(os.path.join(__file__, "../../../../projects/project1/include"))
    output_file = "big_ref_plot.pdf"
    fig.savefig(os.path.join(output_directory, output_file))


def error_plots():
    abs_error_plot()
    rel_error_plot()


def abs_error_plot():
    plt.rcParams.update(get_rc_params())
    import glob
    files = glob.glob(FILENAME)
    files.sort()

    df_ref = pd.read_csv(REFERENCE).sort_values("x")

    fig, ax = plt.subplots(figsize=(6,3))
    for file in files:
        df = pd.read_csv(file).sort_values("x")
        # combined_df = pd.merge_asof(df, df_ref, on="x", suffixes=("_num", "_ref"))
        log_abs_err = np.log10((df["u(x)"] - df["v(x)"]).abs())
        ax.plot(df["x"], log_abs_err, label=f"N={os.path.basename(file).replace('.csv','').split('_')[-1]}")

    ax.set_xlabel(r"$x$")
    ax.set_ylabel(r"$\log_{10}(|u(x) - v(x)|)$")
    ax.legend()

    fig.tight_layout()

    output_directory = os.path.abspath(os.path.join(__file__, "../../../../projects/project1/include"))
    output_file = "abs_error_plot.pdf"
    fig.savefig(os.path.join(output_directory, output_file))


def rel_error_plot():
    plt.rcParams.update(get_rc_params())
    import glob
    files = glob.glob(FILENAME)
    files.sort()

    df_ref = pd.read_csv(REFERENCE).sort_values("x")

    fig, ax = plt.subplots(figsize=(6,3))
    for file in files:
        df = pd.read_csv(file).sort_values("x")
        # combined_df = pd.merge_asof(df, df_ref, on="x", suffixes=("_num", "_ref"))
        # Drop rows where u(x)_ref is zero
        df = df[df["u(x)"] > 1e-10]
        log_rel_err = np.log10((df["u(x)"] - df["v(x)"]).abs() / df["u(x)"].abs())
        ax.plot(df["x"], log_rel_err, label=f"N={os.path.basename(file).replace('.csv','').split('_')[-1]}")

    ax.set_xlabel(r"$x$")
    ax.set_ylabel(r"$\log_{10}(|\frac{u(x) - v(x)}{u(x)}|)$")
    ax.legend()

    fig.tight_layout()

    output_directory = os.path.abspath(os.path.join(__file__, "../../../../projects/project1/include"))
    output_file = "rel_error_plot.pdf"
    fig.savefig(os.path.join(output_directory, output_file))


if __name__ == "__main__":
    if REFERENCE:
        if "*" in FILENAME:
            multiple_reference_plot()
            abs_error_plot()
            rel_error_plot()
        else:
            reference_plot()
    else:
        singular_plot()