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

@author: timof
"""

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/2lay/steps=100100/'
suspath = path + 'suscepts_euk/'
kurpath = path + 'kurts_euk/'
chi = chr(967)
vareps = chr(949)
kappa = chr(954)
rho=chr(961)

vals = [[],[]]

runsteps = 100000

#eps_space = np.linspace(0.005, 0.1, 20)
eps_space = np.linspace(0.005, 0.5, 100)

epses=20

dims = list(range(3,27,2))#+[16]

def __calc_chi(mode='density'):
    ma=[]
    s=[]
    k=[]
    
    suffix = mode

    for dim in dims:
        try:
            with open(suspath+f"radium_{suffix}_dim={dim:02}.txt", 'r', encoding='utf-8') as f:
                magnet = json.load(f)
            with open(suspath+f"susceptibility_{suffix}_dim={dim:02}.txt", 'r', encoding='utf-8') as f:
                suscept = json.load(f)
            with open(kurpath+f"kurtosis_{suffix}_dim={dim:02}.txt", 'r', encoding='utf-8') as f:
                kurt = []+json.load(f)
        except:
            magnet=[]
            suscept = []
            kurt = []
            jumped = False
            print('suscept or kurt file not found')
            if not os.path.exists(suspath):
                os.makedirs(suspath)
            if not os.path.exists(kurpath):
                os.makedirs(kurpath)
            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=='density':
                        mag = np.sum(np.linalg.norm(activation,axis=1))/len(activation)                  # density
                        mag2 = np.sum(np.linalg.norm(activation,axis=1)**2)/len(activation)              # density
                        mag4 = np.sum(np.linalg.norm(activation,axis=1)**4)/len(activation)              # density
                    elif mode=='absolute':
                        mag = np.sum(np.linalg.norm(activation*dim**2,axis=1))/len(activation)                 # density
                        mag2 = np.sum(np.linalg.norm(activation*dim**2,axis=1)**2)/len(activation)              # density
                        mag4 = np.sum(np.linalg.norm(activation*dim**2,axis=1)**4)/len(activation)              # density
                    
                        
                    mag = round(mag,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}")
                    print(f"===========================================================")
                    jumped=True
                except:
                    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(suspath+f"radium_{suffix}_dim={dim:02}.txt", 'w', encoding='utf-8') as f:
                json.dump(magnet, f, ensure_ascii=False, indent=1)
            with open(suspath+f"susceptibility_{suffix}_dim={dim:02}.txt", 'w', encoding='utf-8') as f:
                json.dump(suscept, f, ensure_ascii=False, indent=1)
            with open(kurpath+f"kurtosis_{suffix}_dim={dim:02}.txt", 'w', encoding='utf-8') as f:
                json.dump(kurt, f, ensure_ascii=False, indent=1)
                
        ma.append([1]+magnet[:epses])
        s.append([0]+suscept[:epses])
        k.append([1]+kurt[:epses])
    
    return ma,s,k

ma,s,k = __calc_chi()
#%%
qtplot(f"Magnetisation evolution for different automaton sizes",
        [[0]+list(eps_space[:epses])]*len(ma),
        [[dims[len(dims)-1-d]**2]+list(dims[len(dims)-1-d]**2*np.array(ls[1:])) for d,ls in enumerate(ma[::-1])],
        [f'dim={dim}x{dim}' for dim in dims[::-1]],
        y_tag = f'mean cycle {kappa}',
        export=True,
        path=path+"magnets/",
        filename=f'ep_radius_density_dims=3-9_steps={runsteps}_extra.png',
        close=False)

qtplot(f"Susceptibility evolution for different automaton sizes",
        [[0]+list(eps_space[:epses])]*len(ma),
        [[0]+list(np.array(sus[1:])*dims[len(dims)-1-i]**2) for i,sus in enumerate(s[::-1])],
        [f'dim={dim}x{dim}' for dim in dims[::-1]],
        y_tag = f'susceptibility {chi}',
        export=True,
        path=suspath,
        filename=f'eps+0_susceptibility_density_dims=3-15_odd_steps={runsteps}_extra.png',
        close=False)

qtplot("Kurtosis evolution for different automaton sizes",
        [[0]+list(eps_space[:epses])]*len(k),
        k[::-1],
        [f'dim={dim}x{dim}' for dim in dims[::-1]],
        y_tag = f'kurtosis U',
        y_range = (1,2),
        y_log=True,
        export=True,
        path=kurpath,
        filename=f'eps+0_kurtosis_density_dims=3-15_odd_steps={runsteps}_extra.png',
        close=False)