#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Sep  9 22:23:30 2023

@author: timofej
"""



def sampling(dim,eps,samplemethod,extremes,dtmax,noeff=1):
    print(f'Starting causal simulation with dim={dim}, eps={eps:.3f}, {samplemethod}, {extremes} extremes for dt={dtmax} and noeff={noeff:03d}')
    
    import os
    import json
    import math as m
    import numpy as np
    from numpy.linalg import norm
    from datetime import datetime
    from random import sample as choose

    from plot import qtplot

    from neuropercolation import Simulate4Layers

    eps_space = list(np.linspace(0.01,0.2,20))

    def resultant(sample):
        phase_x = [m.cos(ind) for ind in sample]
        phase_y = [m.sin(ind) for ind in sample]
        
        return (np.average(phase_x), np.average(phase_y))

    def new_folder(path):
        if not os.path.exists(path):
            os.makedirs(path)

    phase = np.vectorize(lambda x,y: (m.atan2(y,x)+m.pi)%(2*m.pi)-m.pi)
    diff = np.vectorize(lambda x,y: (y-x+m.pi)%(2*m.pi)-m.pi)
    H2 = lambda x: -x*m.log2(x)-(1-x)*m.log2(1-x)    
    
    #print(f'Started at {datetime.now()} with eps={eps:.3f}')
    
    eps = round(eps,3)
    path='/cloud/Public/_data/neuropercolation/4lay/cons=27-knight_steps=1000100/dim=09/batch=0/'
    
    try:
        with open(path+f"eps={round(eps,3):.3f}_phase_diff.txt", 'r', encoding='utf-8') as f:
            phase_diff = json.load(f)
    except:
        with open(path+f"eps={round(eps,3):.3f}_activation.txt", 'r', encoding='utf-8') as f:
            activation = json.load(f)[100:]
    
        osc = list(zip(*activation))
        phase_abs = np.array([[np.arctan2(*act[::-1]) for act in osc[i]] for i in range(2)])
        phase_diff = diff(phase_abs[0],phase_abs[1])
        phase_diff = [round(pha,6) for pha in phase_diff]
        
        with open(path+f"eps={round(eps,3):.3f}_phase_diff.txt", 'w', encoding='utf-8') as f:
            json.dump(list(phase_diff), f, indent=1)

    
    all_res = norm(resultant(phase_diff))
    av_diff = np.arccos(all_res)
    
    try:
        with open(path+f"eps={round(eps,3):.3f}_ei.txt", 'r', encoding='utf-8') as f:
            ei = json.load(f)
    except:
        with open(path+f"eps={round(eps,3):.3f}_channels.txt", 'r', encoding='utf-8') as f:
            channels = json.load(f)[100:]
        
        ei = [round(np.sum(cha)*(1-H2(eps)),6) for cha in channels]
        
        with open(path+f"eps={round(eps,3):.3f}_ei.txt", 'w', encoding='utf-8') as f:
            json.dump(ei, f, indent=1)
    
    ei_ind = [i for i,val in enumerate(ei) if val>0]
    
    print(f'{len(ei_ind)} states with positive EI')
    
    samples = choose(ei_ind, extremes)
    
    with open(path+f"eps={round(eps,3):.3f}_states.txt", 'r', encoding='utf-8') as f:
        states = json.load(f)[100:]
    with open(path+f"eps={round(eps,3):.3f}_coupling.txt", 'r', encoding='utf-8') as f:
        coupling = json.load(f)
        coupling = [tuple(edge) for edge in coupling]
    
    phase_pairs = [[] for dt in range(dtmax+1)]
    ei_pairs = [[] for dt in range(dtmax+1)]
    for num,i in enumerate(samples):
        causal_init = states[i].translate(str.maketrans('', '', '.-='))
        
        sim = Simulate4Layers(  dim,
                                eps,
                                coupling=coupling,
                                init=causal_init,
                                noeffect=noeff-1,
                                steps=dtmax,
                                draw=None,
                                )
            
        activation = sim._activations()
        channel = sim._channels()
    
        osc = list(zip(*activation))
        phase_abs = np.array([[np.arctan2(*act[::-1]) for act in osc[i]] for i in range(2)])
        phasediff_c = diff(phase_abs[0],phase_abs[1])
        ei_c = [round(np.sum(cha)*(1-H2(eps)),6) for cha in channel]
        max_ei_c = max([np.sum(cha) for cha in channel])

        sim = Simulate4Layers(  dim,
                                eps,
                                coupling=coupling,
                                init=causal_init,
                                noeffect=-1,
                                steps=dtmax,
                                draw=None,
                                )
            
        activation = sim._activations()
        channel = sim._channels()
        
        osc = list(zip(*activation))
        phase_abs = np.array([[np.arctan2(*act[::-1]) for act in osc[i]] for i in range(2)])
        phasediff_i = diff(phase_abs[0],phase_abs[1])
        ei_i = [round(np.sum(cha)*(1-H2(eps)),6) for cha in channel]
        max_ei_i = max([np.sum(cha) for cha in channel])
        
        for dt in range(1,dtmax+1):
            phase_pairs[dt].append((i, phasediff_i[dt], phasediff_c[dt]))
            ei_pairs[dt].append((i, ei_i[dt], ei_c[dt]))
                    
        if num%1000==999:
            print(f'Done {num:0{len(str(extremes))}d}')
    
    
    for dt in range(1,dtmax+1):
        savepath = path + samplemethod + f'_samples={extremes}/dt={dt}/noeff={noeff:03d}/'
        new_folder(savepath)
        with open(savepath+f"eps={round(eps,3):.3f}_phase_pairs.txt", 'w', encoding='utf-8') as f:
            json.dump(phase_pairs[dt], f, indent=1)
        with open(savepath+f"eps={round(eps,3):.3f}_ei_pairs.txt", 'w', encoding='utf-8') as f:
            json.dump(ei_pairs[dt], f, indent=1)

def plotting(dim,eps,samplemethod,extremes,dt):
    import os
    import json
    import math as m
    import numpy as np
    from numpy.linalg import norm
    from datetime import datetime
    from random import sample as choose

    from plot import qtplot

    from neuropercolation import Simulate4Layers

    eps_space = list(np.linspace(0.01,0.2,20))

    def resultant(sample):
        phase_x = [m.cos(ind) for ind in sample]
        phase_y = [m.sin(ind) for ind in sample]
        
        return (np.average(phase_x), np.average(phase_y))

    def new_folder(path):
        if not os.path.exists(path):
            os.makedirs(path)

    phase = np.vectorize(lambda x,y: (m.atan2(y,x)+m.pi)%(2*m.pi)-m.pi)
    diff = np.vectorize(lambda x,y: (y-x+m.pi)%(2*m.pi)-m.pi)
    H2 = lambda x: -x*m.log2(x)-(1-x)*m.log2(1-x)

    
    print(f'Started at {datetime.now()} with eps={eps:.3f}')
    
    eps = round(eps,3)
    path='/cloud/Public/_data/neuropercolation/4lay/cons=27-knight_steps=1000100/dim=09/batch=0/'
    
    savepath = path + samplemethod + f'_samples={extremes}/dt={dt}/'
    
    try:
        with open(savepath+f"eps={round(eps,3):.3f}_phase_pairs.txt", 'r', encoding='utf-8') as f:
            phase_pairs = json.load(f)
        with open(savepath+f"eps={round(eps,3):.3f}_ei_pairs.txt", 'r', encoding='utf-8') as f:
            ei_pairs = json.load(f)
    except:
        sampling(dim,eps,samplemethod,extremes,dt)
        with open(savepath+f"eps={round(eps,3):.3f}_phase_pairs.txt", 'r', encoding='utf-8') as f:
            phase_pairs = json.load(f)
        with open(savepath+f"eps={round(eps,3):.3f}_ei_pairs.txt", 'r', encoding='utf-8') as f:
            ei_pairs = json.load(f)
        
    t, phases_i, phases_c = zip(*phase_pairs)
    t, ei_i, ei_c = zip(*ei_pairs)
    
    extremes = len(t)
    
    phases_cdiff = [abs(phases_i[i])-abs(phases_c[i]) for i in range(len(t))]
    
    phase_space = np.linspace(-m.pi,m.pi,101)
    absph_space = np.linspace(0,m.pi,50+1)
    cdiff_space = np.linspace(min(phases_cdiff),max(phases_cdiff),51)
    
    phase_dist_i = [len([ph for ph in phases_i if (phase_space[j]+phase_space[j-1])/2<=ph<(phase_space[j]+phase_space[j+1])/2])/len(t)
                    for j in range(100)]
    phase_dist_c = [len([ph for ph in phases_c if (phase_space[j]+phase_space[j-1])/2<=ph<(phase_space[j]+phase_space[j+1])/2])/len(t)
                    for j in range(100)]
     
    absph_dist_i = [len([ph for ph in phases_i if low<=abs(ph)<high])/len(t) for low,high in zip(absph_space[:-1],absph_space[1:])]
    absph_dist_c = [len([ph for ph in phases_c if low<=abs(ph)<high])/len(t) for low,high in zip(absph_space[:-1],absph_space[1:])]
    
    cdiff_dist = [len([diff for diff in phases_cdiff if low<=diff<high])/len(t) for low,high in zip(cdiff_space[:-1],cdiff_space[1:])]
    
    max_cha_i = max([round(e/(1-H2(eps))) for e in ei_i])
    max_cha_c = max([round(e/(1-H2(eps))) for e in ei_c])

    max_cha = max(max_cha_i, max_cha_c)
    ei_space = np.linspace(0,np.max([ei_i,ei_c]),max_cha+1)
    
    ei_dist_i = [len([e for e in ei_i if round(e/(1-H2(eps)))==val])/len(t) for val in range(max_cha)]
    ei_dist_c = [len([e for e in ei_c if round(e/(1-H2(eps)))==val])/len(t) for val in range(max_cha)]

    qtplot(f'Phase distribution for dt={dt} dim={dim} eps={eps:.3f} with 4 layers',
            [absph_space[:-1]]*2,
            [absph_dist_i, absph_dist_c],
            ['Phase dist with ei',
              'Phase dist without ei'],
            x_tag = 'phase',
            y_tag = 'density',
            export=True,
            path=savepath,
            filename=f'Phase dist eps={round(eps,3):.3f} dim={dim} extremes={extremes}.png',
            close=True)
    
    qtplot(f'Phase distribution for dt={dt} dim={dim} eps={eps:.3f} with 4 layers',
            [phase_space[:-1]]*2,
            [phase_dist_i, phase_dist_c],
            ['Phase dist with ei',
             'Phase dist without ei'],
            x_tag = 'phase',
            y_tag = 'density',
            export=True,
            path=savepath,
            filename=f'Phase original dist eps={round(eps,3):.3f} dim={dim} extremes={extremes}.png',
            close=True)
    
    qtplot(f'Phase causal diff distribution for dt={dt} dim={dim} eps={eps:.3f} with 4 layers',
            [cdiff_space[:-1]],
            [cdiff_dist],
            ['Phase causal difference dist with ei'],
            x_tag = 'sync raise',
            y_tag = 'density',
            export=True,
            path=savepath,
            filename=f'Phase causal diff dist eps={round(eps,3):.3f} dim={dim} extremes={extremes}.png',
            close=True)
    
    qtplot(f'EI distribution for dt={dt} dim={dim} eps={eps:.3f} with 4 layers',
            [ei_space[:-1]]*2,
            [ei_dist_i, ei_dist_c],
            ['EI dist with ei',
              'EI dist without ei'],
            x_tag = 'ei',
            y_tag = 'density',
            export=True,
            path=savepath,
            filename=f'EI dist eps={round(eps,3):.3f} dim={dim} extremes={extremes}.png',
            close=True)
    
    ttest(dim,eps,samplemethod,extremes,dt)
    
    print(f'Done eps={eps:.3f} with dim={dim} at {datetime.now()}')
    
def ttest(dim,eps,samplemethod,extremes,dt):
    from scipy.stats import ttest_rel, ttest_1samp, normaltest
    import numpy as np
    import json
    
    eps = round(eps,3)
    path='/cloud/Public/_data/neuropercolation/4lay/cons=27-knight_steps=1000100/dim=09/batch=0/'
    
    savepath = path + samplemethod + f'_samples={extremes}/dt={dt}/'

    try:
        with open(savepath+f"eps={round(eps,3):.3f}_phase_pairs.txt", 'r', encoding='utf-8') as f:
            phase_pairs = json.load(f)
        with open(savepath+f"eps={round(eps,3):.3f}_ei_pairs.txt", 'r', encoding='utf-8') as f:
            ei_pairs = json.load(f)
    except:
        sampling(dim,eps,samplemethod,extremes,dt)
        with open(savepath+f"eps={round(eps,3):.3f}_phase_pairs.txt", 'r', encoding='utf-8') as f:
            phase_pairs = json.load(f)
        with open(savepath+f"eps={round(eps,3):.3f}_ei_pairs.txt", 'r', encoding='utf-8') as f:
            ei_pairs = json.load(f)
    
    t, phases_i, phases_c = zip(*phase_pairs)
    t, ei_i, ei_c = zip(*ei_pairs)
    
    phases_cdiff = [abs(phases_i[i])-abs(phases_c[i]) for i in range(len(t))]
    
    stat = phases_cdiff
    stat_av = np.average(stat)
    
    print('===============')
    print(f'For eps={eps} and dt={dt}: Mean={stat_av}')
    print('normaltest: pval='+str(normaltest(stat).pvalue))
    print('ttest: pval='+str(ttest_1samp(stat,0).pvalue))
    print('===============')
    
def full_stats(dim,eps,samplemethod,extremes,dt,ret='stats',noeff=1):
    from scipy.stats import ttest_rel, ttest_1samp, normaltest, wilcoxon
    import numpy as np
    import json
    import math as m
    from random import random
    from numpy.linalg import norm
    
    def resultant(sample):
        phase_x = [m.cos(ind) for ind in sample]
        phase_y = [m.sin(ind) for ind in sample]
        
        return (np.average(phase_x), np.average(phase_y))

    phase = np.vectorize(lambda x,y: (m.atan2(y,x)+m.pi)%(2*m.pi)-m.pi)
    diff = np.vectorize(lambda x,y: (y-x+m.pi)%(2*m.pi)-m.pi)
    H2 = lambda x: -x*m.log2(x)-(1-x)*m.log2(1-x)

    eps = round(eps,3)
    path='/cloud/Public/_data/neuropercolation/4lay/cons=27-knight_steps=1000100/dim=09/batch=0/'
    
    savepath = path + samplemethod + f'_samples={extremes}/dt={dt}/noeff={noeff:03d}/'

    try:
        with open(savepath+f"eps={round(eps,3):.3f}_phase_pairs.txt", 'r', encoding='utf-8') as f:
            phase_pairs = json.load(f)
        with open(savepath+f"eps={round(eps,3):.3f}_ei_pairs.txt", 'r', encoding='utf-8') as f:
            ei_pairs = json.load(f)
    except:
        sampling(dim,eps,samplemethod,extremes,dt,noeff=noeff)
        with open(savepath+f"eps={round(eps,3):.3f}_phase_pairs.txt", 'r', encoding='utf-8') as f:
            phase_pairs = json.load(f)
        with open(savepath+f"eps={round(eps,3):.3f}_ei_pairs.txt", 'r', encoding='utf-8') as f:
            ei_pairs = json.load(f)
    
    t, phases_i, phases_c = zip(*phase_pairs)
    t, ei_i, ei_c = zip(*ei_pairs)
    
    phases_cdiff = [abs(phases_i[i])-abs(phases_c[i]) if abs(abs(phases_i[i])-abs(phases_c[i]))!=m.pi else m.pi*(-1)**(random()<0.5)
                    for i in range(len(t))]
    
    print('phases with pi diff: '+str(len([i for i in range(len(t)) if abs(abs(phases_i[i])-abs(phases_c[i]))==m.pi])))
    
    stat = phases_cdiff
    mean = np.average(stat)
    std = np.std(stat)
    integral = sum([1 for val in stat if val<0])/len(stat)
    cohend = mean / (np.var(np.abs(phases_i),ddof=1)/2+np.var(np.abs(phases_c),ddof=1)/2)
    norm_p = normaltest(stat).pvalue
    ttest_p = ttest_1samp(stat,0,alternative='less').pvalue
    sign_p = wilcoxon(stat,alternative='less').pvalue
    print('===============')
    print(f'For eps={eps} and dt={dt}: Mean={mean}')
    print('normaltest: pval='+str(norm_p))
    print('ttest: pval='+str(ttest_p))
    print('sign: pval='+str(sign_p))
    print('===============')
    
    if ret=='stats':
        return mean,std,integral,cohend,norm_p,ttest_p,sign_p
    elif ret=='phases':
        return norm(resultant(phases_i)), norm(resultant(phases_c))