import matplotlib.pyplot as plt
import numpy as np
import straxen
from multihist import Hist1d, Histdd
from matplotlib.colors import LogNorm


@straxen.mini_analysis(requires=('peak_basics',))
def plot_peaks_aft_histogram(
        context, run_id, peaks,
        pe_bins=np.logspace(0, 7, 120),
        rt_bins=np.geomspace(2, 1e5, 120),
        extra_labels=tuple(),
        rate_range=(1e-4, 1),
        aft_range=(0, .85),
        figsize=(14, 5)):
    """Plot side-by-side (area, width) histograms of the peak rate
    and mean area fraction top.

    :param pe_bins: Array of bin edges for the peak area dimension [PE]
    :param rt_bins: array of bin edges for the rise time dimension [ns]
    :param extra_labels: List of (area, risetime, text, color) extra labels
    to put on the plot
    :param rate_range: Range of rates to show [peaks/(bin*s)]
    :param aft_range: Range of mean S1 area fraction top / bin to show
    :param figsize: Figure size to use
    """
    livetime_sec = straxen.get_livetime_sec(context, run_id, peaks)

    mh = Histdd(peaks,
                dimensions=(
                    ('area', pe_bins),
                    ('range_50p_area', rt_bins),
                    ('area_fraction_top', np.linspace(0, 1, 100))
                ))

    f, axes = plt.subplots(1, 2, figsize=figsize)

    def std_axes():
        plt.gca().set_facecolor('k')
        plt.yscale('log')
        plt.xscale('log')
        plt.xlabel("Area [PE]")
        plt.ylabel("Range 50% area [ns]")
        labels = [
                     (12, 8, "AP?", 'white'),
                     (3, 150, "1PE\npileup", 'gray'),

                     (30, 200, "1e", 'gray'),
                     (100, 1000, "n-e", 'w'),
                     (2000, 2e4, "Train", 'gray'),

                     (1200, 50, "S1", 'w'),
                     (45e3, 60, "αS1", 'w'),

                     (2e5, 800, "S2", 'w'),
                 ] + list(extra_labels)

        for x, w, text, color in labels:
            plt.text(x, w, text, color=color,
                     verticalalignment='center',
                     horizontalalignment='center')

    plt.sca(axes[0])
    (mh / livetime_sec).sum(axis=2).plot(
        norm=LogNorm(vmin=rate_range[0], vmax=rate_range[1]),
        colorbar_kwargs=dict(extend='both'),
        cblabel='Peaks / (bin * s)')
    std_axes()

    plt.sca(axes[1])
    mh.average(axis=2).plot(
        vmin=aft_range[0], vmax=aft_range[1],
        colorbar_kwargs=dict(extend='max'),
        cmap=plt.cm.jet,
        cblabel='Mean area fraction top')

    std_axes()
    plt.tight_layout()


@straxen.mini_analysis(requires=['event_info'])
def event_scatter(context, run_id, events,
                  show_single=True,
                  s=10,
                  color_range=(None, None),
                  color_dim='s1_area_fraction_top',
                  figsize=(7, 5)):
    """Plot a (cS1, cS2) event scatter plot

    :param show_single: Show events with only S1s or only S2s just besides
    the axes.
    :param s: Scatter size
    :param color_dim: Dimension to use for the color. Must be in event_info.
    :param color_range: Minimum and maximum color value to show.
    :param figsize: (w, h) figure size to use, or leave None to not make a
    new matplotlib figure.
    """
    if figsize is not None:
        plt.figure(figsize=figsize)
    if color_dim == 's1_area_fraction_top' and color_range == (None, None):
        color_range = (0, 0.3)

    plt.scatter(np.nan_to_num(events['cs1']).clip(.9, None),
                np.nan_to_num(events['cs2']).clip(.9, None),
                clip_on=not show_single,
                c=events[color_dim],
                vmin=color_range[0], vmax=color_range[1],
                s=s,
                cmap=plt.cm.jet,
                marker='.', edgecolors='none')

    plt.xlabel('cS1 [PE]')
    plt.xscale('log')
    plt.xlim(1, None)

    plt.ylabel('cS2 [PE]')
    plt.yscale('log')
    plt.ylim(1, None)

    p = context.get_single_plugin(run_id, 'energy_estimates')
    ax = plt.gca()
    el_lim = p.cs1_to_e(np.asarray(ax.get_xlim()))
    ec_lim = p.cs2_to_e(np.asarray(ax.get_ylim()))

    ax2 = ax.twiny()
    ax2.set_xlim(*el_lim)
    ax2.set_xscale('log')
    ax2.set_xlabel("E_light [keVee]")

    ax3 = ax2.twinx()
    ax3.set_ylim(*ec_lim)
    ax3.set_yscale('log')
    ax3.set_ylabel("E_charge [keVee]")

    plt.sca(ax3)
    plt.plot(el_lim, el_lim, c='k', alpha=0.2)
    x = np.geomspace(*el_lim, num=1000)
    e_label = 1.2e-3
    for e_const, label in [
        (0.1, ''), (1, '1\nkeV'), (10, '10\nkeV'),
        (100, '100\nkeV'), (1e3, '1\nMeV'), (1e4, '')]:
        plt.plot(x, e_const - x, c='k', alpha=0.2)
        plt.text(e_const - e_label, e_label, label,
                 bbox=dict(facecolor='white', alpha=0.5, edgecolor='none'),
                 horizontalalignment='center', verticalalignment='center',
                 color='k', alpha=0.5)

    plt.sca(ax)
    if color_range[0] is None:
        extend = 'neither' if color_range[1] is None else 'max'
    else:
        extend = 'min' if color_range[1] is None else 'both'
    if color_dim == 's1_area_fraction_top':
        plt.colorbar(label="S1 area fraction top",
                     extend=extend,
                     ax=[ax, ax3])
    else:
        plt.colorbar(label=color_dim,
                     extend=extend,
                     ax=[ax, ax3])


