# author : S. Mandalia
#          s.p.mandalia@qmul.ac.uk
#
# date   : March 19, 2018

"""
Plotting functions for the BSM flavour ratio analysis
"""

from __future__ import absolute_import, division

import os
from copy import deepcopy

import numpy as np
import matplotlib as mpl
mpl.use('Agg')
from matplotlib import rc
from matplotlib import pyplot as plt
from matplotlib.offsetbox import AnchoredText

from getdist import plots, mcsamples

from utils import misc as misc_utils
from utils.enums import EnergyDependance, Likelihood, ParamTag, StatCateg
from utils.fr import angles_to_u, angles_to_fr

rc('text', usetex=False)
rc('font', **{'family':'serif', 'serif':['Computer Modern'], 'size':18})


def centers(x):
    return (x[:-1]+x[1:])*0.5


def get_units(dimension):
    if dimension == 3: return r' / GeV'
    if dimension == 4: return r''
    if dimension == 5: return r' / GeV^{-1}'
    if dimension == 6: return r' / GeV^{-2}'
    if dimension == 7: return r' / GeV^{-3}'
    if dimension == 8: return r' / GeV^{-4}'


def calc_nbins(x):
    n =  (np.max(x) - np.min(x)) / (2 * len(x)**(-1./3) * (np.percentile(x, 75) - np.percentile(x, 25)))
    return np.floor(n)


def calc_bins(x):
    nbins = calc_nbins(x)
    return np.linspace(np.min(x), np.max(x)+2, num=nbins+1)


def most_likely(arr):
    """Return the densest region given a 1D array of data."""
    binning = calc_bins(arr)
    harr = np.histogram(arr, binning)[0]
    return centers(binning)[np.argmax(harr)]


def interval(arr, percentile=68.):
    """Returns the *percentile* shortest interval around the mode."""
    center = most_likely(arr)
    sarr = sorted(arr)
    delta = np.abs(sarr - center)
    curr_low = np.argmin(delta)
    curr_up = curr_low
    npoints = len(sarr)
    while curr_up - curr_low < percentile/100.*npoints:
        if curr_low == 0:
            curr_up += 1
        elif curr_up == npoints-1:
            curr_low -= 1
        elif sarr[curr_up]-sarr[curr_low-1] < sarr[curr_up+1]-sarr[curr_low]:
            curr_low -= 1
        elif sarr[curr_up]-sarr[curr_low-1] > sarr[curr_up+1]-sarr[curr_low]:
            curr_up += 1
        elif (curr_up - curr_low) % 2:
            # they are equal so step half of the time up and down
            curr_low -= 1
        else:
            curr_up += 1
    return sarr[curr_low], center, sarr[curr_up]


def flat_angles_to_u(x):
    """Convert from angles to mixing elements."""
    return abs(angles_to_u(x)).astype(np.float32).flatten().tolist()


def plot_Tchain(Tchain, axes_labels, ranges):
    """Plot the Tchain using getdist."""
    Tsample = mcsamples.MCSamples(
        samples=Tchain, labels=axes_labels, ranges=ranges
    )

    Tsample.updateSettings({'contours': [0.90, 0.99]})
    Tsample.num_bins_2D=500
    Tsample.fine_bins_2D=500
    Tsample.smooth_scale_2D=0.03

    g = plots.getSubplotPlotter()
    g.settings.num_plot_contours = 2
    g.settings.axes_fontsize = 10
    g.settings.figure_legend_frame = False
    g.triangle_plot(
        [Tsample], filled=True,
    )
    return g


