"""
Set the colorscheme for matplotlib plots and latex.

Calling this script generates util/colorscheme.tex containing xcolor definitions.
"""
import matplotlib as mpl
import matplotlib.pyplot as plt
from cycler import cycler

skipasserts = False

GRUVBOX = {
    "bg0":          "#282828",
    "bg0-hard":     "#1d2021",
    "bg0-soft":     "#32302f",
    "bg1":          "#3c3836",
    "bg2":          "#504945",
    "bg3":          "#665c54",
    "bg4":          "#7c6f64",
    "fg0":          "#fbf1c7",
    "fg0-hard":     "#f9f5d7",
    "fg0-soft":     "#f2e5bc",
    "fg1":          "#ebdbb2",
    "fg2":          "#d5c4a1",
    "fg3":          "#bdae93",
    "fg4":          "#a89984",
    "dark-red":     "#cc241d",
    "dark-green":   "#98971a",
    "dark-yellow":  "#d79921",
    "dark-blue":    "#458588",
    "dark-purple":  "#b16286",
    "dark-aqua":    "#689d6a",
    "dark-orange":  "#d65d0e",
    "dark-gray":    "#928374",
    "light-red":    "#fb4934",
    "light-green":  "#b8bb26",
    "light-yellow": "#fabd2f",
    "light-blue":   "#83a598",
    "light-purple": "#d3869b",
    "light-aqua":   "#8ec07c",
    "light-orange": "#f38019",
    "light-gray":   "#a89984",
    "alt-red":      "#9d0006",
    "alt-green":    "#79740e",
    "alt-yellow":   "#b57614",
    "alt-blue":     "#076678",
    "alt-purple":   "#8f3f71",
    "alt-aqua":     "#427b58",
    "alt-orange":   "#af3a03",
    "alt-gray":     "#7c6f64",
}

FORMULASHEET_COLORSCHEME = GRUVBOX

colors = ["red", "orange", "yellow", "green", "aqua", "blue", "purple", "gray"]

# default order for matplotlib
color_order = ["blue", "orange", "green", "red", "purple", "yellow", "aqua", "gray"]

def set_mpl_colorscheme(palette: dict[str, str], variant="dark"):
    P = palette
    if variant == "dark":
        FIG_BG_COLOR =          P["bg0"]
        PLT_FG_COLOR =          P["fg0"]
        PLT_BG_COLOR =          P["bg0"]
        PLT_GRID_COLOR =        P["bg2"]
        LEGEND_FG_COLOR =       PLT_FG_COLOR
        LEGEND_BG_COLOR =       P["bg1"]
        LEGEND_BORDER_COLOR =   P["bg2"]
    else:
        FIG_BG_COLOR =          P["fg0"]
        PLT_FG_COLOR =          P["bg0"]
        PLT_BG_COLOR =          P["fg0"]
        PLT_GRID_COLOR =        P["fg2"]
        LEGEND_FG_COLOR =       PLT_FG_COLOR
        LEGEND_BG_COLOR =       P["fg1"]
        LEGEND_BORDER_COLOR =   P["fg2"]
    COLORS = [P[f"{variant}-{c}"] for c in color_order]


    color_rcParams = {
        'axes.edgecolor': PLT_FG_COLOR,
        'axes.facecolor': PLT_BG_COLOR,
        'axes.labelcolor': PLT_FG_COLOR,
        'axes.titlecolor': 'auto',
        # 'axes.prop_cycle': cycler('color', ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']),
        'axes.prop_cycle': cycler('color', COLORS),
        # 'axes3d.xaxis.panecolor': (0.95, 0.95, 0.95, 0.5),
        # 'axes3d.yaxis.panecolor': (0.9, 0.9, 0.9, 0.5),
        # 'axes3d.zaxis.panecolor': (0.925, 0.925, 0.925, 0.5),
        # 'boxplot.boxprops.color': 'black',
        # 'boxplot.capprops.color': 'black',
        # 'boxplot.flierprops.color': 'black',
        # 'boxplot.flierprops.markeredgecolor': 'black',
        # 'boxplot.flierprops.markeredgewidth': 1.0,
        # 'boxplot.flierprops.markerfacecolor': 'none',
        # 'boxplot.meanprops.color': 'C2',
        # 'boxplot.meanprops.markeredgecolor': 'C2',
        # 'boxplot.meanprops.markerfacecolor': 'C2',
        # 'boxplot.meanprops.markersize': 6.0,
        # 'boxplot.medianprops.color': 'C1',
        # 'boxplot.whiskerprops.color': 'black',
        'figure.edgecolor': PLT_BG_COLOR,
        'figure.facecolor': PLT_BG_COLOR,
        # 'figure.figsize': [6.4, 4.8],
        # 'figure.frameon': True,
        # 'figure.labelsize': 'large',
        'grid.color': PLT_GRID_COLOR,
        # 'hatch.color': 'black',
        'legend.edgecolor': LEGEND_BORDER_COLOR,
        'legend.facecolor': LEGEND_BG_COLOR,
        'xtick.color': PLT_FG_COLOR,
        'ytick.color': PLT_FG_COLOR,
        'xtick.labelcolor': PLT_FG_COLOR,
        'ytick.labelcolor': PLT_FG_COLOR,
        # 'lines.color': 'C0',
        'text.color': PLT_FG_COLOR,
    }

    for k, v in color_rcParams.items():
        plt.rcParams[k] = v

    # override single char codes
    # TODO: use color name with variant from palette instead of order
    mpl.colors.get_named_colors_mapping()["b"] = COLORS[0]
    mpl.colors.get_named_colors_mapping()["o"] = COLORS[1]
    mpl.colors.get_named_colors_mapping()["g"] = COLORS[2]
    mpl.colors.get_named_colors_mapping()["r"] = COLORS[3]
    mpl.colors.get_named_colors_mapping()["m"] = COLORS[4]
    mpl.colors.get_named_colors_mapping()["y"] = COLORS[5]
    mpl.colors.get_named_colors_mapping()["c"] = COLORS[6]
    mpl.colors.get_named_colors_mapping()["k"] = P["fg0"]
    mpl.colors.get_named_colors_mapping()["w"] = P["bg0"]



def color_latex_def(name, color):
    name = "{" + name.replace("-", "_") + "}"
    color = "{" + color.strip("#") + "}"
    return f"\\definecolor{name:10}{{HTML}}{color}"

def generate_latex_colorscheme(palette, variant="light"):
    s = ""
    for n, c in palette.items():
        s += color_latex_def(n, c) + "\n"
    return s