import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.widgets import SpanSelector
from matplotlib.colors import LogNorm
import numpy as np
import h5py


def plot_all_chan_spectrum(spectrum, bins, *, ax=None, **kwargs):

    def integrate_to_angles(spectrum, bins, lo, hi):
        lo_ind, hi_ind = bins.searchsorted([lo, hi])
        return spectrum[lo_ind:hi_ind].sum(axis=0)

    if ax is None:
        fig, ax = plt.subplots(figsize=(13.5, 9.5))
    else:
        fig = ax.figure

    div = make_axes_locatable(ax)
    ax_r = div.append_axes('right', 2, pad=0.1, sharey=ax)
    ax_t = div.append_axes('top', 2, pad=0.1, sharex=ax)

    ax_r.yaxis.tick_right()
    ax_r.yaxis.set_label_position("right")
    ax_t.xaxis.tick_top()
    ax_t.xaxis.set_label_position("top")

    im = ax.imshow(spectrum, origin='lower', aspect='auto',
                   extent=(-.5, 383.5,
                           bins[0], bins[-1]),
                   norm=LogNorm())

    e_line, = ax_r.plot(spectrum.sum(axis=1), bins[:-1] + np.diff(bins))
    p_line, = ax_t.plot(spectrum.sum(axis=0))
    label = ax_t.annotate('[0, 70] kEv', (0, 1), (10, -10),
                          xycoords='axes fraction',
                          textcoords='offset pixels',
                          va='top', ha='left')

    def update(lo, hi):
        p_data = integrate_to_angles(spectrum, bins, lo, hi)
        p_line.set_ydata(p_data)
        ax_t.relim()
        ax_t.autoscale(axis='y')

        label.set_text(f'[{lo:.1f}, {hi:.1f}] keV')
        fig.canvas.draw_idle()

    span = SpanSelector(ax_r, update, 'vertical', useblit=True,
                        rectprops={'alpha': .5, 'facecolor': 'red'},
                        span_stays=True)

    ax.set_xlabel('channel [#]')
    ax.set_ylabel('E [keV]')

    ax_t.set_xlabel('channel [#]')
    ax_t.set_ylabel('total counts')

    ax_r.set_ylabel('E [keV]')
    ax_r.set_xlabel('total counts')
    ax.set_xlim(-.5, 383.5)
    ax.set_ylim(bins[0], bins[-1])
    ax_r.set_xlim(xmin=0)

    return spectrum, bins, {'center': {'ax': ax, 'im': im},
                            'top': {'ax': ax_t, 'p_line': p_line},
                            'right': {'ax': ax_r, 'e_line': e_line,
                                      'span': span}}


with h5py.File('germ.h5', 'r') as fin:
    spectrum = fin['spectrum'][:]
    bins = fin['bins'][:]

ret = plot_all_chan_spectrum(spectrum, bins)
plt.show()


# Exercise (15 minutes)
# - add span selector to top axes to change curve in right axes