def gen_figtext(args):
    """Generate the figure text."""
    t = ''
    mr1, mr2, mr3 = misc_utils.solve_ratio(args.measured_ratio)
    if args.fix_source_ratio:
        sr1, sr2, sr3 = misc_utils.solve_ratio(args.source_ratio)
        t += 'Source flavour ratio = [{0}, {1}, {2}]\nIC ' \
                'observed flavour ratio = [{3}, {4}, ' \
                '{5}]\nDimension = {6}'.format(
                    sr1, sr2, sr3, mr1, mr2, mr3, args.dimension,
                    int(args.energy)
                )
    else:
        t += 'IC observed flavour ratio = [{0}, {1}, ' \
                '{2}]\nDimension = {3}'.format(
                    mr1, mr2, mr3, args.dimension, int(args.energy)
                )
    if args.fix_scale:
        t += 'Scale = {0}'.format(args.scale)
    if args.likelihood is Likelihood.GAUSSIAN:
        t += '\nSigma = {0:.3f}'.format(args.sigma_ratio)
    if args.energy_dependance is EnergyDependance.SPECTRAL:
        if not args.fold_index:
            t += '\nSpectral Index = {0}'.format(int(args.spectral_index))
        t += '\nBinning = [{0}, {1}] TeV - {2} bins'.format(
            int(args.binning[0]/1e3), int(args.binning[-1]/1e3), len(args.binning)-1
        )
    elif args.energy_dependance is EnergyDependance.MONO:
        t += '\nEnergy = {0} TeV'.format(int(args.energy/1e3))
    return t


def chainer_plot(infile, outfile, outformat, args, llh_paramset):
    """Make the triangle plot."""
    if not args.plot_angles and not args.plot_elements:
        return

    raw = np.load(infile)
    print 'raw.shape', raw.shape

    misc_utils.make_dir(outfile)
    fig_text = gen_figtext(args)

    axes_labels = llh_paramset.labels
    ranges = llh_paramset.ranges

    if args.plot_angles:
        print "Making triangle plots"
        Tchain = raw
        g = plot_Tchain(Tchain, axes_labels, ranges)

        mpl.pyplot.figtext(0.5, 0.7, fig_text, fontsize=15)

        for i_ax_1, ax_1 in enumerate(g.subplots):
            for i_ax_2, ax_2 in enumerate(ax_1):
                if i_ax_1 == i_ax_2:
                    itv = interval(Tchain[:,i_ax_1], percentile=90.)
                    for l in itv:
                        ax_2.axvline(l, color='gray', ls='--')
                        ax_2.set_title(r'${0:.2f}_{{{1:.2f}}}^{{+{2:.2f}}}$'.format(
                            itv[1], itv[0]-itv[1], itv[2]-itv[1]
                        ), fontsize=10)

        # if not args.fix_mixing:
        #     sc_index = llh_paramset.from_tag(ParamTag.SCALE, index=True)
        #     itv = interval(Tchain[:,sc_index], percentile=90.)
        #     mpl.pyplot.figtext(
        #         0.5, 0.3, 'Scale 90% Interval = [1E{0}, 1E{1}], Center = '
        #         '1E{2}'.format(itv[0], itv[2], itv[1])
        #     )

        for of in outformat:
            g.export(outfile+'_angles.'+of)

    if args.plot_elements:
        print "Making triangle plots"
        if args.fix_mixing_almost:
            raise NotImplementedError
        nu_index = llh_paramset.from_tag(ParamTag.NUISANCE, index=True)
        fr_index = llh_paramset.from_tag(ParamTag.MMANGLES, index=True)
        sc_index = llh_paramset.from_tag(ParamTag.SCALE, index=True)
        if not args.fix_source_ratio:
            sr_index = llh_paramset.from_tag(ParamTag.SRCANGLES, index=True)

        nu_elements = raw[:,nu_index]
        fr_elements = np.array(map(flat_angles_to_u, raw[:,fr_index]))
        sc_elements = raw[:,sc_index]
        if not args.fix_source_ratio:
            sr_elements = np.array(map(angles_to_fr, raw[:,sr_index]))
        if args.fix_source_ratio:
            Tchain = np.column_stack(
                [nu_elements, fr_elements, sc_elements]
            )
        else:
            Tchain = np.column_stack(
                [nu_elements, fr_elements, sc_elements, sr_elements]
            )

        trns_ranges = np.array(ranges)[nu_index,].tolist()
        trns_axes_labels = np.array(axes_labels)[nu_index,].tolist()
        if not args.fix_mixing:
            trns_axes_labels += \
                [r'\mid \tilde{U}_{e1} \mid'    , r'\mid \tilde{U}_{e2} \mid'    , r'\mid \tilde{U}_{e3} \mid'     , \
                 r'\mid \tilde{U}_{\mu1} \mid'  , r'\mid \tilde{U}_{\mu2} \mid'  , r'\mid \tilde{U}_{\mu3} \mid'   , \
                 r'\mid \tilde{U}_{\tau1} \mid' , r'\mid \tilde{U}_{\tau2} \mid' , r'\mid \tilde{U}_{\tau3} \mid']
            trns_ranges += [(0, 1)] * 9
        if not args.fix_scale:
            trns_axes_labels += [np.array(axes_labels)[sc_index].tolist()]
            trns_ranges += [np.array(ranges)[sc_index].tolist()]
        if not args.fix_source_ratio:
            trns_axes_labels += [r'\phi_e', r'\phi_\mu', r'\phi_\tau']
            trns_ranges += [(0, 1)] * 3

        g = plot_Tchain(Tchain, trns_axes_labels, trns_ranges)

        mpl.pyplot.figtext(0.5, 0.7, fig_text, fontsize=15)
        for of in outformat:
            g.export(outfile+'_elements.'+of)


