import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from plotstyle import plot_style


base_path = Path('data')
data_path = base_path / 'cells'
sims_path = base_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, freqs, chi2, nsegs):
    ax.set_aspect('equal')
    i0 = np.argmin(freqs < 0)
    i1 = np.argmax(freqs > 300)
    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)
    ax.set_xlim(0, 300)
    ax.set_ylim(0, 300)
    ax.set_xticks_delta(100)
    ax.set_yticks_delta(100)
    ax.set_xlabel('$f_1$', 'Hz')
    ax.set_ylabel('$f_2$', 'Hz')
    ax.text(1, 1.1, f'$N=10^{{{np.log10(nsegs):.0f}}}$',
            ha='right', transform=ax.transAxes)
    
    cax = ax.inset_axes([1.04, 0, 0.05, 1])
    cax.set_spines_outward('lrbt', 0)
    cb = fig.colorbar(pc, cax=cax)
    cb.outline.set_color('none')
    cb.outline.set_linewidth(0)
    cax.set_ylabel(r'$|\chi_2|$ [Hz]')
    cax.set_yticks_delta(ten)


def plot_chi2_contrast(ax1, ax2, s, cell_name, contrast, nsmall, nlarge):
    data_files = sims_path.glob(f'chi2-noisen-{cell_name}-{1000*contrast:03.0f}-*.npz')
    files, nums = sort_files(cell_name, data_files, 2)
    for ax, n in zip([ax1, ax2], [nsmall, nlarge]):
        i = nums.index(n)
        data = np.load(files[i])
        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))
        plot_chi2(ax, s, freqs, chi2, n)


def plot_chi2_split(ax1, ax2, s, cell_name, nsmall, nlarge):
    data_files = sims_path.glob(f'chi2-split-{cell_name}-*.npz')
    files, nums = sort_files(cell_name, data_files, 1)
    for ax, n in zip([ax1, ax2], [nsmall, nlarge]):
        i = nums.index(n)
        data = np.load(files[i])
        n = data['n']
        alpha = data['alpha']
        noise_frac = data['noise_frac']
        freqs = data['freqs']
        pss = data['pss']
        chi2 = np.abs(data['prss'])*0.5/np.sqrt(pss.reshape(1, -1)*pss.reshape(-1, 1))
        plot_chi2(ax, s, freqs, chi2, n)
    return alpha, noise_frac

        
def plot_chi2_data(ax, s, cell_name, run):
    data_file = data_path / f'{cell_name}-spectral-s{run:02d}.npz'
    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))
    print(f'Measured cell {data_file.name} at {100*alpha:.1f}% contrast')
    plot_chi2(ax, s, freqs, chi2, n)
    return alpha

    
def plot_noise_split(ax, contrast, noise_contrast, noise_frac):
    axr, axs, axn = ax.subplots(3, 1, hspace=0.1)
    tmax = 50
    
    axr.show_spines('l')
    axr.set_xlim(0, tmax)
    axr.set_ylim(-8, 8)
    axr.set_yticks_delta(6)
    axr.set_ylabel('\\%')
    
    axs.show_spines('l')
    axs.set_xlim(0, tmax)
    axs.set_ylim(-8, 8)
    axs.set_yticks_delta(6)
    axs.set_ylabel('\\%')
    
    axn.set_ylim(-6, 6)
    axn.set_xlim(0, tmax)
    axn.set_yticks_delta(6)
    axn.set_yticks_blank()
    axn.set_xticks_delta(25)
    axn.set_xlabel('Time', 'ms')

    
if __name__ == '__main__':
    cell_name = '2012-07-03-ak-invivo-1'
    nsmall = 100
    nlarge = 1000000
    contrast = 0.03
                 
    s = plot_style()
    fig, axs = plt.subplots(2, 4, cmsize=(s.plot_width, 0.4*s.plot_width),
                            width_ratios=[1, 0, 1, 1, 1])
    fig.subplots_adjust(leftm=7, rightm=8, topm=2, bottomm=3.5,
                        wspace=0.4, hspace=0.6)
    axs[1, 0].set_visible(False)
    data_contrast = plot_chi2_data(axs[0, 0], s, cell_name[:13], 0)
    plot_noise_split(axs[0, 1], data_contrast, 0, 1)
    plot_chi2_contrast(axs[0, 2], axs[0, 3], s, cell_name, contrast, nsmall, nlarge)
    noise_contrast, noise_frac = plot_chi2_split(axs[1, 2], axs[1, 3], s,
                                                 cell_name, nsmall, nlarge)
    plot_noise_split(axs[1, 1], contrast, noise_contrast, noise_frac)
    fig.common_xticks(axs[:, 2])
    fig.common_xticks(axs[:, 3])
    fig.common_yticks(axs[0, 2:])
    fig.common_yticks(axs[1, 2:])
    #fig.tag(axs, xoffs=-4.5, yoffs=1.8)
    fig.savefig()
    print()