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

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 H2(x):
    return -x*m.log2(x)-(1-x)*m.log2(1-x)

extremes = None
maxdt = 300

for dim in [9]:
    for eps in eps_space[:10]:
        path=f'/cloud/Public/_data/neuropercolation/4lay/cons=7-knight_steps=1000100/dim=07/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 = np.array([[np.arctan2(*act[::-1]) for act in osc[i]] for i in range(2)])
            phase_diff = (phase[1]-phase[0]+m.pi)%(2*m.pi)-m.pi
            
            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
        orth_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])
        
        thres = 0.04*m.pi
        last_sync = lambda i: i if abs(phase_diff[i])<0.08*m.pi else last_sync(i-1) 
        last_orth = lambda i: i if 0.42*m.pi<abs(phase_diff[i])<0.58*m.pi else last_orth(i-1) 
        
        trav_out = [(i,pha) for i,pha in enumerate(phase_diff[:-maxdt]) if abs(pha)<thres and abs(phase_diff[i+1])>thres]
        diff00_ind = [last_sync(i) for i,pha in trav_out]
        diff00m_ind = [last_sync(i) for i,pha in trav_out if pha<0]
        diff00p_ind = [last_sync(i) for i,pha in trav_out if pha>0]
        
        trav_in = [(i,pha) for i,pha in enumerate(phase_diff[:-maxdt]) if m.pi/2-thres<abs(pha)<m.pi/2+thres and abs(phase_diff[i+1])<m.pi/2-thres]
        diff90_ind = [last_orth(i) for i,pha in trav_in]
        diff90m_ind = [last_orth(i) for i,pha in trav_in if pha<0]
        diff90p_ind = [last_orth(i) for i,pha in trav_in if pha>0]
        
        print(f'{len(diff00_ind)}/{len(trav_out)} synced states and {len(diff90_ind)}/{len(trav_in)} orthogonal states')
        #print(all_res, av_diff)
        
        in_ei = []
        out_ei = []
        
        inm_res = []
        inp_res = []
        outm_res = []
        outp_res = []
        for dt in range(maxdt):
            inm_res.append( resultant([phase_diff[i+dt] for i in diff90m_ind]) )
            outm_res.append( resultant([phase_diff[i+dt] for i in diff00m_ind]) )
            
            inp_res.append( resultant([phase_diff[i+dt] for i in diff90p_ind]) )
            outp_res.append( resultant([phase_diff[i+dt] for i in diff00p_ind]) )
            
            in_ei.append( np.average([ei[i+dt] for i in diff90_ind]) )
            out_ei.append( np.average([ei[i+dt] for i in diff00_ind]) )
            
            if dt%100==99:
                print(f'Done dt={dt}')
            
        im_res = list(zip(*inm_res))
        om_res = list(zip(*outm_res))
        
        ip_res = list(zip(*inp_res))
        op_res = list(zip(*outp_res))
        
        space = np.linspace(0,2*m.pi,101)
        
        av_ei = np.average(ei)
        
        qtplot(f'Transitions resultants for eps={round(eps,3):.3f} with dim={dim} with 4 layers',
                [im_res[0], om_res[0], ip_res[0], op_res[0], np.cos(space)],
                [im_res[1], om_res[1], ip_res[1], op_res[1], np.sin(space)],
                ['Resultant ev inward -', 'Resultant ev outward -',
                 'Resultant ev inward +', 'Resultant ev outward +',
                 'Unit Circle'],
                x_tag = 'x',
                y_tag = 'y',
                lw=1,
                export=True,
                path=path,
                filename=f'Transitions resultant pm for eps={round(eps,3):.3f} dim={dim}.png',
                close=False)
        
        qtplot(f'Transitions ei for eps={round(eps,3):.3f} with dim={dim} with 4 layers',
                [np.array(range(maxdt))]*3,
                [in_ei, out_ei, [av_ei]*maxdt],
                ['EI ev inward', 'EI ev outward', 'Average EI'],
                x_tag = 'dt',
                y_tag = 'average ei',
                export=True,
                path=path,
                filename=f'Transitions EI for eps={round(eps,3):.3f} dim={dim} .png',
                close=False)
        
        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)