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, style=None):
    fig_path.mkdir(parents=True, exist_ok=True)
    if fig is None:
        fig = plt.gcf()

    with hiro_style(style):
        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 lighten_color(color, amount=0.5):
    """
    Lightens the given color by multiplying (1-luminosity) by the given amount.
    Input can be matplotlib color string, hex string, or RGB tuple.

    Examples:
    >> lighten_color('g', 0.3)
    >> lighten_color('#F034A3', 0.6)
    >> lighten_color((.3,.55,.1), 0.5)
    """
    import matplotlib.colors as mc
    import colorsys

    try:
        c = mc.cnames[color]
    except:
        c = color
    c = colorsys.rgb_to_hls(*mc.to_rgb(c))
    return colorsys.hls_to_rgb(c[0], 1 - amount * (1 - c[1]), c[2])


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(width="thesis", fraction=1, subplots=(1, 1)):
    """Set figure dimensions to avoid scaling in LaTeX.

    Parameters
    ----------
    width: float or string
            Document width in points, or string of predined document type
    fraction: float, optional
            Fraction of the width which you wish the figure to occupy
    subplots: array-like, optional
            The number of rows and columns of subplots.
    Returns
    -------
    fig_dim: tuple
            Dimensions of figure in inches
    """
    if width == "thesis":
        width_pt = 330.62111
    elif width == "poster":
        width_pt = 957.13617
    elif width == "beamer":
        width_pt = 307.28987
    else:
        width_pt = width

    # Width of figure (in pts)
    fig_width_pt = width_pt * fraction
    # Convert from pt to inches
    inches_per_pt = 1 / 72.27

    # Golden ratio to set aesthetic figure height
    # https://disq.us/p/2940ij3
    golden_ratio = (5 ** 0.5 - 1) / 2

    # Figure width in inches
    fig_width_in = fig_width_pt * inches_per_pt
    # Figure height in inches
    fig_height_in = fig_width_in * golden_ratio * (subplots[0] / subplots[1])

    return (fig_width_in, fig_height_in)


MPL_RC = {
    "font.family": "serif",
    "text.usetex": False,
    "pgf.rcfonts": False,
    "lines.linewidth": 1,
    "font.size": 10,
    "axes.labelsize": 10,
    "axes.titlesize": 8,
    "legend.fontsize": 8,
    "xtick.labelsize": 8,
    "ytick.labelsize": 8,
}

MPL_RC_POSTER = {
    "font.family": "sans",
    "text.usetex": False,
    "pgf.rcfonts": False,
    "lines.linewidth": 2,
    "font.size": 17.28,
    "axes.labelsize": 17.28,
    "axes.titlesize": 17.28,
    "legend.fontsize": 17.28,
    "xtick.labelsize": 14.4,
    "ytick.labelsize": 14.4,
}


@contextmanager
def hiro_style(style=None):
    with plt.style.context("ggplot"):
        with matplotlib.rc_context(MPL_RC_POSTER if style == "poster" else MPL_RC):
            yield True


plt.style.use("ggplot")
matplotlib.rcParams.update(MPL_RC)


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


def fancy_error(x, y, err, ax, **kwargs):
    line = ax.plot(
        x,
        y,
        **kwargs,
    )

    err = ax.fill_between(
        x,
        y + err,
        y - err,
        color=lighten_color(line[0].get_color(), 0.5),
        alpha=0.5,
    )

    return line, err


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

    fancy_error(
        x,
        transform(y[-1][1][bath] if bath is not None else y[-1][1]),
        (y[-1][2][bath] if bath is not None else y[-1][2]).real,
        label=f"{label}" + (fr"$N={y[-1][0]}$" if len(ys) > 1 else ""),
        linestyle=linestyle,
        color=color,
        ax=ax,
        **kwargs,
    )

    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}}|$")