# -*- 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
            
path = '/cloud/Public/_data/neuropercolation/1lay/steps=1000100/'
suffix = ''

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

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]

def __calc_chi(mode='density'):
    ma=[]
    s=[]
    k=[]
    mk=[]
    lastkurt=None
    for dim in dims:
        try:
            1/0
            with open(path+f"magnets/magnetisation_{mode}_dim={dim:02}.txt", 'r', encoding='utf-8') as f:
                magnet = json.load(f)
            with open(path+f"suscepts/susceptibility_{mode}_dim={dim:02}.txt", 'r', encoding='utf-8') as f:
                suscept = json.load(f)
            with open(path+f"kurts/kurtosis_{mode}_dim={dim:02}.txt", 'r', encoding='utf-8') as f:
                kurt = json.load(f)
        except:
            magnet = []
            suscept = []
            kurt = []
            jumped=False
            print('magnets or suscept or kurt file not found')
            if not os.path.exists(path+"magnets/"):
                os.makedirs(path+"magnets/")
            if not os.path.exists(path+"suscepts/"):
                os.makedirs(path+"suscepts/")
            if not os.path.exists(path+"kurts/"):
                os.makedirs(path+"kurts/")
            for epsilon in eps_space:
                try:
                    with open(path+f"dim={dim:02}/eps={round(epsilon,3):.3f}_activation.txt", 'r', encoding='utf-8') as f:
                        activation = np.array(json.load(f)[100:])
                    
                    if mode=='absolute':
                        dmag = np.sum(activation*dim**2)/len(activation)
                        mag = np.sum(np.abs(activation*dim**2))/len(activation)
                        mag2 = np.sum(activation**2*dim**4)/len(activation)
                        mag4 = np.sum(activation**4*dim**8)/len(activation)
                    elif mode=='density':
                        dmag = np.sum(activation)/len(activation)
                        mag = np.sum(np.abs(activation))/len(activation)
                        mag2 = np.sum(activation**2)/len(activation)
                        mag4 = np.sum(activation**4)/len(activation)
                    else:
                        raise NotImplementedError
                        
                    mag = round(abs(dmag),6)
                    fluct = round(mag2-mag**2,6)
                    tail = round(mag4/mag2**2,6)
                    
                    magnet.append(mag)
                    suscept.append(fluct)
                    kurt.append(tail)
                    print(f"Done dim={dim:02} eps={round(epsilon,3):.3f}: mag={mag}")
                    print(f"Done dim={dim:02} eps={round(epsilon,3):.3f}: fluct={fluct}")
                    print(f"Done dim={dim:02} eps={round(epsilon,3):.3f}: tail={tail}")
                    jumped=True
                except:
                    if not jumped:
                        magnet.append(1)
                    elif jumped:
                        magnet.append(0)
                    suscept.append(0)
                    if not jumped:
                        kurt.append(1)
                    elif jumped:
                        kurt.append(3)
                    print(f"Missing dim={dim:02} eps={round(epsilon,3):.3f}")
            with open(path+f"magnets/magnetisation_{mode}_dim={dim:02}.txt", 'w', encoding='utf-8') as f:
                json.dump(magnet, f, ensure_ascii=False, indent=1)
            with open(path+f"suscepts/susceptibility_{mode}_dim={dim:02}.txt", 'w', encoding='utf-8') as f:
                json.dump(suscept, f, ensure_ascii=False, indent=1)
            with open(path+f"kurts/kurtosis_{mode}_dim={dim:02}.txt", 'w', encoding='utf-8') as f:
                json.dump(kurt, f, ensure_ascii=False, indent=1)
            
        if lastkurt is not None:
            pos=0
            while kurt[pos]<lastkurt[pos]:
                pos+=1
            currdiff = kurt[pos]-lastkurt[pos]
            lastdiff = kurt[pos-1]-lastkurt[pos-1]
            crossres = currdiff/(currdiff-lastdiff)
            crosspoint = eps_space[pos] - crossres*(eps_space[pos]-eps_space[pos-1])
            assert eps_space[pos-1]<crosspoint<eps_space[pos]
            mk.append(crosspoint)
        lastkurt = kurt
        ma.append(magnet)
        s.append(suscept)
        k.append(kurt)
        
    
    return ma,s,k,mk

mode = 'absolute'
ma,s,k,mk = __calc_chi(mode=mode)
#%%
qtplot(f"Magnetisation {mode} evolution for different automaton sizes",
        [[0]+list(eps_space)]*len(ma),
        [[1]+ls for d,ls in enumerate(ma[::-1])],
        [f'dim={dim}x{dim}' for dim in dims[::-1]],
        x_tag = f'noise parameter {vareps}',
        y_tag = f'magnetisation {m}',
        export=True,
        path=path+"magnets/",
        filename=f'eps_magnetisation_{mode}_dims=3-9_steps={runsteps}_extra.png',
        close=False)

qtplot(f"Susceptibility evolution for different automaton sizes",
        [[0]+list(eps_space)]*len(s),
        [[0]+ls for ls in s[::-1]],
        [f'dim={dim}x{dim}' for dim in dims[::-1]],
        x_tag = f'noise parameter {vareps}',
        y_tag = f'susceptibility {chi}',
        export=True,
        path=path+"suscepts/",
        filename=f'eps_susceptibility_{mode}_dims=3-9_steps={runsteps}_extra.png',
        close=False)

qtplot("Kurtosis evolution for different automaton sizes",
        [[0]+list(eps_space)]*len(k),
        [[1]+ls for ls in k[::-1]],
        [f'dim={dim}x{dim}' for dim in dims[::-1]],
        x_tag = f'noise parameter {vareps}',
        y_tag = f'kurtosis U',
        export=True,
        path=path+"kurts/",
        filename=f'eps_kurtosis_dims=3-9_steps={runsteps}_extra.png',
        close=False)

qtplot("Kurtosis crossing for different automaton sizes",
        [[dim+0.5 for dim in dims[:-1]]],
        [mk],
        ['epsilon of kurtosis crosspoint'],
        x_tag = 'size',
        y_tag = f'epsilon {vareps}',
        export=True,
        path=path+"kurts/",
        filename=f'kurtosis_cross_dims=3-9_steps={runsteps}_extra.png',
        close=False)