import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import pearsonr, linregress, gaussian_kde
from thunderlab.tabledata import TableData
from pathlib import Path
from spectral import whitenoise, diag_projection, peakedness
from plotstyle import plot_style, labels_params, significance_str


model_cells = ['2017-07-18-ai-invivo-1',  # strong triangle
               '2012-12-13-ao-invivo-1',  # triangle
               '2012-12-20-ac-invivo-1',  # weak border triangle
               '2013-01-08-ab-invivo-1']  # no triangle

data_path = Path('data')
sims_path = data_path / 'simulations'


def sort_files(cell_name, all_files, n):
    files = [fn for fn in all_files if '-'.join(fn.stem.split('-')[2:-n]) == cell_name]
    if len(files) == 0:
        return None, 0
    nums = [int(fn.stem.split('-')[-1]) for fn in files]
    idxs = np.argsort(nums)
    files = [files[i] for i in idxs]
    nums = [nums[i] for i in idxs]
    return files, nums


def plot_chi2(ax, s, data_file, rate):
    fcutoff = 300
    data = np.load(data_file)
    n = data['n']
    alpha = data['alpha']
    freqs = data['freqs']
    pss = data['pss']
    chi2 = np.abs(data['prss'])*0.5/np.sqrt(pss.reshape(1, -1)*pss.reshape(-1, 1))
    ax.set_visible(True)
    ax.set_aspect('equal')
    i0 = np.argmin(freqs < -fcutoff)
    i0 = np.argmin(freqs < 0)
    i1 = np.argmax(freqs > fcutoff)
    if i1 == 0:
        i1 = len(freqs)
    freqs = freqs[i0:i1]
    chi2 = chi2[i0:i1, i0:i1]
    vmax = np.quantile(chi2, 0.996)
    ten = 10**np.floor(np.log10(vmax))
    for fac, delta in zip([1, 2, 3, 4, 6, 8, 10],
                          [0.5, 1, 1, 2, 3, 4, 5]):
        if fac*ten >= vmax:
            vmax = fac*ten
            ten *= delta
            break
    pc = ax.pcolormesh(freqs, freqs, chi2, vmin=0, vmax=vmax,
                       rasterized=True)
    if 'noise_frac' in data:
        ax.set_title('$c$=0\\,\\%', fontsize='medium')
    else:
        ax.set_title(f'$c$={100*alpha:g}\\,\\%', fontsize='medium')
    ax.set_xlim(0, fcutoff)
    ax.set_ylim(0, fcutoff)
    ax.set_xticks_delta(100)
    ax.set_yticks_delta(100)
    ax.set_xlabel('$f_1$', 'Hz')
    ax.set_ylabel('$f_2$', 'Hz')
    dfreqs, diag = diag_projection(freqs, chi2, 2*fcutoff)
    nli, nlif = peakedness(dfreqs, diag, rate, median=False)
    ax.text(0.95, 0.88, f'SI($r$)={nli:.1f}', ha='right', zorder=50,
            color='white', fontsize='medium', transform=ax.transAxes)
    cax = ax.inset_axes([1.04, 0, 0.05, 1])
    cax.set_spines_outward('lrbt', 0)
    if alpha == 0.1:
        cb = fig.colorbar(pc, cax=cax, label=r'$|\chi_2|$ [Hz]')
    else:
        cb = fig.colorbar(pc, cax=cax)
    cb.outline.set_color('none')
    cb.outline.set_linewidth(0)
    cax.set_yticks_delta(ten)

    
def plot_chi2_contrasts(axs, s, cell_name):
    d = sims_path / f'baseline-{cell_name}.npz'
    data = np.load(d)
    rate = float(data['rate'])
    cv = float(data['cv'])
    print(f'    {cell_name}: r={rate:3.0f}Hz, CV={cv:4.2f}')
    files, nums = sort_files(cell_name,
                             sims_path.glob(f'chi2-split-{cell_name}-*.npz'), 1)
    plot_chi2(axs[0], s, files[-1], rate)
    for k, alphastr in enumerate(['010', '030', '100']):
        files, nums = sort_files(cell_name,
                                 sims_path.glob(f'chi2-noisen-{cell_name}-{alphastr}-*.npz'), 2)
        plot_chi2(axs[k + 1], s, files[-1], rate)


