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

from plot import qtplot

eps_space = list(np.linspace(0.005,0.5,100))
eps_space= eps_space[1::2]

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)

dims=list(range(3,10))

R=[]
EI=[]
for dim in dims:
    cons = [(n,(2*n+m)%dim) for n in range(dim) for m in range(0,dim-2,3)]
    path=f'/cloud/Public/_data/neuropercolation/4lay/cons=dimtimesdimby3_steps=100100/dim={dim:02d}_cons={len(cons)}/batch=0/'
    
    diff_res = [1]
    av_ei = [0]
    for eps in eps_space:
        try:
            with open(path+f"eps={round(eps,3):.3f}_phase_diff.txt", 'r', encoding='utf-8') as f:
                phase_diff = json.load(f)[100:]
            with open(path+f"eps={round(eps,3):.3f}_ei.txt", 'r', encoding='utf-8') as f:
                ei = json.load(f)[100:]
        except:
            with open(path+f"eps={round(eps,3):.3f}_activation.txt", 'r', encoding='utf-8') as f:
                activation = json.load(f)[100:]
            with open(path+f"eps={round(eps,3):.3f}_channels.txt", 'r', encoding='utf-8') as f:
                channels = 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
            ei = [np.sum(cha)*(1-H2(eps)) for cha in channels]
        
        res = np.linalg.norm(resultant(phase_diff))
        diff_res.append(res)
        
        av_ei.append(np.average(ei))
        
        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)

        with open(path+f"eps={round(eps,3):.3f}_ei.txt", 'w', encoding='utf-8') as f:
            json.dump(ei, f, indent=1)
        print(f'Done eps={eps:.3f} with dim={dim} at {datetime.now()}')
    R.append(diff_res)
    EI.append(av_ei)

    #%%
savepath=f'/cloud/Public/_data/neuropercolation/4lay/cons=dimtimesdimby3_steps=100100/plots/'

qtplot(f'Total resultant of phase-difference',
       [[0]+eps_space]*len(dims),
       R[::-1],
       [f'{dim}x{dim}' for dim in dims[::-1]],
       y_tag = 'concentration',
       export=False,
       path=savepath,
       filename=f'Resultant_ev.png',
       close=False)
      
qtplot(f'Average Effect Integrated Information',
       [[0]+eps_space]*len(dims),
       EI[::-1],
       [f'{dim}x{dim}' for dim in dims[::-1]],
       y_tag='integrated information',
       export=False,
       path=savepath,
       filename=f'EI_ev.png',
       close=False)