import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from spectral import diag_projection, peak_size
from plotstyle import plot_style
from plotstyle import plot_chi2
from punitexamplecell import load_baseline, load_noise, load_spectra
from punitexamplecell import plot_response_spectrum, plot_response
from punitexamplecell import plot_gain, plot_diagonals, plot_isih_small


example_cell = [['2012-05-15-ac', 3],
                ['2012-05-15-ac', 1]]

example_cells = [
    ['2010-11-26-an', 0],
    ['2010-11-08-aa', 1],
    ['2011-02-18-ab', 1],
    ['2014-01-16-aj', 5],
]

data_path = Path('data') / 'cells'


def plot_isih(ax, s, rate, cv, isis, pdf):
    ax.show_spines('b')
    ax.fill_between(1000*isis, pdf, facecolor=s.cell_color1)
    ax.set_xlim(0, 12)
    ax.set_xticks_delta(4)
    ax.set_xlabel('ISI', 'ms')
    ax.text(0.95, 1.08, f'CV$_{{\\rm base}}$={cv:.2f}, $r={rate:.0f}$Hz',
            transform=ax.transAxes)

    
if __name__ == '__main__':
    """
    # find a nice example cell:
    from thunderlab.tabledata import TableData
    data = TableData('data/Apteronotus_leptorhynchus-Ampullary-data.csv')
    data = data[(data['fcutoff'] > 140) & (data['fcutoff'] < 160), :]
    data = data[(data['sinorm_nmax'] > 5) & (data['sinorm_nmax'] < 50), :]
    data = data[(data['contrast'] > 0.04) & (data['contrast'] < 0.06), :]
    data = data[(data['respmod2'] > 0) & (data['respmod2'] < 100), :]
    data = data[(data['cvbase'] > 0) & (data['cvbase'] < 0.2), :]
    data = data[(data['ratebase'] > 50) & (data['ratebase'] < 180), :]
    for k in range(data.rows()):
        print(f'{data[k, "cell"]:<22s} s{data[k, "stimindex"]:02.0f}: '
              f'{100*data[k, "contrast"]:3g}%, r={data[k, "ratebase"]:3.0f}Hz, '
              f'CV={data[k, "cvbase"]:4.2f}, '
              f'rmod={data[k, "respmod2"]:3.0f}Hz, '
              f'SI={data[k, "sinorm_nmax"]:5.2f}')
    print()
    #exit()
    """
    
    #mode = 'all'
    mode = '100'

    cell_name = example_cell[0][0]
    print('Example Ampullary cell:')
    eodf, rate, cv, isis, pdf, freqs, prr = load_baseline(data_path, cell_name)
    print(f'    {cell_name:<22s}:         fbase={rate:3.0f}Hz, CV={cv:.2f}')
    contrast1, time1, stimulus1, spikes1 = load_noise(data_path,
                                                      *example_cell[0])
    contrast2, time2, stimulus2, spikes2 = load_noise(data_path,
                                                      *example_cell[1])
    fcutoff1, contrast1, freqs1, gain1, chi21 = load_spectra(data_path, mode,
                                                             *example_cell[0])
    fcutoff2, contrast2, freqs2, gain2, chi22 = load_spectra(data_path, mode,
                                                             *example_cell[1])
    print(f'    contrast1: {100*contrast1:4.1f}%   contrast2: {100*contrast2:4.1f}%')
    print(f'    fcutoff1 :  {fcutoff1:3.0f}Hz  fcutoff2 :  {fcutoff2:3.0f}Hz')
    print(f'    duration1: {time1[-1]:4.1f}s   duration2: {time2[-1]:4.1f}s')
    
    s = plot_style()
    s.cell_color1 = s.ampul_color1
    s.cell_color2 = s.ampul_color2
    s.lsC1 = s.lsA1
    s.lsC2 = s.lsA2
    s.psC1 = s.psA1
    s.psC2 = s.psA2
    fig, (ax1, ax2, ax3) = \
        plt.subplots(3, 1, height_ratios=[3, 0, 3, 0.3, 4.7],
                     cmsize=(s.plot_width, 0.85*s.plot_width))
    fig.subplots_adjust(leftm=8, rightm=2, topm=2, bottomm=3.5,
                        wspace=0.4, hspace=0.42)
    axi, axp, axr = ax1.subplots(1, 3, width_ratios=[2, 3, 0, 10, 0.2])
    axg, axc1, axc2, axd = ax2.subplots(1, 4, wspace=0.2,
                                        width_ratios=[3.5, 0.5, 4, 4, 0.8, 3.5])
    axs = ax3.subplots(2, 4, wspace=0.4, hspace=0.35,
                       width_ratios=[1, 0.1, 1, 1, 1, 0.1],
                       height_ratios=[1, 4])
    
    axi.text(0, 1.08, 'Ampullary:', transform=axi.transAxes,
             color=s.cell_color1, fontsize='large')
    plot_isih(axi, s, rate, cv, isis, pdf)
    plot_response_spectrum(axp, s, eodf, rate, freqs, prr)
    plot_response(axr, s, eodf, time1, stimulus1, contrast1, spikes1,
                  contrast2, spikes2, am=False)
    
    plot_gain(axg, s, contrast1, freqs1, gain1,
              contrast2, freqs2, gain2, fcutoff1, ymax=12, dy=4)
    axg.axvline(rate, **s.lsGrid)
    axg.text(rate, 12.5, '$r$', ha='center')    
    axc = plot_chi2(axc1, s, freqs2, chi22, fcutoff2, None, 12)
    axc.remove()
    axc1.set_title(f'$c$={100*contrast2:g}\\,\\%',
                   fontsize='medium', color=s.cell_color2)
    plot_chi2(axc2, s, freqs1, chi21, fcutoff1, None, 12)
    axc2.set_title(f'$c$={100*contrast1:g}\\,\\%',
                   fontsize='medium', color=s.cell_color1)
    plot_diagonals(axd, s, rate, contrast1, freqs1, chi21,
                   contrast2, freqs2, chi22, fcutoff1, ymax=17, toffs=1)
    
    fig.common_yticks(axc1, axc2)
    fig.tag([axi, axp, axr], xoffs=-3, yoffs=0)
    fig.tag([axg, axc1, axc2, axd], xoffs=-3, yoffs=2)
    print()

    print('Additional example cells:')
    axs[0, 0].text(0, 1.6, 'Ampullary cells:', transform=axs[0, 0].transAxes,
                   color=s.cell_color1, fontsize='large')
    for k, (cell, run) in enumerate(example_cells):
        eodf, rate, cv, isis, pdf, _, _ = load_baseline(data_path, cell)
        fcutoff, contrast, freqs, gain, chi2 = load_spectra(data_path, mode,
                                                            cell, run)
        print(f'    {cell:<22s}: run={run:2d}, contrast={100*contrast:3.2g}%, '
              f'fbase={rate:3.0f}Hz, CV={cv:.2f}')
        plot_isih_small(axs[0, k], s, contrast, rate, cv, isis, pdf)
        vmax = 15 if k > 0 else 3
        axc = plot_chi2(axs[1, k], s, freqs, chi2, fcutoff, rate, vmax)
        if k % 3 != 0:
            axc.remove()
        if k == 0:
            axc.set_ylabel('')
    fig.common_yticks(axs[1, :])
    fig.tag([axs[0, :3]], xoffs=-3, yoffs=1)
    fig.tag(axs[0, 3:])
    
    fig.savefig()
    print()