def myround(x, base=5, up=False, down=False):
    if up == down and up is True: assert 0
    if up: return int(base * np.round(float(x)/base-0.5))
    elif down: return int(base * np.round(float(x)/base+0.5))
    else: int(base * np.round(float(x)/base))


def plot_statistic(data, outfile, outformat, args, scale_param, label=None):
    """Make MultiNest factor or LLH value plot."""
    print 'Making Statistic plot'
    fig_text = gen_figtext(args)
    if label is not None: fig_text += '\n' + label

    print 'data', data
    print 'data.shape', data.shape
    scales, statistic = data.T
    min_idx = np.argmin(scales)
    null = statistic[min_idx]
    if args.stat_method is StatCateg.BAYESIAN:
        reduced_ev = -(statistic - null)
    elif args.stat_method is StatCateg.FREQUENTIST:
        reduced_ev = -2*(statistic - null)

    fig = plt.figure(figsize=(7, 5))
    ax = fig.add_subplot(111)

    ax.set_xlim(np.log10(args.scale_region))
    ax.set_xlabel('$'+scale_param.tex+'$')
    if args.stat_method is StatCateg.BAYESIAN:
        ax.set_ylabel(r'Bayes Factor')
    elif args.stat_method is StatCateg.FREQUENTIST:
        ax.set_ylabel(r'$-2\Delta {\rm LLH}$')

    ax.plot(scales, reduced_ev)

    for ymaj in ax.yaxis.get_majorticklocs():
        ax.axhline(y=ymaj, ls=':', color='gray', alpha=0.3, linewidth=1)
    for xmaj in ax.xaxis.get_majorticklocs():
        ax.axvline(x=xmaj, ls=':', color='gray', alpha=0.3, linewidth=1)

    at = AnchoredText(
        '\n'+fig_text, prop=dict(size=7), frameon=True, loc=2
    )
    at.patch.set_boxstyle("round,pad=0.1,rounding_size=0.5")
    ax.add_artist(at)
    misc_utils.make_dir(outfile)
    for of in outformat:
        print 'Saving as {0}'.format(outfile+'.'+of)
        fig.savefig(outfile+'.'+of, bbox_inches='tight', dpi=150)


