import matplotlib as mpl
from matplotlib import rcParams
import matplotlib.transforms as mtransforms
import matplotlib.pyplot as plt


def _exptext(exp="",
             text="",
             loc=0,
             ax=None,
             fontname=None,
             fontsize=None,
             italic=(False, False),
             pad=0):
    """
    pad in units of axis fraction
    """

    _font_size = rcParams['font.size'] if fontsize is None else fontsize
    fontname = 'TeX Gyre Heros' if fontname is None else fontname

    if ax is None:
        ax = plt.gca()

    loc1_dict = {
        0: {'xy': (0.001, 1 + pad),
            'va': 'bottom',
            },
        1: {'xy': (0.05, 0.95 - pad),
            'va': 'top',
            },
    }

    loc2_dict = {
        0: {'xy': (0.001, 1.005 + pad),
            'va': 'bottom',
            },
        1: {'xy': (0.05, 0.9550 - pad),
            'va': 'bottom',
            },
        2: {'xy': (0.05, 0.9450 - pad),
            'va': 'top',
            },
        3: {'xy': (0.05, 0.95 - pad),
            'va': 'top',
            },
    }

    if loc not in [0, 1, 2, 3]:
        raise ValueError("loc must be in {0, 1, 2}:\n"
                         "0 : Above axes, left aligned\n"
                         "1 : Top left corner\n"
                         "2 : Top left corner, multiline\n"
                         "3 : Split EXP above axes, rest of label in"
                         "top left corner")

    def pixel_to_axis(extent, ax=None):
        # Transform pixel bbox extens to axis fractions
        if ax is None:
            ax = plt.gca()

        extent = extent.transformed(ax.transData.inverted())

        def dist(tup):
            return abs(tup[1] - tup[0])

        dimx, dimy = dist(ax.get_xlim()), dist(ax.get_ylim())
        x, y = ax.get_xlim()[0], ax.get_ylim()[0]
        x0, y0, x1, y1 = extent.extents

        return extent.from_extents(
            abs(x0 - x) / dimx,
            abs(y0 - y) / dimy,
            abs(x1 - x) / dimx,
            abs(y1 - y) / dimy)

    if loc in [0, 3]:
        _exp_loc = 0
    else:
        _exp_loc = 1
    _formater = ax.get_yaxis().get_major_formatter()
    if type(mpl.ticker.ScalarFormatter()) == type(_formater) and _exp_loc == 0:
        ax.figure.canvas.draw()
        _sci_box = pixel_to_axis(ax.get_yaxis().offsetText.get_window_extent())
        _sci_offset = _sci_box.width * 1.1
        loc1_dict[_exp_loc]['xy'] = (_sci_offset, loc1_dict[_exp_loc]['xy'][-1])
        if loc == 0:
            loc2_dict[_exp_loc]['xy'] = (_sci_offset, loc2_dict[_exp_loc]['xy'][-1])

    exp = ax.text(*loc1_dict[_exp_loc]['xy'], s=exp,
                  transform=ax.transAxes,
                  ha='left',
                  va=loc1_dict[_exp_loc]['va'],
                  fontsize=_font_size * 1.3,
                  fontweight='bold',
                  fontstyle='italic' if italic[0] else 'normal',
                  fontname=fontname
                  )

    ax.figure.canvas.draw()
    _dpi = ax.figure.dpi
    _exp_xoffset = exp.get_window_extent().width / _dpi * 1.05
    if loc == 0:
        _t = mtransforms.offset_copy(exp._transform,
                                     x=_exp_xoffset,
                                     units='inches', fig=ax.figure)
    elif loc == 1:
        _t = mtransforms.offset_copy(exp._transform,
                                     x=_exp_xoffset,
                                     y=-exp.get_window_extent().height / _dpi,
                                     units='inches', fig=ax.figure)
    elif loc == 2:
        _t = mtransforms.offset_copy(exp._transform,
                                     y=-exp.get_window_extent().height / _dpi,
                                     units='inches', fig=ax.figure)
    elif loc == 3:
        _t = mtransforms.offset_copy(exp._transform, units='inches', fig=ax.figure)

    ax.text(*loc2_dict[loc]['xy'],
            s=text,
            transform=_t,
            ha='left',
            va=loc2_dict[loc]['va'],
            fontsize=_font_size,
            fontname=fontname,
            fontstyle='italic' if italic[1] else 'normal'
            )

    return ax


# Lumi text
def lumitext(text="",
             ax=None,
             fontname=None,
             fontsize=None):

    _font_size = rcParams['font.size'] if fontsize is None else fontsize
    fontname = 'TeX Gyre Heros' if fontname is None else fontname

    if ax is None:
        ax = plt.gca()

    ax.text(x=1, y=1.005, s=text,
            transform=ax.transAxes,
            ha='right',
            va='bottom',
            fontsize=_font_size * 0.95,
            fontweight='normal',
            fontname=fontname,
            )


# Wrapper
def _explabel(ax=None, loc=0, data=False, paper=False, supplementary=False,
              year=2017, lumi=None, llabel=None, rlabel=None, fontname=None,
              fontsize=None, pad=0, exp="", italic=(False, False)):

    # Right label
    if rlabel is not None:
        _lumi = rlabel
    else:
        if lumi is not None:
            _lumi = r'{lumi}, {year} (13 TeV)'.format(lumi=str(lumi) +
                                                      r' $\mathrm{fb^{-1}}$',
                                                      year=str(year))
        else:
            _lumi = '{} (13 TeV)'.format(str(year))

    lumitext(text=_lumi,
             ax=ax,
             fontname=fontname,
             fontsize=fontsize)

    # Left label
    if llabel is not None:
        _label = llabel
    else:
        _label = ""
        if not data:
            _label = " ".join(["Simulation", _label])
        if not paper:
            _label = " ".join([_label, "Preliminary"])
        if supplementary:
            _label = " ".join([_label, "Supplementary"])

        _label = " ".join(_label.split())

    _exptext(exp=exp,
             text=_label,
             loc=loc,
             ax=ax,
             fontname=fontname,
             fontsize=fontsize,
             italic=italic,
             pad=pad)

    return ax
