#!/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 numpy.linalg import norm
from datetime import datetime
from random import sample as choose

from plot import qtplot

eps_space = list(np.linspace(0.01,0.2,20))

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)

extremes = 50000
maxdt = 500

for dim in [7]:
    for eps in eps_space:
        path=f'/cloud/Public/_data/neuropercolation/4lay/cons=14_steps=1000100/dim={dim:02}/batch=0/'
        
        try:
            with open(path+f"eps={round(eps,3):.3f}_phase_diff.txt", 'r', encoding='utf-8') as f:
                phase_diff = json.load(f)
        except:
            with open(path+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(path+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))
        
        try:
            with open(path+f"eps={round(eps,3):.3f}_channelsum.txt", 'r', encoding='utf-8') as f:
                chasum = json.load(f)
        except:
            with open(path+f"eps={round(eps,3):.3f}_channels.txt", 'r', encoding='utf-8') as f:
                channels = json.load(f)[100:]
            
            chasum = [int(np.sum(cha)) for cha in channels]
            
            with open(path+f"eps={round(eps,3):.3f}_channelsum.txt", 'w', encoding='utf-8') as f:
                json.dump(chasum, f, indent=1)
        
        maxchas = max(chasum)
        
        res = [norm(resultant([pha for i, pha in enumerate(phase_diff) if chasum[i]==n])) for n in range(maxchas+1) if len([pha for i, pha in enumerate(phase_diff) if chasum[i]==n])>=20]
        #dist = [len([pha for i, pha in enumerate(phase_diff) if chasum[i]==n]) for n in range(maxchas+1)]
        
        qtplot(f'Channel-Resultant for dim={dim}',
                [list(range(len(res)))],
                [res],
                [f'Resultant against open channels'],
                x_tag = 'open channels',
                y_tag = 'resultant',
                export=True,
                path=path,
                filename=f'Channel-Resultant eps={round(eps,3):.3f} dim={dim}.png',
                close=True)
        
        print(f'Done eps={eps:.3f} with dim={dim} at {datetime.now()}')
        
    # qtplot(f'Resultant and EI evolution for dim={dim} with 4 layers',
    #         [[0]+eps_space]*2,
    #         [max(av_ei)*diff_res, av_ei],
    #         ['Resultant', 'avEI'],
    #         export=True,
    #         path=path,
    #         filename=f'Resultant and EI for dim={dim}.png',
    #         close=True)