def plot_sens_full(data, outfile, outformat, args):
    print 'Making FULL sensitivity plot'

    colour = {0:'red', 1:'blue', 2:'green', 3:'purple', 4:'orange', 5:'black'}
    xticks = ['{0}'.format(x) for x in range(np.min(args.dimensions),
                                             np.max(args.dimensions)+1)]
    yranges = [np.inf, -np.inf]
    legend_handles = []

    fig = plt.figure(figsize=(7, 5))
    ax = fig.add_subplot(111)

    ax.set_xlim(args.dimensions[0]-1, args.dimensions[-1]+1)
    ax.set_xticklabels([''] + xticks + [''])
    ax.set_xlabel(r'BSM operator dimension ' + r'$d$')
    ax.set_ylabel(r'${\rm log}_{10} \Lambda^{-1} / GeV^{-d+4}$')

    argsc = deepcopy(args)
    for idim in xrange(len(data)):
        dim = args.dimensions[idim]
        print 'dim', dim
        argsc.dimension = dim
        for isrc in xrange(len(data[idim])):
            src = args.source_ratios[isrc]
            argsc.source_ratio = src

            # fig_text = gen_figtext(argsc)
            scales, statistic = data[idim][isrc].T
            min_idx = np.argmin(scales)
            null = statistic[min_idx]
            if args.stat_method is StatCateg.BAYESIAN:
                reduced_ev = -(statistic - null)
                # al = scales[reduced_ev > np.log(10**(3/2.))] # Strong degree of belief
                al = scales[reduced_ev > 0.4] # Testing
            elif args.stat_method is StatCateg.FREQUENTIST:
                reduced_ev = -2*(statistic - null)
                al = scales[reduced_ev > 2.71] # 90% CL for 1 DOF via Wilks
            if len(al) == 0:
                print 'No points for DIM {0} FRS {1} NULL {2}!'.format(
                    dim, misc_utils.solve_ratio(src), null
                )
                print 'Reduced EV {0}'.format(reduced_ev)
                continue
            lim = al[0]
            label = '[{0}, {1}, {2}]'.format(*misc_utils.solve_ratio(src))
            if lim < yranges[0]: yranges[0] = lim
            if lim > yranges[1]: yranges[1] = lim+4
            line = plt.Line2D(
                (dim-0.1, dim+0.1), (lim, lim), lw=3, color=colour[isrc], label=label
            )
            ax.add_line(line)
            if idim == 0: legend_handles.append(line)
            x_offset = isrc*0.05 - 0.05
            ax.annotate(
                s='', xy=(dim+x_offset, lim), xytext=(dim+x_offset, lim+3),
                arrowprops={'arrowstyle': '<-', 'lw': 1.2, 'color':colour[isrc]}
            )

    try:
        yranges = (myround(yranges[0], up=True), myround(yranges[1], down=True))
        ax.set_ylim(yranges)
    except: pass

    ax.legend(handles=legend_handles, prop=dict(size=8), loc='upper right')
    for ymaj in ax.yaxis.get_majorticklocs():
        ax.axhline(y=ymaj, ls=':', color='gray', alpha=0.4, linewidth=1)
    for xmaj in ax.xaxis.get_majorticklocs():
        ax.axvline(x=xmaj, ls=':', color='gray', alpha=0.4, linewidth=1)

    misc_utils.make_dir(outfile)
    for of in outformat:
        print 'Saving plot as {0}'.format(outfile+'.'+of)
        fig.savefig(outfile+'.'+of, bbox_inches='tight', dpi=150)


