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

@author: astral
"""

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

phase = lambda x,y: (m.atan2(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

for dim in [7]:
    for eps in eps_space:
        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/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)
        
        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_in_ind = [i for i in dev_ind if not from_sync(i)]
        dev_out_ind = [i for i in dev_ind if from_sync(i)]
        
        devm_in_ind = [i for i in dev_in_ind if phase_diff[i]<0]
        devp_in_ind = [i for i in dev_in_ind if phase_diff[i]>0]
        devm_out_ind = [i for i in dev_out_ind if phase_diff[i]<0]
        devp_out_ind = [i for i in dev_out_ind if phase_diff[i]>0]
        
        if extremes is None:
            extremes = min(len(dev_in_ind),len(dev_out_ind))//1000*100

        print(len(dev_out_ind), len(dev_in_ind))
        #print(all_res, av_diff)
        
        
        bot_in_ind = dev_in_ind[         :extremes]
        top_in_ind = dev_in_ind[-extremes:        ]
        bot_out_ind = dev_out_ind[         :extremes]
        top_out_ind = dev_out_ind[-extremes:        ]
        
        bot_ind = bot_in_ind + bot_out_ind
        top_ind = top_in_ind + top_out_ind
        
        bot_res = []
        top_res = []
        dev_res = []
        tot_res = []
        
        bot_ph = []
        top_ph = []
        dev_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_ind]
            top_pha = [abs(phase_diff[i+dt]) for i in top_ind]
            dev_pha = [abs(phase_diff[i+dt]) for i in dev_ind]
            tot_pha = np.abs(phase_diff[dt:dt-maxdt])
            
            bot_res.append( norm(resultant(bot_pha)) )
            top_res.append( norm(resultant(top_pha)) )
            dev_res.append( norm(resultant(dev_pha)) )
            tot_res.append( norm(resultant(tot_pha)) )
            
            bot_ph.append( phase(*resultant(bot_pha)) )
            top_ph.append( phase(*resultant(top_pha)) )
            dev_ph.append( phase(*resultant(dev_pha)) )
            tot_ph.append( phase(*resultant(tot_pha)) )
            
            bot_ei.append( np.average([ei[i+dt] for i in bot_ind]) )
            top_ei.append( np.average([ei[i+dt] for i in top_ind]) )
            dev_ei.append( np.average([ei[i+dt] for i in dev_ind]) )
            tot_ei.append( np.average(ei[dt:dt-maxdt]) )
            
            if dt%100==99:
                print(f'Done dt={dt}')
            
        qtplot(f'Diachronic resultant for dim={dim} with 4 layers',
                [np.array(range(maxdt))]*4,
                [bot_res, top_res, dev_res, tot_res],
                ['Resultant ev of bottom {extremes} ei', 'Resultant ev of top {extremes} ei',
                  'Resultant ev of phase filtered ei', 'Average Resultant'],
                x_tag = 'dt',
                y_tag = 'concentration',
                export=True,
                path=path+'plots/',
                filename=f'Diachronic Resultant Norm balanced eps={round(eps,3):.3f} dim={dim} extremes={extremes} roll{pha_dev:.3f}.png',
                close=True)
        
        qtplot(f'Diachronic resultant phase for dim={dim} with 4 layers',
                [np.array(range(maxdt))]*4,
                [bot_ph, top_ph, dev_ph, tot_ph],
                ['bottom {extremes} ei', 'top {extremes} ei',
                  'all filtered ei', 'Average'],
                x_tag = 'dt',
                y_tag = 'concentration',
                export=True,
                path=path+'plots/',
                filename=f'Diachronic Resultant Phase balanced 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)