import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from spectral import whitenoise
from plotstyle import plot_style, noise_files, plot_chi2


example_cell = ['2017-07-18-ai-invivo-1', 1]
model_cell = example_cell[0]

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


def plot_overn(ax, s, files, nmax=1e6):
    ns = []
    stats = []
    for fname in files:
        data = np.load(fname)
        fcutoff = data['fcutoff']
        n = data['nsegs']
        if nmax is not None and n > nmax:
            continue
        alpha = data['contrast']
        freqs = data['freqs']
        pss = data['pss']
        prss = data['prss']
        chi2 = np.abs(prss)/0.5/(pss.reshape(1, -1)*pss.reshape(-1, 1))
        chi2 *= 1e-4 # Hz/%^2
        ns.append(n)
        i0 = np.argmin(freqs < 0)
        i1 = np.argmax(freqs > fcutoff)
        if i1 == 0:
            i1 = len(freqs)
        chi2 = chi2[i0:i1, i0:i1]
        stats.append(np.quantile(chi2, [0, 0.001, 0.05, 0.25, 0.5,
                                        0.75, 0.95, 0.998, 1.0]))
    ns = np.array(ns)
    stats = np.array(stats)
    indx = np.argsort(ns)
    ns = ns[indx]
    stats = stats[indx]
    ax.set_visible(True)
    ax.plot(ns, stats[:, 7], zorder=50, label='99.8\\%', **s.lsMax)
    ax.fill_between(ns, stats[:, 2], stats[:, 6], fc='0.85', zorder=40, label='5--95\\%')
    ax.fill_between(ns, stats[:, 3], stats[:, 5], fc='0.5', zorder=45, label='25-75\\%')
    ax.plot(ns, stats[:, 4], zorder=50, label='median', **s.lsMedian)
    #ax.plot(ns, stats[:, 8], '0.0')
    ax.set_xlim(1e2, nmax)
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_yticks_log(numticks=5)
    if nmax > 1e6:
        ax.set_ylim(3e-1, 5e3)
        ax.set_minor_yticks_log(numticks=5)
        ax.set_xticks_log(numticks=4)
        ax.set_minor_xticks_log(numticks=8)
    else:
        ax.set_ylim(1e-1, 2e3)
        #ax.set_minor_yticks_log(numticks=5)
        ax.set_minor_yticks_off()
        ax.set_xticks_log(numticks=6)
        #ax.set_minor_xticks_log(numticks=6)
    ax.set_xlabel('segments')
    ax.set_ylabel(r'$|\chi_2|$', r'Hz/\%$^2$')

        
def plot_chi2_data(ax, s, cell_name, run):
    vmax = 15
    data_file = data_path / f'{cell_name}-baseline.npz'
    data = np.load(data_file)
    eodf = float(data['eodf'])
    ratebase = float(data['ratebase/Hz'])
    cvbase = float(data['cvbase'])
    data_file = data_path / f'{cell_name}-spectral-100-s{run:02d}.npz'
    data = np.load(data_file)
    nsegs = data['nsegs']
    fcutoff = data['fcutoff']
    nfft = data['nfft']
    deltat = data['deltat']
    alpha = data['contrast']
    freqs = data['freqs']
    pss = data['pss']
    prss = data['prss']
    chi2 = np.abs(prss)*0.5/(pss.reshape(1, -1)*pss.reshape(-1, 1))
    print(f'Measured cell {"-".join(data_file.name.split("-")[:-3])} at {100*alpha:4.1f}% contrast:')
    print(f'    r={ratebase:3.0f}Hz, CV={cvbase:4.2f}, dt={1000*deltat:4.2f}ms, nfft={nfft}, win={1000*deltat*nfft:6.1f}ms, nsegs={nsegs}')
    print()
    ax.text(1, 1.1, f'$N={nsegs}$',
            ha='right', transform=ax.transAxes)
    plot_chi2(ax, s, freqs, chi2, fcutoff, ratebase, vmax)
    return alpha, ratebase, eodf