def plot_sens_fixed_angle(data, outfile, outformat, args):
    print 'Making FIXED_ANGLE sensitivity plot'

    colour = {0:'red', 1:'blue', 2:'green', 3:'purple', 4:'orange', 5:'black'}
    xticks = [r'$\mathcal{O}_{12}$', r'$\mathcal{O}_{13}$', r'$\mathcal{O}_{23}$']
    argsc = deepcopy(args)
    for idim in xrange(len(data)):
        dim = args.dimensions[idim]
        print '= dim', dim
        argsc.dimension = dim

        yranges = [np.inf, -np.inf]
        legend_handles = []

        fig = plt.figure(figsize=(7, 5))
        ax = fig.add_subplot(111)
        ax.set_xlim(0, len(xticks)+1)
        ax.set_xticklabels([''] + xticks + [''])
        ax.set_xlabel(r'BSM operator angle')
        ax.set_ylabel(r'${\rm log}_{10} \Lambda^{-1} / GeV^{-d+4}$')

        for isrc in xrange(len(data[idim])):
            src = args.source_ratios[isrc]
            argsc.source_ratio = src
            print '== src', src

            for ian in xrange(len(data[idim][isrc])):
                print '=== an', ian
                scales, statistic = data[idim][isrc][ian].T
                min_idx = np.argmin(scales)
                null = statistic[min_idx]
                if args.stat_method is StatCateg.BAYESIAN:
                    reduced_ev = -(statistic - null)
                    al = scales[reduced_ev > np.log(10**(3/2.))] # Strong degree of belief
                elif args.stat_method is StatCateg.FREQUENTIST:
                    reduced_ev = -2*(statistic - null)
                    al = scales[reduced_ev > 2.71] # 90% CL for 1 DOF via Wilks
                print 'reduced_ev', reduced_ev
                if len(al) == 0:
                    print 'No points for DIM {0} FRS {1}\nreduced_ev {2}!'.format(
                        dim, src, reduced_ev
                    )
                    continue
                lim = al[0]
                print 'limit = {0}'.format(lim)
                label = '[{0}, {1}, {2}]'.format(*misc_utils.solve_ratio(src))
                if lim < yranges[0]: yranges[0] = lim
                if lim > yranges[1]: yranges[1] = lim+4
                line = plt.Line2D(
                    (ian+1-0.1, ian+1+0.1), (lim, lim), lw=3, color=colour[isrc], label=label
                )
                ax.add_line(line)
                if len(legend_handles) < isrc+1:
                    legend_handles.append(line)
                x_offset = isrc*0.05 - 0.05
                ax.annotate(
                    s='', xy=(ian+1+x_offset, lim), xytext=(ian+1+x_offset, lim+3),
                    arrowprops={'arrowstyle': '<-', 'lw': 1.2, 'color':colour[isrc]}
                )

        try:
            yranges = (myround(yranges[0], up=True), myround(yranges[1], down=True))
            ax.set_ylim(yranges)
        except: pass

        ax.legend(handles=legend_handles, prop=dict(size=8), loc='upper right')
        for ymaj in ax.yaxis.get_majorticklocs():
            ax.axhline(y=ymaj, ls=':', color='gray', alpha=0.4, linewidth=1)
        for xmaj in ax.xaxis.get_majorticklocs():
            ax.axvline(x=xmaj, ls=':', color='gray', alpha=0.4, linewidth=1)

        out = outfile + '_DIM{0}'.format(dim)
        misc_utils.make_dir(out)
        for of in outformat:
            print 'Saving plot as {0}'.format(out+'.'+of)
            fig.savefig(out+'.'+of, bbox_inches='tight', dpi=150)


        for ymaj in ax.yaxis.get_majorticklocs():
            ax.axhline(y=ymaj, ls=':', color='gray', alpha=0.4, linewidth=1)
        for xmaj in ax.xaxis.get_majorticklocs():
            ax.axvline(x=xmaj, ls=':', color='gray', alpha=0.4, linewidth=1)


