import numpy as np
import matplotlib.pyplot as plt
from sys import argv
import os
import re
from scipy.optimize import curve_fit
import matplotlib as mpl

import mpl_colors


def get_name(L, n_steps, eps, chi_max):
    return f"L={L:02}_nSteps={n_steps:04}_eps={eps:e}_chiMax={chi_max:02}"

def power2_(x, a, b):
    return a * 2**x + b

def power2(x, a, b, c):
    return a * 2**(x*b) + c

def powerX(x, a, b, c, d):
    return a * b**(x*c) + d

def power2squared(x, a, b, c):
    return a * 2**(x**2 *b) + c

def exp(x, a, b):
    return a * np.exp(b * x)

def total4D(x, a, b, c, d, e):
    L, n, eps, chi = np.split(x, 4, axis=1)
    return (a * n * e**(L*d) * (1 + b * eps) * 2**chi + c).flatten()
    # fit_and_plot(Ldata, 0, powerX, axL)
    # fit_and_plot(Ndata, 1, lambda x, a, b: a * x + b, axN)
    # fit_and_plot(Epsdata, 2, lambda x, a: x * 0 + a, axEps)
    # fit_and_plot(Chidata, 3, power2, axChi)


def fit(f, x, y):
    popt, pcov = curve_fit(f, x, y)
    yfit = f(x, *popt)
    return yfit, popt, pcov

platform_colors = {
    "firefox": "o",
    "chromium": "b",
    "native": "g",
}


def process_benchmarks_1AT(benchmarks):
    L_default = 12
    n_steps_default = 200
    eps_default = 1e-10
    chi_max_default = 20

    get_benchmarks_lambdas = [
        lambda a: a[(a[:,1] == n_steps_default) & (a[:,2] == eps_default) & (a[:,3] == chi_max_default)],
        lambda a: a[(a[:,0] == L_default) & (a[:,2] == eps_default) & (a[:,3] == chi_max_default)],
        lambda a: a[(a[:,0] == L_default) & (a[:,1] == n_steps_default) & (a[:,3] == chi_max_default)],
        lambda a: a[(a[:,0] == L_default) & (a[:,1] == n_steps_default) & (a[:,2] == eps_default)],
    ]
    fit_fs = [
        # powerX,
        power2,
        lambda x, a, b: a * x + b,
        lambda x, a: x * 0 + a,
        power2,
    ]


    params = ["L", "N", "Eps", "Chi"]
    N_params = len(params)
    xlabels = ["$L$", r"$n_\text{steps}$", r"$\epsilon$", r"$\chi_\text{max}$"]
    colors = ["b", "g", "o", "r"]
    fig_ax_platform = [plt.subplots(figsize=(4,3)) for _ in range(N_params)]
    linestyles = {"desktop": "solid", "laptop": "dashed"}
    markers = {"desktop": "o", "laptop": "x"}
    for device, platforms in benchmarks.items():
        linestyle = linestyles[device]
        marker = markers[device]
        for platform, opt_levels in platforms.items():
            fig_ax = [plt.subplots() for _ in range(N_params)]

            for opt, data in sorted(opt_levels.items()):
                score = int(np.sum(data[:,4], axis=0))
                print(f"{platform} O{opt}: {score}")

                for i in range(N_params):
                    c = colors[opt]
                    d = get_benchmarks_lambdas[i](data)
                    x = d[:,i]
                    y = d[:,4]
                    y_fit, popt, pcov = fit(fit_fs[i], x, y)
                    print(f"{platform:10} O{opt} {params[i]:3}: {popt}")
                    ax = fig_ax[i][1]
                    ax.scatter(x, y, color=c, marker=marker)
                    ax.plot(x, y_fit, color=c, label=f"O{opt}", linestyle=linestyle)
                    ax.set_xlabel(xlabels[i])
                    if params[i] in ["Eps"]: ax.set_xscale("log")

                    if opt == 2:  # platform comparison
                        c = platform_colors[platform]
                        ax = fig_ax_platform[i][1]
                        ax.scatter(x, y, color=c, gid=f"{platform}", marker=marker)
                        ax.plot(x, y_fit, color=c, label=f"{platform} {device}", gid=f"{platform}", linestyle=linestyle)
                        ax.set_xlabel(xlabels[i])
                        if params[i] in ["Eps"]: ax.set_xscale("log")
                        # if params[i] in ["L"]: ax.set_yscale("log")


            for i, (fig, ax) in enumerate(fig_ax):
                ax.legend()
                ax.set_ylabel("$t$ [s]")
                ax.set_title(f"{xlabels[i]} Benchmark")
                fig.tight_layout()
                # fig.savefig(f"plt_{platform}_{params[i]}.svg")

    from matplotlib.lines import Line2D
    legend_elements = []
    for platform in platform_colors.keys():
       legend_elements.append(Line2D([0], [0], linestyle="", marker='o', color=platform_colors[platform], label=platform, markersize=12))
    for device in benchmarks.keys():
       legend_elements.append(Line2D([0], [0], linestyle=linestyles[device], color=mpl_colors.P["fg0"], label=device))

    for i, (fig, ax) in enumerate(fig_ax_platform):
        # ax.legend()
        ax.legend(handles=legend_elements)
        ax.set_ylabel("$t$ [s]")
        ax.set_title(f"{xlabels[i]} Benchmark")
        fig.tight_layout()
        fig.savefig(f"plt_platforms_{params[i]}.svg")
    plt.close('all')


    def process_benchmarks_all(benchmarks):
        for platform, opt_levels in benchmarks.items():
            figL, axL = plt.subplots()
            figN, axN = plt.subplots()
            figEps, axEps = plt.subplots()
            figChi, axChi = plt.subplots()
            axs = [axL, axN, axEps, axChi]


            colors = ["blue", "green", "orange", "red"]
            for opt, data in sorted(opt_levels.items()):
                c = colors[opt]
                score = int(np.sum(data[:,4], axis=0))
                print(f"{platform} O{opt}: {score}")
                # fit
                x = data[:,0:4]
                y = data[:,4].flatten()
                y_fit, popt, pcov = fit(total4D, x, y)
                print(popt, pcov)
                # for plots with only one param varied
                L, n, eps, chi = np.split(x, 4, axis=1)
                fixed_vals = [np.max(L), np.max(n), np.min(eps), np.max(chi)]
                for i in range(4):
                    ax = axs[i]
                    # selected data
                    sel_data = data
                    sel_y_fit = y_fit
                    for j in range(4):
                        if i == j: continue
                        selection = np.where(sel_data[:,j] == fixed_vals[j])
                        sel_data = sel_data[selection]
                        sel_y_fit = sel_y_fit[selection]
                    sel_x = sel_data[:,i]
                    sel_y = sel_data[:,4].flatten()
                    ax.scatter(sel_x, sel_y, color=c)
                    ax.plot(sel_x, sel_y_fit, color=c, label=f"O{opt}")

            for ax in axs:
                ax.legend()
                ax.set_ylabel("$t$ [s]")
                ax.set_title(f"{platform}")
                # ax.set_yscale("log")
            axL.set_xlabel("$L$")
            axN.set_xlabel(r"$n_\text{steps}$")
            axEps.set_xlabel(r"$\epsilon$")
            axEps.set_xscale("log")
            axChi.set_xlabel(r"$\chi_\text{max}$")
        plt.show()

