#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Aug 21 14:59:22 2023

@author: timofej
"""

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

eps_space = list(np.linspace(0.005,0.5,100))

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

extremes = None
maxdt = 300

path=f'/cloud/Public/_data/neuropercolation/4lay/steps=100000_diares_manydim/'


dims = list(range(7,9))

resus = []
phalocks = []
for dim in dims:
    resu = []
    phalock = []
    for eps in eps_space:
        dimpath=f'/cloud/Public/_data/neuropercolation/4lay/steps=100000_diares_manydim/dim={dim:02d}/'
        
        try:
            with open(dimpath+f"eps={round(eps,3):.3f}_phase_diff.txt", 'r', encoding='utf-8') as f:
                phase_diff = json.load(f)
        except:
            with open(dimpath+f"eps={round(eps,3):.3f}_activation.txt", 'r', encoding='utf-8') as f:
                activation = 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
            
            with open(dimpath+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)
        all_pha = np.arctan2(*resultant(phase_diff)[::-1])
        
        resu.append(all_res)
        phalock.append(all_pha)
    
    plotpath = new_folder(dimpath + 'resplot/')
    with open(plotpath+f"resultants.txt", 'w', encoding='utf-8') as f:
        json.dump(list(resu), f, indent=1)
    with open(plotpath+f"mainphase.txt", 'w', encoding='utf-8') as f:
        json.dump(list(phalock), f, indent=1)

    resus.append(resu)
    phalocks.append(phalock)

plotspath = new_folder(path + 'resplot/')

#%%
qtplot(f'Resultant evolution with 4 layers',
        [eps_space]*len(dims),
        resus,
        [f'Resultant for size {dim}x{dim}' for dim in dims],
        x_tag = 'epsilon',
        y_tag = 'resultant',
        lw=3,
        export=True,
        path=plotspath,
        filename=f'resultant.png',
        close=False)

qtplot(f'Dominant phase evolution with 4 layers',
        [eps_space]*len(dims),
        phalocks,
        [f'dominant phase for size {dim}x{dim}' for dim in dims],
        x_tag = 'epsilon',
        y_tag = 'resultant',
        lw=3,
        export=True,
        path=plotspath,
        filename=f'mainphase.png',
        close=True)