def plot_nli_cv(ax, s, data, alpha, cells):
    data = data[data['contrast'] == alpha, :]
    r, p = pearsonr(data['cvbase'], data['dnli'])
    l = linregress(data['cvbase'], data['dnli'])
    x = np.linspace(0, 1, 10)
    ax.set_visible(True)
    ax.set_title(f'$c$={100*alpha:g}\\,\\%', fontsize='medium')
    ax.axhline(1, **s.lsLine)
    ax.plot(x, l.slope*x + l.intercept, **s.lsGrid)
    mask = data['triangle'] > 0.5
    ax.plot(data[mask, 'cvbase'], data[mask, 'dnli'],
            clip_on=False, zorder=30, label='strong', **s.psA1m)
    mask = data['border'] > 0.5
    ax.plot(data[mask, 'cvbase'], data[mask, 'dnli'],
            zorder=20, label='weak', **s.psA2m)
    ax.plot(data['cvbase'], data['dnli'], clip_on=False,
            zorder=10, label='none', **s.psB1m)

    for cell_name in cells:
        mask = data[:, 'cell'] == cell_name
        color = s.psB1m['color']
        if data[mask, 'border']:
            color = s.psA2m['color']
        elif data[mask, 'triangle']:
            color = s.psA1m['color']
        ax.plot(data[mask, 'cvbase'], data[mask, 'dnli'],
                zorder=40, marker='o', ms=s.psB1m['markersize'],
                mfc=color, mec='k', mew=0.8)
    
    ax.set_ylim(0, 8)
    ax.set_xlim(0, 1)
    ax.set_minor_yticks_delta(1)
    ax.set_xlabel('CV$_{\\rm base}$')
    ax.set_ylabel('SI')
    ax.set_yticks_delta(4)
    ax.text(1, 0.9, f'$R={r:.2f}$', transform=ax.transAxes,
            ha='right', fontsize='small')
    ax.text(1, 0.75, significance_str(p), transform=ax.transAxes,
            ha='right', fontsize='small')
    if alpha == 0:
        ax.legend(loc='upper left', bbox_to_anchor=(1.15, 1.05),
                  title='triangle', handlelength=0.5,
                  handletextpad=0.5, labelspacing=0.2)

    kde = gaussian_kde(data['dnli'], 0.15/np.std(data['dnli'], ddof=1))
    nli = np.linspace(0, 8, 100)
    pdf = kde(nli)
    dax = ax.inset_axes([1.04, 0, 0.3, 1])
    dax.show_spines('')
    dax.fill_betweenx(nli, pdf, **s.fsB1a)
    dax.plot(pdf, nli, clip_on=False, **s.lsB1m)

        
def plot_summary_contrasts(axs, s, cells):
    nli_thresh = 1.2
    data = TableData(data_path / 'Apteronotus_leptorhynchus-Punit-models.csv')
    plot_nli_cv(axs[0], s, data, 0, cells)
    print('noise split:')
    cdata = data[data['contrast'] == 0, :]
    nli_split = cdata['dnli']
    print(f'    mean SI = {np.mean(nli_split):.2f}, stdev = {np.std(nli_split):.2f}')
    n = np.sum(nli_split > nli_thresh)
    print(f'    {n} cells ({100*n/len(nli_split):.1f}%) have SI > {nli_thresh:.1f}:')
    for name, cv in cdata[nli_split > nli_thresh, ['cell', 'cvbase']].row_data():
        print(f'        {name:<22} CV={cv:4.2f}')
    print(f'    triangle cells have SI >= {np.min(nli_split[cdata["triangle"] > 0.5]):.2f}')
    print()
    for i, a in enumerate([0.01, 0.03, 0.1]):
        plot_nli_cv(axs[1 + i], s, data, a, cells)
        print(f'contrast {100*a:2g}%:')
        cdata = data[data['contrast'] == a, :]
        nli = cdata['dnli']
        r, p = pearsonr(nli_split, nli)
        print(f'    correlation with split: r={r:.2f}, p={p:.1e}')
        print(f'    mean SI = {np.mean(nli):.2f}, stdev = {np.std(nli):.2f}')
        n = np.sum(nli > nli_thresh)
        print(f'    {n} cells ({100*n/len(nli):.1f}%) have SI > {nli_thresh:.1f}:')
        for name, cv in cdata[nli > nli_thresh, ['cell', 'cvbase']].row_data():
            print(f'        {name:<22} CV={cv:4.2f}')
        print(f'    triangle cells have SI >= {np.min(nli[cdata["triangle"] > 0.5]):.2f}')
        print()
    print('overall lowest baseline CV:',
          ' '.join([f'{cv:.2f}' for cv in np.unique(data['cvbase'])[:5]]))

    
if __name__ == '__main__':
    s = plot_style()
    #labels_params(xlabelloc='right', ylabelloc='top')
    fig, axs = plt.subplots(5, 4, cmsize=(s.plot_width, 0.95*s.plot_width),
                            height_ratios=[3, 3, 3, 3, 0, 3])
    fig.subplots_adjust(leftm=7, rightm=9, topm=2, bottomm=3.5,
                        wspace=1, hspace=0.5)
    print('Example cells:')
    for k in range(len(model_cells)):
        plot_chi2_contrasts(axs[k], s, model_cells[k])
    for k in range(4):
        fig.common_yticks(axs[k, :])
        fig.common_xticks(axs[:4, k])
    print()
    plot_summary_contrasts(axs[4], s, model_cells)
    fig.common_yticks(axs[4, :])
    fig.tag(axs, xoffs=-4.5, yoffs=1.8)
    fig.savefig()
    print()