@straxen.mini_analysis(requires=('event_info',))
def plot_energy_spectrum(
        events,
        color='b', label=None,
        unit=None, exposure_kg_sec=None,
        error_alpha=0.5, errors='fc',
        n_bins=100, min_energy=1, max_energy=100, geomspace=True):
    """Plot an energy spectrum histogram, with 1 sigma
    Poisson confidence intervals around it.

    :param exposure_kg_sec: Exposure in kg * sec
    :param unit: Unit to plot spectrum in. Can be either:
      - events (events per bin)
      - kg_day_kev (events per kg day keV)
      - tonne_day_kev (events per tonne day keV)
      - tonne_year_kev (events per tonne year keV)
    Defaults to kg_day_kev if exposure_kg_sec is provided,
    otherwise events.

    :param min_energy: Minimum energy of the histogram
    :param max_energy: Maximum energy of the histogram
    :param geomspace: If True, will use a logarithmic energy binning.
    Otherwise will use a linear scale.
    :param n_bins: Number of energy bins to use
    
    :param color: Color to plot in
    :param label: Label for the line
    :param error_alpha: Alpha value for the statistical error band
    :param errors: Type of errors to draw, passed to 'errors'
    argument of Hist1d.plot.
    """
    if unit is None:
        if exposure_kg_sec is not None:
            unit = 'kg_day_kev'
        else:
            unit = 'events'

    h = Hist1d(events['e_ces'],
               bins=(np.geomspace if geomspace else np.linspace)(
                   min_energy, max_energy, n_bins))

    if unit == 'events':
        scale, ylabel = 1, 'Events per bin'
    else:
        if exposure_kg_sec is None:
            raise ValueError('you did not specify exposure_kg_sec')
        exposure_kg_day = exposure_kg_sec / (3600 * 24)
        if unit == 'kg_day_kev':
            scale = exposure_kg_day
            ylabel = 'Events / (kg day keV)'
        elif unit == 'tonne_day_kev':
            scale = exposure_kg_day / 1000
            ylabel = 'Events / (tonne day keV)'
        elif unit == 'tonne_year_kev':
            scale = exposure_kg_day / 1000
            ylabel = 'Events / (tonne year keV)'
        else:
            raise ValueError(f"Invalid unit {unit}")
        scale *= h.bin_volumes()

    h.plot(errors=errors,
           error_style='band',
           color=color,
           label=label,
           linewidth=1,
           scale_histogram_by=1 / scale,
           error_alpha=error_alpha)
    plt.yscale('log')
    if geomspace:
        straxen.log_x(min_energy, max_energy, scalar_ticks=True)
    else:
        plt.xlim(min_energy, max_energy)
    plt.ylabel(ylabel)
    plt.xlabel("Energy [keV_ee], CES")


@straxen.mini_analysis(requires=('peak_basics',))
def plot_peak_classification(peaks, s=1):
    """Make an (area, rise_time) scatter plot of peaks
    :param s: Size of dot for each peak
    """
    for cl, color in enumerate('kbg'):
        d = peaks[peaks['type'] == cl]
        plt.scatter(d['area'], d['rise_time'], c=color,
                    s=s, marker='.', edgecolors='none',
                    label={0: 'Unknown', 1: 'S1', 2: 'S2'}[cl])
    plt.legend(loc='lower right', markerscale=10)

    plt.xscale('log')
    plt.yscale('log')
    plt.xlim(1, 2e6)
    plt.ylim(3, 1e4)
    plt.xlabel("Area [PE]")
    plt.ylabel("Rise time [ns]")
