import scipy
import scipy.signal
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import pint

from ..data import *
from ..fitting import *


# Setup pint
SI = pint.UnitRegistry()
SI.setup_matplotlib()
SI.default_format = '~P'


# Colors
# https://arxiv.org/abs/2107.02270
petroff_colors = ["#3f90da", "#ffa90e", "#bd1f01", "#94a4a2", "#832db6", "#a96b59", "#e76300", "#b9ac70", "#717581", "#92dadd"]
cmap_petroff_10 = mpl.colors.ListedColormap(petroff_colors, 'Petroff 10')
mpl.rcParams['axes.prop_cycle'] = mpl.cycler(color=petroff_colors)

cmap_petroff_gradient = mpl.colors.LinearSegmentedColormap.from_list('Petroff gradient', [petroff_colors[i] for i in (9,0,4,2,6,1)])
cmap_petroff_gradient.set_under(petroff_colors[3])
cmap_petroff_gradient.set_over(petroff_colors[7])
#mpl.rcParams['image.cmap'] = cmap_petroff_gradient

cmap_petroff_bipolar = mpl.colors.LinearSegmentedColormap.from_list('Petroff bipolar', [petroff_colors[i] for i in (2,6,1,3,9,0,4)])
cmap_petroff_bipolar.set_under(petroff_colors[5])
cmap_petroff_bipolar.set_over(petroff_colors[8])


# Matplotlib options
mpl.rcParams.update({
    'figure.figsize': (8, 5),  # Note: overwritten when using the %matplotlib magics
    'figure.dpi': 120,         # Note: overwritten when using the %matplotlib magics
    'figure.constrained_layout.use': True,
    'legend.fontsize': 'x-small',
    'legend.title_fontsize': 'small',
    'grid.color': '#DDD',
})




# Utility methods



def ax_set(ax, swap_xy=False, **kwargs):
    if swap_xy:
        nkwargs = {}
        for k in kwargs:
            nk = (('y' if k[0]=='x' else 'x') + k[1:]) if k[0] in 'xy' else k
            nkwargs[nk] = kwargs[k]
        ax.set(**nkwargs)
    else:
        ax.set(**kwargs)