#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Aug 21 14:59:22 2023

@author: timofej
"""

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 random import random

from plot import qtplot

from neuropercolation import Simulate4Layers


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 coherence(sample):
    phase_x = [m.cos(ind) for ind in sample]
    
    return np.average(phase_x)


def bootstrap_stat(whole, sub, strength=10000, estimator='resultant'):
    k = len(sub)
    
    if estimator=='resultant':
        whole_est = norm(resultant(whole))
        sub_est = norm(resultant(sub))
        
        boot_dist = []
        for i in range(strength):
            boot_dist.append(norm(resultant(choose(whole,k))))
            if i%1000==999:
                print(f'Done {i:0{len(str(strength))}d} bootstraps')
        
        confidence = len([val for val in boot_dist if (val-whole_est)<max(sub_est-whole_est,0)])/len(boot_dist)
        
    return confidence

def bootstrap_res(whole, sub, strength=10000):
    k = len(sub)
    
    resultants = []
    for i in range(strength):
        resultants.append(norm(resultant(choose(whole,k))))
        if i%1000==999:
            print(f'Done {i:0{len(str(strength))}d} bootstraps')
    return resultants

def sample_parts(wholes, subs):
    ks = [len(s) for s in subs]
    sample = [pha for j in range(len(wholes)) for pha in choose(wholes[j],ks[j])]
            
    return sample

def bootstrap_parts(wholes, subs, strength=10000, estimator='resultant',sided='right'):
    ks = [len(s) for s in subs]
    
    whole = [w for whole in wholes for w in whole]
    sub   = [s for sub   in subs   for s in sub]
    
    if estimator=='resultant':
        whole_est = norm(resultant(whole))
        sub_est = norm(resultant(sub))
        
        boot_dist = []
        for i in range(strength):
            sample = [pha for j in range(len(wholes)) for pha in choose(wholes[j],ks[j])]
            boot_dist.append(norm(resultant(sample)))
            if i%1000==999:
                print(f'Done {i:0{len(str(strength))}d} bootstraps')
        
    if sided=='right':
        confidence = len([val for val in boot_dist if val<sub_est])/len(boot_dist)
    elif sided=='double':
        confidence = len([val for val in boot_dist if abs(val-whole_est)<abs(sub_est-whole_est)])/len(boot_dist)
    elif sided=='left':
        confidence = len([val for val in boot_dist if val>sub_est])/len(boot_dist)
    else:
        raise NotImplemented
    return confidence

def bootstrap_res_parts(wholes, subs, strength=10000):
    ks = [len(s) for s in subs]
    
    whole = [w for whole in wholes for w in whole]    
    whole_est = norm(resultant(whole))
    
    resultants = []
    for i in range(strength):
        sample = [pha for j in range(len(wholes)) for pha in choose(wholes[j],ks[j])]
        resultants.append(norm(resultant(sample)))
        if i%1000==999:
            print(f'Done {i:0{len(str(strength))}d} bootstraps')
        
    return resultants

def new_folder(path):
    if not os.path.exists(path):
        os.makedirs(path)
    return 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)

extremes = None
maxdt = 250

stp = 1000100
batch = 0

eps_space = list(np.linspace(0.01,0.5,50))

print(f'Started at {datetime.now()}')

for dim in [9]:
    for eps in eps_space[:1]:#list(eps_space[3:8])+list(eps_space[:3])+list(eps_space[8:]):
        eps = round(eps,3)
        path='/cloud/Public/_data/neuropercolation/4lay/cons=dimtimesdimby3_steps=100100/dim=09_cons=27/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)
        
        strength = 100000
        extremes = 10000 #[l//2 for l in lens]
        ext_rat = extremes/(stp-100)
        
        circparts = 100
        pha_dev = m.pi/circparts
        pha_max = np.max(np.abs(phase_diff))
        
        phase_in_part = lambda ph, i: abs(ph)<=pha_max/circparts if i==0 else i*pha_max/circparts<abs(ph)<=(i+1)*pha_max/circparts
        
        dev_parts = [sorted([i for i,val in enumerate(ei[:-maxdt]) if phase_in_part(phase_diff[i],j)],
                            key = lambda i: ei[i]) for j in range(circparts)]
        
        ext_fracs = [round(len(dev_parts[i])*ext_rat,6) for i in range(circparts)]
        ext_parts = [int(ext_fracs[i]) for i in range(circparts)]
        ext_probs = [round(ext_fracs[i]%1,6) for i in range(circparts)]
        
        exts = [ext_parts[i]+(random()<ext_probs[i]) for i in range(circparts)]
    
        top_parts = [dev_parts[i][-exts[i]:] if exts[i]>0 else [] for i in range(circparts)]
        
        pha_parts = [[phase_diff[i] for i in part] for part in dev_parts]
        syn_parts = [[phase_diff[i] for i in part] for part in top_parts]
        ran_sam = sample_parts(dev_parts, top_parts)
        top_sam = sample_parts(top_parts, top_parts)
        
        dev = []
        for part in dev_parts:
            dev.extend(part)
        top = []
        for part in top_parts:
            top.extend(part)
        
        assert sorted(top)==sorted(top_sam)
        
        ran_res = []
        top_res = []
        tot_res = []
        
        bot_ei = []
        top_ei = []
        
        maxdt = 250
        for dt in range(maxdt):
            tot_pha = phase_diff[dt:dt-maxdt]
            ran_pha = [(phase_diff[i+dt]) for i in ran_sam]
            top_pha = [(phase_diff[i+dt]) for i in top_sam]
            
            tot_res.append( norm(resultant(tot_pha)) )
            ran_res.append( norm(resultant(ran_pha)) )
            top_res.append( norm(resultant(top_pha)) )
            

        sampling = 'samedist_varmaxpha'
        plotpath = new_folder(path+f'{sampling}_causal_roll{pha_dev:.3f}/')#'bootstrap={strength}/'
        savepath=new_folder(plotpath+f'extremes={extremes}_bootstrength={strength}/')
        
        newpath = new_folder(path+f'extremes={extremes}_bootstrength={strength}/')
        resultant_stat = bootstrap_res(phase_diff,top,strength=strength)
        
        confs = []
        confdt = 250
        for dt in range(confdt+1):
            pha_evs = [[phase_diff[i+dt] for i in part] for part in dev_parts]
            syn_evs = [[phase_diff[i+dt] for i in part] for part in top_parts]
            phas = [phase_diff[i+dt] for i in top]
            res_pha = norm(resultant(phas))
            
            conf = len([val for val in resultant_stat if val<res_pha])/len(resultant_stat)
            confs.append((dt, conf))
            print(f'Confidence for dim={dim} dt={dt}: {conf}')
        with open(newpath+f"eps={round(eps,3):.3f}_dt={confdt}_simpbootstrap.txt", 'w', encoding='utf-8') as f:
            json.dump(list(confs), f, indent=1)

        qtplot(f'Diachronic resultant for dim={dim} eps={eps:.3f} with 4 layers',
                [np.array(range(maxdt))]*3,
                [tot_res, ran_res, top_res],
                ['Average Resultant', f'random sample of {extremes}',
                    f'top sample of {extremes} ei'],
                x_tag = 'dt',
                y_tag = 'concentration',
                export=True,
                path=newpath,
                filename=f'Random Diachronic Resultant Norm eps={round(eps,3):.3f} dim={dim} extremes={extremes}.png',
                close=True)
        
        # qtplot(f'Diachronic resultant phase for dim={dim} eps={eps:.3f} with 4 layers',
        #         [np.array(range(maxdt))]*3,
        #         [top_ph[0], dis_ph[0], tot_ph],
        #         ['top {extremes} ei', 'dis {extremes} ei',
        #          'Average'],
        #         x_tag = 'dt',
        #         y_tag = 'phase',
        #         export=True,
        #         path=plotpath,
        #         filename=f'All Diachronic Resultant Phase eps={round(eps,3):.3f} dim={dim} extremes={extremes} roll{pha_dev:.3f}.png',
        #         close=True)
        
        # qtplot(f'Diachronic ei for dim={dim} with 4 layers',
        #         [np.array(range(maxdt))]*4,
        #         [bot_ei, top_ei, dev_ei, tot_ei],
        #         ['EI ev of bottom {extremes} ei', 'EI ev of top {extremes} ei',
        #          'EI ev of phase filtered ei', 'Average EI'],
        #         x_tag = 'dt',
        #         y_tag = 'average ei',
        #         export=True,
        #         path=path+'plots/',
        #         filename=f'Diachronic EI balanced for eps={round(eps,3):.3f} dim={dim} extremes={extremes} roll{pha_dev:.3f}.png',
        #         close=True)
        
        print(f'Done eps={eps:.3f} with dim={dim} at {datetime.now()}')
        
    # qtplot(f'Resultant and EI evolution for dim={dim} with 4 layers',
    #         [[0]+eps_space]*2,
    #         [max(av_ei)*diff_res, av_ei],
    #         ['Resultant', 'avEI'],
    #         export=True,
    #         path=path,
    #         filename=f'Resultant and EI for dim={dim}.png',
    #         close=True)