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

@author: timofej
"""

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))

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 = 300

stp = 10100
batch = 0

for dim in [9]:
    for eps in eps_space[1:41:2]:
        path=f'/cloud/Public/_data/neuropercolation/4lay/cons=27-knight_steps={stp}/dim={dim:02}/batch={batch}/'
        
        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])
            
            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 = [np.sum(cha)*(1-H2(eps)) 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)
        
        pha_center = av_diff
        pha_dev = m.pi/32
        
        from_sync = lambda i: True if abs(phase_diff[i])<0.08*m.pi else False if 0.42*m.pi<abs(phase_diff[i])<0.58*m.pi else from_sync(i-1) if i>0 else None
        to_sync = lambda i: True if abs(phase_diff[i])<0.08*m.pi else False if 0.42*m.pi<abs(phase_diff[i])<0.58*m.pi else to_sync(i+1) if i+1<len(phase_diff) else None
        
        dev_ind = sorted([i for i,val in enumerate(ei[:-maxdt]) if (pha_center-pha_dev)<=abs(phase_diff[i])<=(pha_center+pha_dev)], key = lambda i: ei[i])
        dev_00 = [i for i in dev_ind if from_sync(i)          and to_sync(i)         ]
        dev_01 = [i for i in dev_ind if from_sync(i)          and to_sync(i) is False]
        dev_10 = [i for i in dev_ind if from_sync(i) is False and to_sync(i)         ]
        dev_11 = [i for i in dev_ind if from_sync(i) is False and to_sync(i) is False]
        
        lens = [len(dev_00),len(dev_01),len(dev_10),len(dev_11)]
        
        #if not extremes:
        extremes = [100]*4 #[l//2 for l in lens]

        print(lens)
        #print(all_res, av_diff)
        
        
        # bot_00 = dev_00[:extremes[0]]
        # bot_01 = dev_01[:extremes[1]]
        # bot_10 = dev_10[:extremes[2]]
        # bot_11 = dev_11[:extremes[3]]
        
        top_00 = dev_00[-extremes[0]:]
        top_01 = dev_01[-extremes[1]:]
        top_10 = dev_10[-extremes[2]:]
        top_11 = dev_11[-extremes[3]:]
        
        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]
        
        for top,name in [(top_00,'top_00'),(top_01,'top_01'),(top_10,'top_10'),(top_11,'top_11')]:
            for i in top:
                causal_init = states[i].translate(str.maketrans('', '', '.-='))
                path_c = path+f'causal_maxdt={maxdt}/{name}/{i:0{len(str(stp))}}/'
                
                try:
                    with open(path_c+f'eps={round(eps,3):.3f}_phase_diff.txt', 'r', encoding='utf-8') as f:
                        phasediff = json.load(f)
                except:
                    sim=Simulate4Layers(dim,
                                        eps,
                                        coupling=coupling,
                                        init=causal_init,
                                        noeffect=0,
                                        steps=maxdt,
                                        draw=None,
                                        save='all',
                                        path=path_c,
                                        )
                    
                    with open(path_c+f"eps={round(eps,3):.3f}_activation.txt", 'r', encoding='utf-8') as f:
                        activation = json.load(f)
                
                    osc = list(zip(*activation))
                    phase_abs = np.array([[np.arctan2(*act[::-1]) for act in osc[i]] for i in range(2)])
                    phasediff = diff(phase_abs[0],phase_abs[1])
                    
                    with open(path_c+f"eps={round(eps,3):.3f}_phase_diff.txt", 'w', encoding='utf-8') as f:
                        json.dump(list(phasediff), f, indent=1)
        
        for top,name in [(top_00,'top_00'),(top_01,'top_01'),(top_10,'top_10'),(top_11,'top_11')]:
            for i in top:
                causal_init = states[i].translate(str.maketrans('', '', '.-='))
                path_c = path+f'original_maxdt={maxdt}/{name}/{i:0{len(str(stp))}}/'
                
                try:
                    with open(path_c+f'eps={round(eps,3):.3f}_phase_diff.txt', 'r', encoding='utf-8') as f:
                        phasediff = json.load(f)
                except:
                    sim=Simulate4Layers(dim,
                                        eps,
                                        coupling=coupling,
                                        init=causal_init,
                                        noeffect=-1,
                                        steps=maxdt,
                                        draw=None,
                                        save='all',
                                        path=path_c,
                                        )
                    
                    with open(path_c+f"eps={round(eps,3):.3f}_activation.txt", 'r', encoding='utf-8') as f:
                        activation = json.load(f)
                
                    osc = list(zip(*activation))
                    phase_abs = np.array([[np.arctan2(*act[::-1]) for act in osc[i]] for i in range(2)])
                    phasediff = diff(phase_abs[0],phase_abs[1])
                    
                    with open(path_c+f"eps={round(eps,3):.3f}_phase_diff.txt", 'w', encoding='utf-8') as f:
                        json.dump(list(phasediff), f, indent=1)
        
            
        
        
        
        # bot_res = []
        top_res = []
        dis_res = []
        tot_res = []
        
        # bot_ph = []
        top_ph = []
        dis_ph = []
        tot_ph = []
        
        # bot_ei = []
        top_ei = []
        # dev_ei = []
        tot_ei = []
        for dt in range(maxdt):
            # bot_pha = [[abs(phase_diff[i+dt]) for i in bot_00],
            #            [abs(phase_diff[i+dt]) for i in bot_01],
            #            [abs(phase_diff[i+dt]) for i in bot_10],
            #            [abs(phase_diff[i+dt]) for i in bot_11]]
            
            top_pha = [[abs(phase_diff[i+dt]) for i in top_00],
                       [abs(phase_diff[i+dt]) for i in top_01],
                       [abs(phase_diff[i+dt]) for i in top_10],
                       [abs(phase_diff[i+dt]) for i in top_11]]
            
            dis_00 = []
            dis_01 = []
            dis_10 = []
            dis_11 = []
            for i in top_00:
                path_c = path+f'causal_maxdt={maxdt}/top_00/{i:0{len(str(stp))}}/'
                with open(path_c+f'eps={round(eps,3):.3f}_phase_diff.txt', 'r', encoding='utf-8') as f:
                    dis_00.append(abs(json.load(f)[dt]))
            for i in top_01:
                path_c = path+f'causal_maxdt={maxdt}/top_01/{i:0{len(str(stp))}}/'
                with open(path_c+f'eps={round(eps,3):.3f}_phase_diff.txt', 'r', encoding='utf-8') as f:
                    dis_01.append(abs(json.load(f)[dt]))
            for i in top_10:
                path_c = path+f'causal_maxdt={maxdt}/top_10/{i:0{len(str(stp))}}/'
                with open(path_c+f'eps={round(eps,3):.3f}_phase_diff.txt', 'r', encoding='utf-8') as f:
                    dis_10.append(abs(json.load(f)[dt]))
            for i in top_11:
                path_c = path+f'causal_maxdt={maxdt}/top_11/{i:0{len(str(stp))}}/'
                with open(path_c+f'eps={round(eps,3):.3f}_phase_diff.txt', 'r', encoding='utf-8') as f:
                    dis_11.append(abs(json.load(f)[dt]))
            
            dis_pha = [dis_00, dis_01, dis_10, dis_11]
            
            tot_pha = np.abs(phase_diff[dt:dt-maxdt])
            
            # bot_res.append( [norm(resultant(bot_pha[i])) for i in range(4)] )
            top_res.append( [norm(resultant(top_pha[i])) for i in range(4)] )
            dis_res.append( [norm(resultant(dis_pha[i])) for i in range(4)] )
            tot_res.append( norm(resultant(tot_pha)) )
            
            # bot_ph.append( [phase(*resultant(bot_pha[i])) for i in range(4)] )
            top_ph.append( [phase(*resultant(top_pha[i])) for i in range(4)] )
            dis_ph.append( [phase(*resultant(dis_pha[i])) for i in range(4)] )
            tot_ph.append( phase(*resultant(tot_pha)) )
            
            # bot_ei.append( [np.average([ei[i+dt] for i in bot]) for bot in [bot_00,bot_01,bot_10,bot_11]] )
            top_ei.append( [np.average([ei[i+dt] for i in top]) for top in [top_00,top_01,top_10,top_11]] )
            # dev_ei.append( [np.average([ei[i+dt] for i in dev]) for dev in [dev_00,dev_01,dev_10,dev_11]] )
            tot_ei.append( np.average(ei[dt:dt-maxdt]) )
            
            if dt%10==0:
                print(f'Done dt={dt}')
        
        # bot_res = list(zip(*bot_res))
        top_res = list(zip(*top_res))
        dis_res = list(zip(*dis_res))
        
        # bot_ph = list(zip(*bot_ph))
        top_ph = list(zip(*top_ph))
        dis_ph = list(zip(*dis_ph))

        # bot_ei = list(zip(*bot_ei))
        top_ei = list(zip(*top_ei))
        # dev_ei = list(zip(*dev_ei))
        
        plotpath = path+'4waycausal/'

        qtplot(f'Diachronic resultant sync to sync for dim={dim} with 4 layers',
                [np.array(range(maxdt))]*3,
                [top_res[0], dis_res[0], tot_res],
                
                ['sync to sync top {extremes} ei',
                 'sync to sync dis {extremes} ei', 'Average Resultant'],
                x_tag = 'dt',
                y_tag = 'concentration',
                export=True,
                path=plotpath,
                filename=f'Diachronic Resultant Norm eps={round(eps,3):.3f} sts dim={dim} extremes={extremes} roll{pha_dev:.3f}.png',
                close=True)
        qtplot(f'Diachronic resultant sync to orth for dim={dim} with 4 layers',
                [np.array(range(maxdt))]*3,
                [top_res[1], dis_res[1], tot_res],
                ['sync to orth top {extremes} ei', 'sync to orth dis {extremes} ei',
                 'Average Resultant'],
                x_tag = 'dt',
                y_tag = 'concentration',
                export=True,
                path=plotpath,
                filename=f'Diachronic Resultant Norm eps={round(eps,3):.3f} sto dim={dim} extremes={extremes} roll{pha_dev:.3f}.png',
                close=True)
        qtplot(f'Diachronic resultant orth to sync for dim={dim} with 4 layers',
                [np.array(range(maxdt))]*3,
                [top_res[2], dis_res[2], tot_res],
                ['orth to sync top {extremes} ei', 'orth to sync dis {extremes} ei',
                 'Average Resultant'],
                x_tag = 'dt',
                y_tag = 'concentration',
                export=True,
                path=plotpath,
                filename=f'Diachronic Resultant Norm eps={round(eps,3):.3f} ots dim={dim} extremes={extremes} roll{pha_dev:.3f}.png',
                close=True)
        qtplot(f'Diachronic resultant orth to orth for dim={dim} with 4 layers',
                [np.array(range(maxdt))]*3,
                [top_res[3], dis_res[3], tot_res],
                ['orth to orth top {extremes} ei', 'orth to orth dis {extremes} ei',
                 'Average Resultant'],
                x_tag = 'dt',
                y_tag = 'concentration',
                export=True,
                path=plotpath,
                filename=f'Diachronic Resultant Norm eps={round(eps,3):.3f} oto dim={dim} extremes={extremes} roll{pha_dev:.3f}.png',
                close=True)
        
        qtplot(f'Diachronic resultant phase sync to sync for dim={dim} with 4 layers',
                [np.array(range(maxdt))]*3,
                [top_ph[0], dis_ph[0], tot_ph],
                ['sync to sync top {extremes} ei', 'sync to sync dis {extremes} ei',
                 'Average'],
                x_tag = 'dt',
                y_tag = 'phase',
                export=True,
                path=plotpath,
                filename=f'Diachronic Resultant Phase eps={round(eps,3):.3f} sts dim={dim} extremes={extremes} roll{pha_dev:.3f}.png',
                close=True)
        qtplot(f'Diachronic resultant phase sync to orth for dim={dim} with 4 layers',
                [np.array(range(maxdt))]*3,
                [top_ph[1], dis_ph[1], tot_ph],
                ['sync to orth top {extremes} ei', 'sync to orth dis {extremes} ei',
                 'Average'],
                x_tag = 'dt',
                y_tag = 'phase',
                export=True,
                path=plotpath,
                filename=f'Diachronic Resultant Phase eps={round(eps,3):.3f} sto dim={dim} extremes={extremes} roll{pha_dev:.3f}.png',
                close=True)
        qtplot(f'Diachronic resultant phase orth to sync for dim={dim} with 4 layers',
                [np.array(range(maxdt))]*3,
                [top_ph[2], dis_ph[2], tot_ph],
                ['orth to sync top {extremes} ei', 'orth to sync dos {extremes} ei',
                 'Average'],
                x_tag = 'dt',
                y_tag = 'phase',
                export=True,
                path=plotpath,
                filename=f'Diachronic Resultant Phase eps={round(eps,3):.3f} ots dim={dim} extremes={extremes} roll{pha_dev:.3f}.png',
                close=True)
        qtplot(f'Diachronic resultant phase orth to orth for dim={dim} with 4 layers',
                [np.array(range(maxdt))]*3,
                [top_ph[3], dis_ph[3], tot_ph],
                ['orth to orth top {extremes} ei', 'orth to orth dis {extremes} ei',
                 'Average'],
                x_tag = 'dt',
                y_tag = 'phase',
                export=True,
                path=plotpath,
                filename=f'Diachronic Resultant Phase eps={round(eps,3):.3f} oto 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)