import matplotlib
from matplotlib.figure import Figure
from collections import UserList
from numbers import Number
from contextlib import contextmanager
from copy import copy
import csv
import os


class Scale(UserList):
    def __mul__(self, other):
        if isinstance(other, Number):
            return Scale([s * other for s in self])
        else:
            return Scale([s * o for s, o in zip(self, other)])

    def __rmul__(self, other):
        return self * other


linewidth = 1
bar_line_width = 3
major_tick_size = 3
minor_tick_size = major_tick_size / 2
axes_line_width = .75


pape_major_tick_pad = 1
paper_font_size = 10


class NameDict:
    full_default = False

    def __init__(self, dct):
        self.dct = dct

    def __call__(self, item, dollar=True):
        if not dollar:
            return self.short(item)[1:-1]

        if self.full_default:
            return self.full(item)
        else:
            return self.short(item)

    def long(self, item):
        full_name = self.dct[item][0]
        return full_name

    def short(self, item):
        return self.dct[item][1]

    def full(self, item):
        return self.long(item) + " " + self.short(item)



names_csv_file = os.path.join(os.path.dirname(__file__), "acronyms.csv")

with open(names_csv_file) as f:
    reader = csv.reader(f)
    _names = {}
    for key, long, short, *rest in reader:
        _names[key] = (long, short)



class Paper:
    extension = "pdf"
    plot_superfolder = ""

    heatmap_annot_size = 5

    large_heatmap_scale = 1

    names = _names

    palette = [
        (0, 0, 0),  # Black
        (230, 159, 0),  # Orange
        (86, 180, 233),  # Sky blue
        (0, 158, 115),  # Bluish green
        (240, 228, 66),  # Yellow
        (0, 114, 178),  # Blue
        (213, 94, 0),  # Vermillon
        (204, 121, 167),  # Reddish purple
    ]

    m = Scale([4 / 3, 1])

    s = Scale([4 / 3 * .75, .75])
    sq = Scale([4 / 3 * .75, 4 / 3 * .75])

    xs = Scale([4 / 3 * .5, .5])
    xsq = Scale([4 / 3 * .5, 4 / 3 * .5])

    linewidth = linewidth

    mpl_style = {
        'axes.labelsize': paper_font_size,
        'axes.labelpad': 1,
        'axes.titlesize': paper_font_size,
        'axes.linewidth': axes_line_width,
        'font.size': paper_font_size,
        'xtick.labelsize': paper_font_size,
        'ytick.labelsize': paper_font_size,
        'legend.fontsize': paper_font_size,
        'figure.figsize': (4 / 3, 1),
        'lines.linewidth': linewidth,
        'image.cmap': "viridis",
        'font.family': 'sans-serif',
        'font.sans-serif': "arial",
        'xtick.major.size': major_tick_size,
        'ytick.major.size': major_tick_size,
        'xtick.minor.size': minor_tick_size,
        'ytick.minor.size': minor_tick_size,
        'xtick.major.width': axes_line_width,
        'xtick.minor.width': axes_line_width / 2,
        'ytick.major.width': axes_line_width,
        'ytick.minor.width': axes_line_width / 2,
        'xtick.major.pad': pape_major_tick_pad,
        'ytick.major.pad': pape_major_tick_pad,
        "legend.labelspacing": .2,
        "legend.handlelength": 1,
        "legend.handletextpad": .2,
        "legend.borderpad": 0,
        "savefig.dpi": 600,
        "pdf.fonttype": 42,
        "legend.borderaxespad": .5,
        "axes.formatter.useoffset": False,
        "lines.color": "black",
        "mathtext.default": "regular"
    }

    def __init__(self):
        import seaborn  # b/c of initialization stuff
        seaborn.set_style('ticks')

    def multilabel(self, *args):
        return ", ".join(args)

    def name_dict(self):
        return NameDict

    def acronym_dict(self):
        return self.names

    def monkey_patch(self):
        matplotlib.use("agg", force=True)

    def modify_labelpad(self, x_or_y_axis, labeplad):
        x_or_y_axis.labelpad = labeplad

    def title_y_transoform(self, y):
        return y

    def legend_x_transform(self, x):
        return x

    def maybe_to_black(self, color):
        return color

    def dataset_output_base_folder(self, path):
        return path

    def sanitize_filename(self, filename):
        return filename

    @contextmanager
    def force_pad(self):
        yield

    def maybe_modify_font_size(self, font_size):
        return font_size


thesis_font_size = 11


class PDF:
    extension = "pdf"

class PGF:
    extension = "pgf"