def plot_chi2_contrast(ax1, ax2, s, files, nums, nsmall, nlarge, rate):
    vmax = {0.05: {nsmall: 15, nlarge: 1.2},
            0.01: {nsmall: 400, nlarge: 4}}
    for ax, n in zip([ax1, ax2], [nsmall, nlarge]):
        i = nums.index(n)
        data = np.load(files[i])
        nsegs = data['nsegs']
        fcutoff = float(data['fcutoff'])
        alpha = float(data['contrast'])
        freqs = data['freqs']
        pss = data['pss']
        prss = data['prss']
        if nsegs < 10000:
            ax.text(1, 1.1, f'$N={nsegs}$',
                    ha='right', transform=ax.transAxes)
        else:
            ax.text(1, 1.1, f'$N=10^{{{np.log10(nsegs):.0f}}}$',
                    ha='right', transform=ax.transAxes)
        chi2 = np.abs(prss)*0.5/(pss.reshape(1, -1)*pss.reshape(-1, 1))
        cax = plot_chi2(ax, s, freqs, chi2, fcutoff, rate, vmax[alpha][n])
        cax.set_ylabel('')
        print(f'Modeled  cell {"-".join(files[i].name.split("-")[:-4])} at {100*alpha:4.1f}% contrast: noise_frac={1:3.1f}, nsegs={n}')
    print()


def plot_chi2_split(ax1, ax2, s, files, nums, nsmall, nlarge, rate):
    vmax = {nsmall: 4, nlarge: 1.2}
    for ax, n in zip([ax1, ax2], [nsmall, nlarge]):
        i = nums.index(n)
        data = np.load(files[i])
        nsegs = data['nsegs']
        fcutoff = float(data['fcutoff'])
        alpha = float(data['contrast'])
        noise_frac = data['noise_frac']
        freqs = data['freqs']
        pss = data['pss']
        prss = data['prss']
        chi2 = np.abs(prss)*0.5/(pss.reshape(1, -1)*pss.reshape(-1, 1))
        if nsegs < 10000:
            ax.text(1, 1.1, f'$N={nsegs}$',
                    ha='right', transform=ax.transAxes)
        else:
            ax.text(1, 1.1, f'$N=10^{{{np.log10(nsegs):.0f}}}$',
                    ha='right', transform=ax.transAxes)
        cax = plot_chi2(ax, s, freqs, chi2, fcutoff, rate, vmax[n])
        cax.set_ylabel('')
        print(f'Modeled  cell {"-".join(files[i].name.split("-")[:-3])} at {100*alpha:4.1f}% contrast: noise_frac={noise_frac:3.1f}, nsegs={n}')
    print()
    return alpha, noise_frac

    
def plot_ram(ax, contrast, eodf, wtime, wnoise):
    tmax = 50
    am = 1 + contrast*wnoise
    eod = np.sin(2*np.pi*eodf*wtime)*am

    ax.show_spines('l')
    ax.plot(1e3*wtime, eod, clip_on=False, **s.lsEOD)
    ax.plot(1e3*wtime, +am, clip_on=False, **s.lsAM)
    ax.plot(1e3*wtime, -am, clip_on=False, **s.lsAM)
    ax.set_xlim(0, tmax)
    ax.set_ylim(-1.3, 1.3)
    ax.set_yticks_delta(1)
    ax.set_ylabel('EOD')
    ax.text(1, 1, f'RAM ($c={100*contrast:.0f}$\\,\\%)', ha='right',
            transform=ax.transAxes, color=s.lsAM['color'])

    
def plot_noise_split(ax, contrast, noise_contrast, noise_frac,
                     wtime, wnoise):
    axr, axs, axn = ax.subplots(3, 1, hspace=0.2)
    cmax = 26
    cdelta = 20
    tmax = 50

    axr.show_spines('l')
    axr.axhline(0, **s.lsGrid)
    axr.plot(1e3*wtime, 100*contrast*wnoise, clip_on=False, **s.lsAM)
    axr.set_xlim(0, tmax)
    axr.set_ylim(-cmax, cmax)
    axr.set_yticks_delta(cdelta)
    axr.set_ylabel('\\%')
    axr.text(1, 1, f'RAM ($c={100*contrast:.0f}$\\,\\%)', ha='right',
             transform=axr.transAxes, color=s.lsAM['color'])
    
    axs.show_spines('l')
    axs.axhline(0, **s.lsGrid)
    axs.plot(1e3*wtime, 100*noise_contrast*wnoise, clip_on=False, **s.lsAMsplit)
    axs.set_xlim(0, tmax)
    axs.set_ylim(-cmax, cmax)
    axs.set_yticks_delta(cdelta)
    axs.set_ylabel('\\%')
    if noise_contrast > 0:
        axs.text(1, 1, f'$s_{{\\xi}}(t)$ ($c={100*noise_contrast:.0f}$\\,\\%)',
                 ha='right', transform=axs.transAxes,
                 color=s.lsAMsplit['color'])

    ntime = np.linspace(0, 1e-3*tmax, 800)
    rng = np.random.default_rng(45432)
    nnoise = rng.normal(size=len(ntime))
    axn.show_spines('l')
    axn.axhline(0, **s.lsGrid)
    axn.plot(1e3*ntime, noise_frac*nnoise, clip_on=False, **s.lsNoise)
    axn.set_ylim(-2, 2)
    axn.set_xlim(0, tmax)
    axn.set_yticks_delta(5)
    axn.set_yticks_blank()
    #axn.set_xticks_delta(25)
    #axn.set_xlabel('Time', 'ms')
    y = 0.8 if noise_frac < 1 else 1.2
    axn.text(1, y, f'Intrinsic noise (${100*noise_frac:.0f}$\\,\\%)',
             ha='right', transform=axn.transAxes, color=s.lsNoise['color'])
    if noise_frac < 1:
        axn.xscalebar(1, -0.1, 10, 'ms', ha='right')
        
    return axr


