#!/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 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)

extremes = None
maxdt = 200

stp = 1000100
batch = 0

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

for dim in [9]:
    for eps in eps_space[14:]:
        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)
        
        extremes = 10000 #[l//2 for l in lens]
        ext_rat = extremes/(stp-100)
        
        circparts = 32
        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_parts = [int(np.ceil(len(dev_parts[i])*ext_rat)) for i in range(circparts)]
        
        top_parts = [dev_parts[i][-ext_parts[i]:] if ext_parts[i]>0 else [] for i in range(circparts)]
        bot_parts = [dev_parts[i][:ext_parts[i]] if ext_parts[i]>0 else [] for i in range(circparts)]

        top = []
        for part in top_parts:
            top.extend(part)
        bot = []
        for part in bot_parts:
            bot.extend(part)
        
        print(len(top), len(bot), extremes, 'equal?')
        sampling = 'samedist_varmaxpha'
    
        # pha_center = av_diff
        # pha_dev = m.pi/8
        
        # 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
        
        # infra_phase = lambda ph: (pha_center-pha_dev)<=abs(ph)<=(pha_center        )
        # supra_phase = lambda ph: (pha_center        )< abs(ph)<=(pha_center+pha_dev)

        # dev_inf = sorted([i for i,val in enumerate(ei[:-maxdt]) if infra_phase(phase_diff[i])], key = lambda i: ei[i])
        # dev_sup = sorted([i for i,val in enumerate(ei[:-maxdt]) if supra_phase(phase_diff[i])], key = lambda i: ei[i])
        
        # ext_inf = round(extremes*len(dev_inf)/(len(dev_inf)+len(dev_sup)))
        # ext_sup = round(extremes*len(dev_sup)/(len(dev_inf)+len(dev_sup)))

        # top_inf = dev_inf[-ext_inf:]
        # top_sup = dev_sup[-ext_sup:]
        
        # top = top_inf + top_sup
        
        # print(len(top), extremes, 'equal?')
        # sampling = 'biequal'
        
        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 i in top:
            causal_maxdt=0
            for file in os.listdir(path):
                f = os.path.join(path, file)
                if not os.path.isfile(f):
                    c_maxdt = file.replace('causal_maxdt=','')
                    if c_maxdt != file:
                        causal_maxdt = int(c_maxdt)
            
            if causal_maxdt>=maxdt:
                path_c = path+f'causal_maxdt={causal_maxdt}/eps={round(eps,3):.3f}/{i:0{len(str(stp))}}/'
            else:
                path_c = path+f'causal_maxdt={maxdt}/eps={round(eps,3):.3f}/{i:0{len(str(stp))}}/'
            
            causal_init = states[i].translate(str.maketrans('', '', '.-='))
            
            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='simple',
                                    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 i in []:#top:
            original_maxdt=0
            for file in os.listdir(path):
                f = os.path.join(path, file)
                if not os.path.isfile(f):
                    o_maxdt = file.replace('original_maxdt=','')
                    if c_maxdt != file:
                        original_maxdt = int(o_maxdt)
            
            if original_maxdt>=maxdt:
                path_c = path+f'original_maxdt={original_maxdt}/eps={round(eps,3):.3f}/{i:0{len(str(stp))}}/'
            else:
                path_c = path+f'original_maxdt={maxdt}/eps={round(eps,3):.3f}/{i:0{len(str(stp))}}/'
            
            causal_init = states[i].translate(str.maketrans('', '', '.-='))
            
            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='simple',
                                    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])
                phasediff = [round(pha,6) for pha in phasediff]
                
                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 = []
        dis_ei = []
        tot_ei = []
        
        
        present_dt=0
        for file in os.listdir(path):
            f = os.path.join(path, file)
            if not os.path.isfile(f):
                p_maxdt = file.replace(f'{sampling}_causal_roll{pha_dev:.3f}_maxdt=','')
                if p_maxdt != file:
                    present_dt = int(p_maxdt)
        
        if present_dt>0:
            try:
                datapath = path + f"{sampling}_causal_roll{pha_dev:.3f}_maxdt={present_dt}/data/"
                with open(datapath+f"eps={round(eps,3):.3f}_bot_dia_res.txt", 'r', encoding='utf-8') as f:
                    bot_res = json.load(f)
                with open(datapath+f"eps={round(eps,3):.3f}_top_dia_res.txt", 'r', encoding='utf-8') as f:
                    top_res = json.load(f)
                with open(datapath+f"eps={round(eps,3):.3f}_dis_dia_res.txt", 'r', encoding='utf-8') as f:
                    dis_res = json.load(f)
                with open(datapath+f"eps={round(eps,3):.3f}_bot_dia_ei.txt", 'r', encoding='utf-8') as f:
                    bot_ei = json.load(f)
                with open(datapath+f"eps={round(eps,3):.3f}_top_dia_ei.txt", 'r', encoding='utf-8') as f:
                    top_ei = json.load(f)
                with open(datapath+f"eps={round(eps,3):.3f}_dis_dia_ei.txt", 'r', encoding='utf-8') as f:
                    dis_ei = json.load(f)
            except:
                present_dt=0
                bot_res = []
                top_res = []
                dis_res = []
        
        for dt in range(present_dt,maxdt):
            top_pha = [(phase_diff[i+dt]) for i in top]
            bot_pha = [(phase_diff[i+dt]) for i in bot]
            
            dis = []
            for i in top:
                causal_maxdt=0
                for file in os.listdir(path):
                    f = os.path.join(path, file)
                    if not os.path.isfile(f):
                        c_maxdt = file.replace('causal_maxdt=','')
                        if c_maxdt != file:
                            causal_maxdt = int(c_maxdt)
                
                if causal_maxdt>=maxdt:
                    path_c = path+f'causal_maxdt={causal_maxdt}/eps={round(eps,3):.3f}/{i:0{len(str(stp))}}/'
                else:
                    path_c = path+f'causal_maxdt={maxdt}/eps={round(eps,3):.3f}/{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.append((json.load(f)[dt]))
            
            dei = []
            for i in top:
                causal_maxdt=0
                for file in os.listdir(path):
                    f = os.path.join(path, file)
                    if not os.path.isfile(f):
                        c_maxdt = file.replace('causal_maxdt=','')
                        if c_maxdt != file:
                            causal_maxdt = int(c_maxdt)
                
                if causal_maxdt>=maxdt:
                    path_c = path+f'causal_maxdt={causal_maxdt}/eps={round(eps,3):.3f}/{i:0{len(str(stp))}}/'
                else:
                    path_c = path+f'causal_maxdt={maxdt}/eps={round(eps,3):.3f}/{i:0{len(str(stp))}}/'
                
                with open(path_c+f'eps={round(eps,3):.3f}_channels.txt', 'r', encoding='utf-8') as f:
                    dei.append(np.sum(json.load(f)[dt])*(1-H2(eps)))
            
            dis_pha = dis
            tot_pha = (phase_diff[dt:dt-maxdt])
            
            bot_res.append( [norm(resultant(bot_pha))] )
            top_res.append( [norm(resultant(top_pha))] )
            dis_res.append( [norm(resultant(dis_pha))] )
            tot_res.append( norm(resultant(tot_pha)) )
            
            # bot_ph.append( [phase(*resultant(bot_pha[i])) for i in range(1)] )
            # top_ph.append( [phase(*resultant(top_pha[i])) for i in range(1)] )
            # dis_ph.append( [phase(*resultant(dis_pha[i])) for i in range(1)] )
            # tot_ph.append( phase(*resultant(tot_pha)) )
            
            bot_ei.append( [np.average([ei[i+dt] for i in bot])] )
            top_ei.append( [np.average([ei[i+dt] for i in top])] )
            dis_ei.append( [np.average(dei)] )
            tot_ei.append( np.average(ei[dt:dt-maxdt]) )
            
            if dt%10==9:
                print(f'Done dt={dt:{len(str(maxdt))}d}')
        
        plotpath = path+f'{sampling}_causal_roll{pha_dev:.3f}_maxdt={maxdt}/'
        new_folder(plotpath+'data/')
        
        with open(plotpath+f"data/eps={round(eps,3):.3f}_bot_dia_res.txt", 'w', encoding='utf-8') as f:
            json.dump(list(bot_res), f, indent=1)
        with open(plotpath+f"data/eps={round(eps,3):.3f}_top_dia_res.txt", 'w', encoding='utf-8') as f:
            json.dump(list(top_res), f, indent=1)
        with open(plotpath+f"data/eps={round(eps,3):.3f}_dis_dia_res.txt", 'w', encoding='utf-8') as f:
            json.dump(list(dis_res), f, indent=1)
        with open(plotpath+f"data/eps={round(eps,3):.3f}_bot_dia_ei.txt", 'w', encoding='utf-8') as f:
            json.dump(list(bot_ei), f, indent=1)
        with open(plotpath+f"data/eps={round(eps,3):.3f}_top_dia_ei.txt", 'w', encoding='utf-8') as f:
            json.dump(list(top_ei), f, indent=1)
        with open(plotpath+f"data/eps={round(eps,3):.3f}_dis_dia_ei.txt", 'w', encoding='utf-8') as f:
            json.dump(list(dis_ei), f, indent=1)
        
        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))
        dis_ei = list(zip(*dis_ei))
        
        

        qtplot(f'Diachronic resultant for dim={dim} eps={eps:.3f} with 4 layers',
                [np.array(range(maxdt))]*4,
                [tot_res, bot_res[0], dis_res[0], top_res[0]],
                ['Average Resultant', f'bottom {extremes} ei'
                 f'dis {extremes} ei', f'top {extremes} ei'],
                x_tag = 'dt',
                y_tag = 'concentration',
                export=True,
                path=plotpath,
                filename=f'All 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)