import os
from pathlib import Path
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from contextlib import contextmanager


fig_path = Path(os.getcwd()) / "figures"
val_path = Path(os.getcwd()) / "values"


def cm2inch(*tupl):
    inch = 2.54
    if isinstance(tupl[0], tuple):
        return tuple(i / inch for i in tupl[0])
    else:
        return tuple(i / inch for i in tupl)


def export_fig(name, fig=None):
    fig_path.mkdir(parents=True, exist_ok=True)
    if fig is None:
        fig = plt.gcf()

    fig.tight_layout()
    fig.canvas.draw()

    fig.savefig(fig_path / f"{name}.pdf")
    fig.savefig(fig_path / f"{name}.svg")
    fig.savefig(fig_path / f"{name}.pgf")

    return fig


def scientific_round(val, *err, retprec=False):
    """Scientifically rounds the values to the given errors."""
    val, err = np.asarray(val), np.asarray(err)
    if len(err.shape) == 1:
        err = np.array([err])
        err = err.T
    err = err.T

    if err.size == 1 and val.size > 1:
        err = np.ones_like(val) * err

    if len(err.shape) == 0:
        err = np.array([err])

    if val.size == 1 and err.shape[0] > 1:
        val = np.ones_like(err) * val

    i = np.floor(np.log10(err))
    first_digit = (err // 10 ** i).astype(int)
    prec = (-i + np.ones_like(err) * (first_digit <= 3)).astype(int)
    prec = np.max(prec, axis=1)

    def smart_round(value, precision):
        value = np.round(value, precision)
        if precision <= 0:
            value = value.astype(int)
        return value

    if val.size > 1:
        rounded = np.empty_like(val)
        rounded_err = np.empty_like(err)
        for n, (value, error, precision) in enumerate(zip(val, err, prec)):
            rounded[n] = smart_round(value, precision)
            rounded_err[n] = smart_round(error, precision)

        if retprec:
            return rounded, rounded_err, prec
        else:
            return rounded, rounded_err

    else:
        prec = prec[0]
        if retprec:
            return (smart_round(val, prec), *smart_round(err, prec)[0], prec)
        else:
            return (smart_round(val, prec), *smart_round(err, prec)[0])


def tex_value(val, err=None, unit=None, prefix="", suffix="", prec=0, save=None):
    """Generates LaTeX output of a value with units and error."""

    if err:
        val, err, prec = scientific_round(val, err, retprec=True)
    else:
        val = np.round(val, prec)

    if prec == 0:
        val = int(val)
        if err:
            err = int(err)

    val_string = fr"{val:.{prec}f}" if prec > 0 else str(val)
    if err:
        val_string += fr"\pm {err:.{prec}f}" if prec > 0 else str(err)

    ret_string = r"\(" + prefix

    if unit is None:
        ret_string += val_string
    else:
        ret_string += fr"\SI{{{val_string}}}{{{unit}}}"

    ret_string += suffix + r"\)"

    if save is not None:
        val_path.mkdir(parents=True, exist_ok=True)

        with open(val_path / f"{save}.tex", "w") as f:
            f.write(ret_string)

    return ret_string


###############################################################################
#                                  Plot Porn                                  #
###############################################################################


def wrap_plot(f):
    def wrapped(*args, ax=None, setup_function=plt.subplots, **kwargs):
        fig = None
        if not ax:
            fig, ax = setup_function()

        ret_val = f(*args, ax=ax, **kwargs)
        return (fig, ax, ret_val) if ret_val else (fig, ax)

    return wrapped


def get_figsize(
    columnwidth,
    wf=0.5,
    hf=(5.0 ** 0.5 - 1.0) / 2.0,
):
    """Parameters:
      - wf [float]:  width fraction in columnwidth units
      - hf [float]:  height fraction in columnwidth units.
                     Set by default to golden ratio.
      - columnwidth [float]: width of the column in latex. Get this from LaTeX
                             using \showthe\columnwidth
    Returns:  [fig_width,fig_height]: that should be given to matplotlib
    """
    fig_width_pt = columnwidth * wf
    inches_per_pt = 1.0 / 72.27  # Convert pt to inch
    fig_width = fig_width_pt * inches_per_pt  # width in inches
    fig_height = fig_width * hf  # height in inches
    return [fig_width, fig_height]


@contextmanager
def hiro_style():
    with plt.style.context("ggplot"):
        with matplotlib.rc_context(
            {
                "font.family": "serif",
                "text.usetex": False,
                "pgf.rcfonts": False,
                "lines.linewidth": 1,
                "font.size": 10,
                "axes.labelsize": 10,
                "axes.titlesize": 8,
                "font.size": 10,
                "legend.fontsize": 8,
                "xtick.labelsize": 8,
                "ytick.labelsize": 8,
            }
        ):
            yield True


@wrap_plot
def plot_complex(x, y, *args, ax=None, label="", **kwargs):
    label = label + ", " if (len(label) > 0) else ""
    ax.plot(x, y.real, *args, label=f"{label}real part", **kwargs)
    ax.plot(x, y.imag, *args, label=f"{label}imag part", **kwargs)
    ax.legend()


@wrap_plot
def plot_convergence(
    x,
    y,
    ax=None,
    label="",
    transform=lambda y: y,
    slice=None,
    linestyle="-",
    bath=None,
):
    label = label + ", " if (len(label) > 0) else ""
    slice = (0, -1) if not slice else slice
    for n, val, _ in y[slice[0] : slice[1]]:
        plt.plot(
            x,
            transform(val[bath] if bath is not None else val),
            label=f"{label}$N={n}$",
            alpha=n / y[-1][0],
            linestyle=":",
        )

    ax.errorbar(
        x,
        transform(y[-1][1][bath] if bath is not None else y[-1][1]),
        yerr=(y[-1][2][bath] if bath is not None else y[-1][2]).real,
        ecolor="yellow",
        label=f"{label}$N={y[-1][0]}$",
        color="red",
        linestyle=linestyle,
    )

    return None


@wrap_plot
def plot_diff_vs_sigma(
    x,
    y,
    reference,
    ax=None,
    label="",
    transform=lambda y: y,
    ecolor="yellow",
    ylabel=None,
    bath=None,
):
    label = label + ", " if (len(label) > 0) else ""
    y = y if bath is None else [(n, val[bath], err[bath]) for n, val, err in y]

    ref_traj = transform(reference)
    ref_err = np.abs(y[-1][2].real)

    ax.fill_between(
        x,
        0,
        ref_err,
        color=ecolor,
        label=fr"{label}$\sigma\, (N={y[-1][0]})$",
    )

    for n, val, err in y:
        diff = np.abs(transform(val) - ref_traj)
        within = (diff < ref_err).sum() / y[-1][2].size

        not_last = n < y[-1][0]

        ax.plot(
            x,
            diff,
            label=fr"{label}$N={n}$  $\Delta<\sigma = {within * 100:.1f}\%$",
            alpha=within if not_last else 1,
            linestyle=":" if not_last else "-",
            color=None if not_last else "red",
        )

    if ylabel is not None:
        if ylabel[0] == "$":
            ylabel = ylabel[1:-1]
        else:
            ylabel = fr"\text{{ {ylabel} }}"

        ax.set_ylabel(fr"$|{{{ylabel}}}_{{\mathrm{{ref}}}}-{{{ylabel}}}_{{N_i}}|$")