if __name__ == '__main__':
    nsmall = 100
    nlarge = 1000000
    contrast = 0.01

    wdt = 0.0001
    wnoise = whitenoise(0, 300, wdt, 0.05, rng=np.random.default_rng(51234))
    wtime = np.arange(len(wnoise))*wdt
                 
    s = plot_style()
    fig, axs = plt.subplots(4, 4, cmsize=(s.plot_width, 0.85*s.plot_width),
                            width_ratios=[1, 0, 1, 1, 0.2, 0.85])
    fig.subplots_adjust(leftm=8, rightm=1.5, topm=3.5, bottomm=4,
                        wspace=0.25, hspace=0.6)
    axs[0, 2].set_visible(False)
    axs[0, 3].set_visible(False)
    xt = -2.25
    yt = 1.25
    
    # data:
    axss = axs[0]
    axss[1].text(xt, yt, 'P-unit data', fontsize='large',
                 transform=axss[1].transAxes, color=s.punit_color1)
    data_contrast, ratebase, eodf = plot_chi2_data(axss[1], s, example_cell[0],
                                                   example_cell[1])
    plot_ram(axss[0], data_contrast, eodf, wtime, wnoise)
    axss[1].text(xt + 0.9, yt, f'$r={ratebase:.0f}$\\,Hz',
                 transform=axss[1].transAxes, fontsize='large')
    
    # model 5%:
    axss = axs[1]
    files, nums = noise_files(sims_path, model_cell, data_contrast)
    axss[1].text(xt, yt, 'P-unit model', fontsize='large',
                 transform=axs[1, 1].transAxes, color=s.model_color1)
    plot_chi2_contrast(axss[1], axss[2], s, files, nums, nsmall, nlarge,
                       ratebase)
    axr1 = plot_noise_split(axss[0], data_contrast, 0, 1, wtime, wnoise)
    plot_overn(axss[3], s, files, nmax=1e6)
    axss[3].legend(loc='lower center', bbox_to_anchor=(0.5, 1.2),
                   markerfirst=False, title='$|\\chi_2|$ percentiles')
    
    # model 1%:
    axss = axs[2]
    files, nums = noise_files(sims_path, model_cell, contrast)
    plot_chi2_contrast(axss[1], axss[2], s, files, nums, nsmall, nlarge,
                       ratebase)
    axr2 = plot_noise_split(axss[0], contrast, 0, 1, wtime, wnoise)
    plot_overn(axss[3], s, files, nmax=1e6)

    # model noise split:
    axss = axs[3]
    files, nums = noise_files(sims_path, model_cell)
    axss[1].text(xt, yt, 'P-unit model', fontsize='large',
                 transform=axss[1].transAxes, color=s.model_color1)
    axss[1].text(xt + 0.9, yt, f'(noise split)', fontsize='large',
                 transform=axss[1].transAxes)
    noise_contrast, noise_frac = plot_chi2_split(axss[1], axss[2], s,
                                                 files, nums, nsmall, nlarge,
                                                 ratebase)
    axr3 = plot_noise_split(axss[0], 0, noise_contrast, noise_frac,
                            wtime, wnoise)
    plot_overn(axss[3], s, files, nmax=1e6)
    
    fig.common_xticks(axs[:, 1])
    fig.common_xticks(axs[:, 2])
    fig.common_xticks(axs[:, 3])
    fig.common_yticks(axs[1, 1:3])
    fig.common_yticks(axs[2, 1:3])
    fig.common_yticks(axs[3, 1:3])
    fig.tag([axs[0, :2],
             [axr1] + axs[1, 1:].tolist(),
             [axr2] + axs[2, 1:].tolist(),
             [axr3] + axs[3, 1:].tolist()],
            xoffs=[-4.5, 1, 1, -4.5], yoffs=2)
    fig.savefig()