def plot_sens_corr_angle(data, outfile, outformat, args):
    print 'Making CORR_ANGLE sensitivity plot'

    labels = [r'$sin^2(2\mathcal{O}_{12})$',
              r'$sin^2(2\mathcal{O}_{13})$',
              r'$sin^2(2\mathcal{O}_{23})$']

    argsc = deepcopy(args)
    for idim in xrange(len(data)):
        dim = args.dimensions[idim]
        print '= dim', dim
        argsc.dimension = dim
        for isrc in xrange(len(data[idim])):
            src = args.source_ratios[isrc]
            argsc.source_ratio = src
            print '== src', src
            for ian in xrange(len(data[idim][isrc])):
                print '=== an', ian

                d = data[idim][isrc][ian].reshape(
                    len(data[idim][isrc][ian])**2, 3
                )

                fig = plt.figure(figsize=(7, 5))
                ax = fig.add_subplot(111)
                ax.set_ylim(0, 1)
                ax.set_ylabel(labels[ian])
                ax.set_xlabel(r'${\rm log}_{10} \Lambda^{-1}'+get_units(dim)+r'$')

                xranges = [np.inf, -np.inf]
                legend_handles = []

                y = d[:,0]
                x = d[:,1]
                z = d[:,2]

                print 'x', x
                print 'y', y
                print 'z', z
                null_idx = np.argmin(x)
                null = z[null_idx]
                print 'null = {0}, for scale = {1}'.format(null, x[null_idx])

                if args.stat_method is StatCateg.BAYESIAN:
                    z_r = -(z - null)
                elif args.stat_method is StatCateg.FREQUENTIST:
                    z_r = -2*(z - null)
                print 'scale', x
                print 'reduced ev', z_r

                pdat = np.array([x, y, z_r, np.ones(x.shape)]).T
                print 'pdat', pdat
                pdat_clean = []
                for d in pdat:
                    if not np.any(np.isnan(d)): pdat_clean.append(d)
                pdat = np.vstack(pdat_clean)
                sort_column = 3
                pdat_sorted = pdat[pdat[:,sort_column].argsort()]
                uni, c = np.unique(pdat[:,sort_column], return_counts=True)
                print uni, c
                print len(uni)
                print np.unique(c)
                
                n = len(uni)
                assert len(np.unique(c)) == 1
                c = c[0]
                col_array = []
                for col in pdat_sorted.T:
                    col_array.append(col.reshape(n, c))
                col_array = np.stack(col_array)
                sep_arrays = []
                for x_i in xrange(n):
                    sep_arrays.append(col_array[:,x_i])
                
                print len(sep_arrays)
                sep_arrays = sep_arrays[0][:3]
                print sep_arrays

                if args.stat_method is StatCateg.BAYESIAN:
                    reduced_pdat_mask = (sep_arrays[2] > log(10**(3/2.))) # Strong degree of belief
                elif args.stat_method is StatCateg.FREQUENTIST:
                    reduced_pdat_mask = (sep_arrays[2] > 4.61) # 90% CL for 2 DOFS via Wilks
                reduced_pdat = sep_arrays.T[reduced_pdat_mask].T
                print 'reduced_pdat', reduced_pdat

                ax.tick_params(axis='x', labelsize=11)
                ax.tick_params(axis='y', labelsize=11)

                mini, maxi = np.min(x), np.max(x)
                ax.set_xlim((mini, maxi))
                ax.set_ylim(0, 1)
                ax.grid(b=False)
                
                x_v = reduced_pdat[0].round(decimals=4)
                y_v = reduced_pdat[1].round(decimals=4)
                uniques = np.unique(x_v)
                print 'uniques', uniques
                if len(uniques) == 1: continue
                bw = np.min(np.diff(uniques))
                print 'bw', bw
                print np.diff(uniques)
                uni_x_split = np.split(uniques, np.where(np.diff(uniques) > bw*(1e20))[0] + 1)
                print 'len(uni_x_split)', len(uni_x_split)
                for uni_x in uni_x_split:
                    p_x_l, p_y_l = [], []
                    p_x_u, p_y_u = [], []
                    for uni in uni_x:
                        idxes = np.where(x_v == uni)[0]
                        ymin, ymax = 1, 0
                        for idx in idxes:
                            if y_v[idx] < ymin: ymin = y_v[idx]
                            if y_v[idx] > ymax: ymax = y_v[idx]
                        p_x_l.append(uni)
                        p_y_l.append(ymin)
                        p_x_u.append(uni)
                        p_y_u.append(ymax)
                    p_x_l, p_y_l = np.array(p_x_l, dtype=np.float64), np.array(p_y_l, dtype=np.float64)
                    p_x_u, p_y_u = np.array(list(reversed(p_x_u)), dtype=np.float64), np.array(list(reversed(p_y_u)), dtype=np.float64)
                    p_x = np.hstack([p_x_l, p_x_u])
                    p_y = np.hstack([p_y_l, p_y_u])
                    p_x = np.r_[p_x, p_x[0]]
                    p_y = np.r_[p_y, p_y[0]]
                    print 'p_x', p_x
                    print 'p_y', p_y
                    try:
                        tck, u = interpolate.splprep([p_x, p_y], s=1e-5, per=True)
                        xi, yi = interpolate.splev(np.linspace(0, 1, 1000), tck)
                    except:
                        xi, yi = p_x, p_y
                    ax.fill(xi, yi, 'r', edgecolor='r', linewidth=1)

                out = outfile + '_DIM{0}_SRC{1}_AN{2}'.format(dim, isrc, ian)
                misc_utils.make_dir(out)
                for of in outformat:
                    print 'Saving plot as {0}'.format(out+'.'+of)
                    fig.savefig(out+'.'+of, bbox_inches='tight', dpi=150)