class Thesis:
    output_format = PGF()
    plot_superfolder = "/home/emmanuel/WritingProjects/Thesis/figures/plots"

    large_heatmap_scale = .8

    _names = _names

    palette = [
        (0, 0, 0),  # Black
        (86, 180, 233),  # Sky blue
        (230, 159, 0),  # Orange
        (0, 158, 115),  # Bluish green
        (240, 228, 66),  # Yellow
        (0, 114, 178),  # Blue
        (213, 94, 0),  # Vermillon
        (204, 121, 167),  # Reddish purple
    ]

    m = Paper.m
    s = Paper.s
    sq = Paper.sq
    xs = Paper.xs
    xsq = Paper.xsq

    heatmap_annot_size = thesis_font_size *.8

    _scale_factor = 1.4

    linewidth = linewidth

    mpl_style = {
        'axes.labelsize': thesis_font_size,
        'axes.titlesize': thesis_font_size,
        'axes.linewidth': axes_line_width,
        'font.size': thesis_font_size,
        'xtick.labelsize': thesis_font_size,
        'ytick.labelsize': thesis_font_size,
        'legend.fontsize': thesis_font_size,
        'figure.figsize': (4 / 3, 1),
        'image.cmap': "viridis",
        'xtick.major.size': major_tick_size,
        'ytick.major.size': major_tick_size,
        'xtick.minor.size': minor_tick_size,
        'ytick.minor.size': minor_tick_size,
        'xtick.major.width': axes_line_width,
        'xtick.minor.width': axes_line_width / 2,
        'ytick.major.width': axes_line_width,
        'ytick.minor.width': axes_line_width / 2,
        "legend.labelspacing": .7,
        "legend.handlelength": 1,
        "legend.handletextpad": .2,
        "savefig.dpi": 600,
        "legend.borderpad": 0,
        "legend.borderaxespad": .5,
        "axes.formatter.useoffset": False,
        "lines.color": "black",
        "lines.markersize": 6, # 6 is mpl default value
        "lines.linewidth": 1.5, # 1.5 is the default
        "pgf.preamble": r"\usepackage{amsmath},\usepackage{amssymb}",
        #"text.latex.preamble": r"\usepackage{amsmath},\usepackage{amssymb}"
    }
    @property
    def extension(self):
        return self.output_format.extension

    def maybe_modify_font_size(self, font_size):
        return thesis_font_size

    def dataset_output_base_folder(self, path):
        return os.path.join(self.plot_superfolder, os.path.basename(path))

    def sanitize_filename(self, filename):
        return filename.replace(" ", "_")

    def multilabel(self, *args):
        return ",\n".join(args)

    @contextmanager
    def force_pad(self):
        self._force_pad = True
        yield
        self._force_pad = False

    def __init__(self):
        print("Using Thesis config")
        self._force_pad = False
        self.names = self.acronym_dict(capitalize=True)
        import matplotlib as mpl
        mpl.use("pgf")
        import seaborn  # b/c of initialization stuff
        seaborn.set_style('ticks')
        mpl.rcParams.update(self.mpl_style)

    def modify_labelpad(self, x_or_y_axis, labeplad):
        pass

    def name_dict(self):
        NameDict.full_default = True
        return NameDict

    def acronym_dict(self, capitalize=False):
        names = copy(self._names)
        for key, (long, short) in self._names.items():
            if key in self.replacements:
                new_long, new_short = self.replacements[key]
                new_short = short if new_short == "" else new_short
            else:
                new_long = long
                new_short = short
            if capitalize:
                new_long = new_long[0].upper() + new_long[1:]
            names[key] = (new_long, new_short)
        return names


    def title_y_transoform(self, y):
        return 1

    def legend_x_transform(self, x):
        return 1

    def maybe_to_black(self, color):
        return "k"

    def monkey_patch(self, scale_factor=None):
        if scale_factor is None:
            scale_factor = self._scale_factor
        from matplotlib.backends.backend_pgf import FigureCanvasPgf
        #matplotlib.backend_bases.register_backend('pdf', FigureCanvasPgf)
        import matplotlib.pyplot as plt

        plt.rcParams["font.family"] = "serif"
        plt.rcParams["font.serif"] = []




        def ignore_labelpad(f):
            def ignored(*args, force_pad=False, **kwargs):
                try:
                    if not self._force_pad:
                        del kwargs["labelpad"]
                except KeyError:
                    pass
                return f(*args, **kwargs)

            return ignored

        matplotlib.axes.Axes.set_xlabel = ignore_labelpad(
            matplotlib.axes.Axes.set_xlabel)
        matplotlib.axes.Axes.set_ylabel = ignore_labelpad(
            matplotlib.axes.Axes.set_ylabel)

        def scale_figsize(f):
            def new_set_size_inches(self_orig, w, h=None, *args, **kwargs):
                try:
                    w_scaled = tuple(c * scale_factor for c in w)
                except TypeError:
                    w_scaled = w * scale_factor

                if h is not None:
                    h_scaled = h * scale_factor
                else:
                    h_scaled = h

                return f(self_orig, w_scaled, h_scaled, *args, **kwargs)
            return new_set_size_inches

        Figure.set_size_inches = scale_figsize(Figure.set_size_inches)


with open(names_csv_file) as f:
    reader = csv.reader(f)
    Thesis.replacements = {}
    for key, _long_, _short_, long_repl, short_repl in reader:
        if long_repl != "":
            Thesis.replacements[key] = (long_repl, short_repl)



from itertools import starmap


def rgb_to_rgb(r1, r2, r3):
    return r1 / 255, r2 / 255, r3 / 255


try:
    thesis_run = os.environ["THESIS_PLOTS"].lower() == "true"
    if thesis_run:
        config = Thesis()
    else:
        config = Paper()
except KeyError:
    config = Paper()

import seaborn as sns
color_palette = list(starmap(rgb_to_rgb, config.palette))
sns.set_palette(color_palette)




