#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Sep 27 04:39:54 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 random import random

from plot import qtplot

from neuropercolation import Simulate4Layers

eps_space = list(np.linspace(0.01,0.2,20))
def new_folder(path):
    if not os.path.exists(path):
        os.makedirs(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)

def neighbor(digit0, digit1, lenght):
    layer = int(lenght)
    dim = int(np.sqrt(layer))
    digit0, digit1 = np.array([digit0%dim, digit0//dim]), np.array([digit1%dim, digit1//dim])
    #print(digit0,digit1)
    coord_dif = list(map(abs,digit1 - digit0))
    layer_nbor = 0 in coord_dif and len(set([1,dim-1]).intersection(set(coord_dif))) != 0
    #print(coord_dif, set([1,dim-1]).intersection(set(coord_dif)))
    if layer_nbor:
        return True
    else:
        return False
    
def kcomb(zp,zm):
    if zp>2:
        val=1
    elif zm>2:
        val=0
    elif zp==zm:
        val=0.5
    elif zm==2:
        val=0.5**(3-zp)
    elif zp==2:
        val=1-0.5**(3-zm)
    elif zm==0 and zp==1:
        val=9/16
    elif zp==0 and zm==1:
        val=7/16
    else:
        raise NotImplementedError(zp,zm)
    return val

path = '/cloud/Public/_data/neuropercolation/1lay/steps=100000/'

def phi(dim,statestr,partstr,eps):
    length = dim**2
    eta = 1-eps
    
    state = np.array([int(q) for q in statestr])
    state = list(state.reshape((dim,dim)))
    state = [list([int(cell) for cell in row]) for row in state]

    part = np.array([int(p) for p in partstr])
    part = list(part.reshape((dim,dim)))
    part = [list([int(cell) for cell in row]) for row in part]

    inp = [[q+sum([state[(i+1)%dim][j],
                   state[(i-1)%dim][j],
                   state[i][(j+1)%dim],
                   state[i][(j-1)%dim]
                   ]) for j,q in enumerate(row)] for i,row in enumerate(state)]
    
    beps = [[int(inp[i][j]>2)*eta+int(inp[i][j]<3)*eps for j,q in enumerate(row)] for i,row in enumerate(state)]
    
    zplus = [[q+sum([state[(i+1)%dim][j]*(part[i][j]==part[(i+1)%dim][j]),
                     state[(i-1)%dim][j]*(part[i][j]==part[(i-1)%dim][j]),
                     state[i][(j+1)%dim]*(part[i][j]==part[i][(j+1)%dim]),
                     state[i][(j-1)%dim]*(part[i][j]==part[i][(j-1)%dim])
                     ]) for j,q in enumerate(row)] for i,row in enumerate(state)]
    zminus = [[sum([(1-state[(i+1)%dim][j])*(part[i][j]==part[(i+1)%dim][j]),
                    (1-state[(i-1)%dim][j])*(part[i][j]==part[(i-1)%dim][j]),
                    (1-state[i][(j+1)%dim])*(part[i][j]==part[i][(j+1)%dim]),
                    (1-state[i][(j-1)%dim])*(part[i][j]==part[i][(j-1)%dim])
                    ]) for j,q in enumerate(row)] for i,row in enumerate(state)]
    
    kplus = [[kcomb(zplus[i][j],zminus[i][j]) for j,q in enumerate(row)] for i,row in enumerate(state)]
    
    pi = [[eps*(1-kplus[i][j]) + eta*kplus[i][j] for j,q in enumerate(row)] for i,row in enumerate(state)]
    
    crossent = [[-beps[i][j]*m.log2(pi[i][j])-(1-beps[i][j])*m.log2(1-pi[i][j]) for j,q in enumerate(row)] for i,row in enumerate(state)]
    
    return np.sum(crossent) - length*H2(eps)

def MIP(dim,statestr,eps):
    lophi=np.inf
    mip = []
    for parti in range(1,2**(dim**2-1)):
        partstr = bin(parti)[2:].zfill(dim**2)
        print(partstr)
        curphi = phi(dim,statestr,partstr,eps)
        
        if curphi<lophi:
            lophi=curphi
            mip = [partstr]
        elif curphi==lophi:
            mip.append(partstr)
    
    return mip,lophi