#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Sep  9 22:23:30 2023

@author: astral
"""

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[4:]:
        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]
        
        ei_ind = [i for i,val in enumerate(ei[:-maxdt]) if val>0]
        
        print(f'{len(ei_ind)} states with positive EI')
        
        samples = choose(ei_ind, extremes)
        sampling = 'allpos_ei'
        
        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]
        
        phase_pairs = []
        ei_pairs = []
        for num,i in enumerate(samples):
            causal_init = states[i].translate(str.maketrans('', '', '.-='))
            
            sim = Simulate4Layers(  dim,
                                    eps,
                                    coupling=coupling,
                                    init=causal_init,
                                    noeffect=0,
                                    steps=1,
                                    draw=None,
                                    )
                
            activation = sim._activations()
            channel = sim._channels()
        
            osc = list(zip(*activation))
            phase_abs = np.array([[np.arctan2(*act[::-1]) for act in osc[i]] for i in range(2)])
            phasediff_c = np.round(diff(phase_abs[0],phase_abs[1]),6)
            ei_c = [round(np.sum(cha)*(1-H2(eps)),6) for cha in channel]
            max_ei_c = max([np.sum(cha) for cha in channel])

            sim = Simulate4Layers(  dim,
                                    eps,
                                    coupling=coupling,
                                    init=causal_init,
                                    noeffect=-1,
                                    steps=1,
                                    draw=None,
                                    )
                
            activation = sim._activations()
            channel = sim._channels()
            
            osc = list(zip(*activation))
            phase_abs = np.array([[np.arctan2(*act[::-1]) for act in osc[i]] for i in range(2)])
            phasediff_i = np.round(diff(phase_abs[0],phase_abs[1]),6)
            ei_i = [round(np.sum(cha)*(1-H2(eps)),6) for cha in channel]
            max_ei_i = max([np.sum(cha) for cha in channel])

            phase_pairs.append((phasediff_i[-1], phasediff_c[-1]))
            ei_pairs.append((ei_i[-1], ei_c[-1]))
            
            savepath = path + sampling + '/'
            new_folder(savepath)
            
            if num%100==99:
                print(f'Done {num:0{len(str(extremes))}d}')
        
        with open(savepath+f"eps={round(eps,3):.3f}_phase_pairs.txt", 'w', encoding='utf-8') as f:
            json.dump(phase_pairs, f, indent=1)
        with open(savepath+f"eps={round(eps,3):.3f}_ei_pairs.txt", 'w', encoding='utf-8') as f:
            json.dump(ei_pairs, f, indent=1)
            
        phases_i, phases_c = zip(*phase_pairs)
        ei_i, ei_c = zip(*ei_pairs)
        
        phase_space = np.linspace(0,m.pi,100+1)
        ei_space = np.linspace(0,np.max([ei_i,ei_c]),100+1)
        
        phase_dist_i = [len([ph for ph in phases_i if low<=ph<high])/extremes for low,high in zip(phase_space[:-1],phase_space[1:])]
        phase_dist_c = [len([ph for ph in phases_c if low<=ph<high])/extremes for low,high in zip(phase_space[:-1],phase_space[1:])]
        
        max_ei = max(max_ei_i, max_ei_c)
        
        ei_dist_i = [len([e for e in ei_i if round(e/(1-H2(eps)))==val])/extremes for val in range(max_ei)]
        ei_dist_c = [len([e for e in ei_c if round(e/(1-H2(eps)))==val])/extremes for val in range(max_ei)]

        qtplot(f'Phase distribution for dim={dim} eps={eps:.3f} with 4 layers',
                [phase_space[:-1]]*2,
                [phase_dist_i, phase_dist_c],
                ['Phase dist with ei',
                 'Phase dist without ei'],
                x_tag = 'phase',
                y_tag = 'density',
                export=True,
                path=savepath,
                filename=f'Phase dist eps={round(eps,3):.3f} dim={dim} extremes={extremes}.png',
                close=True)
        
        qtplot(f'EI distribution for dim={dim} eps={eps:.3f} with 4 layers',
                [range(max_ei)]*2,
                [ei_dist_i, ei_dist_c],
                ['EI dist with ei',
                 'EI dist without ei'],
                x_tag = 'ei',
                y_tag = 'density',
                export=True,
                path=savepath,
                filename=f'EI dist eps={round(eps,3):.3f} dim={dim} extremes={extremes}.png',
                close=True)
      
        print(f'Done eps={eps:.3f} with dim={dim} at {datetime.now()}')