# -*- coding: utf-8 -*-
"""
Created on Tue Aug 30 14:25:12 2022

@author: timofej
"""

import sys
import os
import json

from plot import qtplot


import math as m
import numpy as np


vect = np.vectorize

@vect
def log2(x):
    try:
        return m.log2(x)
    except ValueError:
        if x==0:
            return float(0)
        else:
            raise

def new_folder(path):
    if not os.path.exists(path):
        os.makedirs(path)
    return 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)


path = '/cloud/Public/_data/neuropercolation/4lay/cons=dimtimesdimby3_steps=100100/'
suffix = ''

chi = chr(967)
vareps = chr(949)
varphi = chr(981)

vals = [[],[]]

runsteps = 1000100

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

dims = list(range(3,10))#+[16,49]

mode='density'
ma=[]
s=[]
k=[]
mk=[]
PHI=[]
lastkurt=None
for dim in dims:
    phis=[]
    con_gap = 3
    cons = [(n,(2*n+m)%dim) for n in range(dim) for m in range(0,dim-2,con_gap)]
    dimpath = new_folder(path + f'dim={dim:02}_cons={len(cons)}/batch=0/')
    for epsilon in eps_space:
        try:
            with open(dimpath+f"eps={round(epsilon,3):.3f}_ei.txt", 'r', encoding='utf-8') as f:
                ei = np.array(json.load(f)[100:])
        except:
            print('Calcing phi')
            with open(dimpath+f"eps={round(epsilon,3):.3f}_channels.txt", 'r', encoding='utf-8') as f:
                channels = np.array(json.load(f)[100:])
                
            ei = channels*(1-H2(epsilon))
            
            with open(dimpath+f"eps={round(epsilon,3):.3f}_ei.txt", 'r', encoding='utf-8') as f:
                json.dump(ei,f,indent=1)
                
        phi=np.average(ei)
        phis.append(phi)
    PHI.append(phis)
    #%%
qtplot(f"Mean effect integration over noise level",
        [eps_space]*len(dims),
        PHI[::-1],
        [f'dim={dim:02d}x{dim:02d}' for dim in dims[::-1]],
        y_tag = f'effect integration {varphi}',
        export=True,
        path=dimpath+"",
        filename=f'eps={round(epsilon,3):.3f}_evolution.png',
        close=False) 

mode = 'density'
#%%