def get_benchmarks(filename_begin):
    files = [f for f in os.listdir('.') if f.startswith(filename_begin) and f.endswith('.csv')]

    if len(files) == 0:
        print("No benchmarks found")
        exit(1)

    benchmarks = {}
    # read the files
    for f in files:
        match = re.fullmatch(filename_begin + r"-?(\w+)-(\w+)(?:-\d+)?-O(\d).csv", f)
        if not match:
            print("Failed to parse benchmark: " + f)
            continue
        device = match.group(1)
        platform = match.group(2)
        opt_level = int(match.group(3))
        if not device in benchmarks: benchmarks[device] = {}
        if not platform in benchmarks[device]: benchmarks[device][platform] = {}
        benchmarks[device][platform][opt_level] = []
        with open(f, 'r') as file:
            for line in file:
                n, time = line.strip("\n").split(",")
                time = float(time)
                match = re.fullmatch(r"L=(\d+)_nSteps=(\d+)_eps=(.+)_chiMax=(\d+)", n.replace("\"", ""))
                if not match:
                    print("Failed to parse benchmark: " + n)
                    exit(1)
                L = int(match.group(1))
                n_steps = int(match.group(2))
                eps = float(match.group(3))
                chi_max = int(match.group(4))
                benchmarks[device][platform][opt_level].append([L, n_steps, eps, chi_max, time])
        benchmarks[device][platform][opt_level] = np.array(benchmarks[device][platform][opt_level])
    return benchmarks



def repl_colors():
    """
    In all svgs in img_dir, replace all occurances of hex colors defined in mpl_colors.P with
    var(--<color-name>)
    so that when the svg is inlined into a html it can be fully (and dynamically!) styled with css
    """
    img_dir = "."

    hex_to_name = {v.lower(): k for k, v in mpl_colors.P.items()}
    hex_color_pattern = re.compile(r"#([0-9a-fA-F]{6})")

    for filename in os.listdir(img_dir):
        if not filename.endswith(".svg"): continue
        print(f"Replacing colors in {filename}")
        file_path = os.path.join(img_dir, filename)
        with open(file_path, "r") as file:
            content = file.read()
        hex_colors = hex_color_pattern.findall(content)
        modified_content = content
        for hex_color in hex_colors:
            hex_code = f"#{hex_color.lower()}"

            # Check if the hex code is in our map
            if hex_code in hex_to_name:
                color_name = hex_to_name[hex_code]
                css_name = f"var(--{color_name})"
                modified_content = modified_content.replace(hex_code, css_name)
            else:
                # Print a warning if the color is not defined
                print(f"Warning: Color {hex_code} in {filename} is not defined in the color map.")

        # Write the modified content back to the SVG file
        with open(file_path, "w", encoding="utf-8") as file:
            file.write(modified_content)


if __name__ == '__main__':
    # correct fonts
    mpl.rcParams["mathtext.fontset"] = "stix"
    mpl.rcParams["font.family"] = "Ubuntu"
    # set replaceable colorscheme
    mpl_colors.set_colorscheme("dark")
    benchmarks1 =get_benchmarks("benchmarks-1AT")
    process_benchmarks_1AT(benchmarks1)
    repl_colors()

    # benchmarks2 =get_benchmarks("benchmarks-all")
    # process_benchmarks_all(benchmarks2)
