import numpy as np
from scipy.stats import linregress
import matplotlib.pyplot as plt
from pathlib import Path
from plotstyle import plot_style, labels_params
from plotstyle import noise_files, plot_chi2
from modelsusceptcontrasts import load_chi2


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


def plot_overn(ax, s, files, nmax=1e6, title=False):
    ns = []
    stats = []
    for fname in files:
        data = np.load(fname)
        if not 'nsegs' in data:
            return
        n = data['nsegs']
        if nmax is not None and n > nmax:
            continue
        noise_frac = data['noise_frac']
        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))
        ns.append(n)
        i0 = np.argmin(freqs < 0)
        i1 = np.argmax(freqs > 300)
        if i1 == 0:
            i1 = len(freqs)
        chi2 = 1e-4*chi2[i0:i1, i0:i1]   # Hz/%^2
        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')
    if title:
        if noise_frac < 1:
            ax.set_title('$c$=0\\,\\%', fontsize='medium')
        else:
            ax.set_title(f'$c$={100*alpha:g}\\,\\%', fontsize='medium')
    ax.set_xlim(1e1, nmax)
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_yticks_log(numticks=3)
    ax.set_ylim(1e-1, 3e3)
    ax.set_minor_yticks_log(numticks=5)
    if nmax > 1e6:
        ax.set_xticks_log(numticks=4)
        ax.set_minor_xticks_log(numticks=8)
    else:
        ax.set_xticks_log(numticks=3)
        ax.set_minor_xticks_log(numticks=6)
    ax.set_xlabel('segments')
    ax.set_ylabel('$|\\chi_2|$ [Hz]')
    if alpha == 0.10:
        ax.legend(loc='upper left', bbox_to_anchor=(1.4, 1.3),
                  markerfirst=False, title='$|\\chi_2|$ percentiles')

    
def plot_chi2_overn(axs, s, cell_name):
    print(cell_name)
    files, nums = noise_files(sims_path, cell_name)
    for k, nsegs in enumerate([1e2, 1e3, 1e4, 1e6]):
        freqs, chi2, fcutoff, contrast, n = load_chi2(sims_path, cell_name,
                                                      None, nsegs)
        ns = f'$N={n}$' if n < 1000 else f'$N=10^{np.log10(n):.0f}$'
        cax = plot_chi2(axs[k], s, freqs, chi2, fcutoff)
        if k < len(axs) - 2:
            cax.set_ylabel('')
        axs[k].set_title(ns, fontsize='medium')
    plot_overn(axs[-1], s, files)

    
if __name__ == '__main__':
    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
    s = plot_style()
    fig, axs = plt.subplots(6, 6, cmsize=(s.plot_width, 0.9*s.plot_width),
                            width_ratios=[1, 1, 1, 1, 0, 1],
                            height_ratios=[1, 1, 1, 1, 0, 1])
    fig.subplots_adjust(leftm=8, rightm=0.5, topm=2, bottomm=4,
                        wspace=1, hspace=0.8)
    for ax in axs.flat:
        ax.set_visible(False)
    for k in range(len(cells)):
        plot_chi2_overn(axs[k], s, cells[k])
    cell_name = cells[0]
    files, nums = noise_files(sims_path, cell_name)
    plot_overn(axs[-1, 0], s, files, 1e7, True)
    for k, alpha in enumerate([0.01, 0.03, 0.1]):
        files, nums = noise_files(sims_path, cell_name, alpha)
        plot_overn(axs[-1, k + 1], s, files, 1e7, True)
    for k in range(4):
        fig.common_yticks(axs[k, :4])
        fig.common_xticks(axs[:4, k])
    fig.common_xticks(axs[:4, -1])
    fig.align_ylabels(axs[:4, -1], dist=12)
    fig.common_yticks(axs[-1, :4])
    fig.tag(axs, xoffs=-2.5, yoffs=1.8)
    fig.savefig()
    print()