#-*- coding: utf-8 -*-
"""
Created on Tue Jun 28 16:45:06 2022

@author: timof
"""
import math as m
import numpy as np
from numba import jit, njit

import sys
import os
import json
import pyqtgraph as qt
import pyqtgraph.exporters

#from scipy.special import binom

log2 = np.log2

@njit
def digit(num, pos):
    div, den = num%(2**(pos)), 2**(pos-1)
    #print(div,den)
    return div//den

@njit
def binstr(num,size):
    return np.array([digit(num,pos) for pos in range(size,0,-1)], dtype=np.bool_)
        
@njit
def binint(state):
    return int(sum(2**np.linspace(len(state)-1,0, len(state))*state))


@njit
def binpart(state, num, part, part_comp):
    num_str = binstr(num, len(part_comp))
    num_dict = {pos:num_str[numpos] for numpos, pos in enumerate(part_comp)}
    num_state = [sta if pos in part else num_dict[pos] for pos, sta in enumerate(state)]
    num_state = np.array(num_state)
    #print(binint(num_state)==sum([num_state[pos]*2**(len(state)-pos-1) for pos in range(len(state))]))
    return sum([num_state[pos]*2**(len(state)-pos-1) for pos in range(len(state))])

@njit
def distance(b0, b1, lenght):
    return sum([abs(binstr(b0,lenght)[pos]-binstr(b1,lenght)[pos]) for pos in range(lenght)])

@njit
def neighbor(digit0, digit1, lenght):
    layer = int(lenght/2)
    dim = int(np.sqrt(layer))
    different_layers = digit0//layer != digit1//layer
    digit_diff = abs(digit1 - digit0)
    digit0, digit1 = digit0%layer, digit1%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 different_layers and digit_diff == int(lenght/2) or not different_layers and layer_nbor:
        return True
    else:
        return False
    
@njit
def evolve(state, lenght):
    
    #lenght = len(state)
    layer = int(lenght/2)
    dim = int(m.sqrt(layer))
    
    state = binstr(state,lenght)
    
    lay0 = state[:layer]
    lay1 = state[layer:]
    
    lay0 = lay0.reshape((dim,dim))
    lay1 = lay1.reshape((dim,dim))
    
    assert len(lay0)==len(lay1)
    assert lenght == 2*dim**2
    
    new_lay0 = lay0.copy()
    new_lay1 = lay1.copy()
    
    
    for pos0, row in enumerate(lay0):
        for pos1, cell in enumerate(row):
            alive = sum([int(lay0[pos[0]][pos[1]]) for pos in [(pos0-1, pos1), (pos0+1-dim, pos1),
                                                               (pos0, pos1-1), (pos0, pos1+1-dim)]])
            alive += 1-int(lay1[pos0][pos1])
            
            new_lay0[pos0][pos1] = alive > 2.5
              
    
    for pos0, row in enumerate(lay1):
        for pos1, cell in enumerate(row):
            alive = sum([int(lay1[pos[0]][pos[1]]) for pos in [(pos0-1, pos1), (pos0+1-dim, pos1),
                                                               (pos0, pos1-1), (pos0, pos1+1-dim)]])
            alive += int(lay0[pos0][pos1])
            
            new_lay1[pos0][pos1] = alive > 2.5
                
    new_state = np.hstack((np.reshape(new_lay0,layer),np.reshape(new_lay1,layer)))
    return binint(new_state)



@njit
def update(dim):
    size = 2*dim**2
    sig = 2**size
    return np.array([evolve(phi, size) for phi in range(sig)], dtype=np.uint32)

@njit
def binom(n, k):
    return 1 if k == 0 or k == n else np.prod(np.array([(n + 1 - i) / i for i in range(1, k + 1)]))

@njit
def H_full(eps, lenght):
    if eps == 0:
        return 0
    else:
        return -sum([binom(lenght, num)*((1-eps)**(lenght-num)*eps**num)*np.log2((1-eps)**(lenght-num)*eps**num) for num in range(lenght+1)])

#print(H_full(0.1,4))

@njit#(fastmath=True)
def ei(dim, state, partition, eps=0):
    lenght = 2*dim**2#int(np.log2(transition[-1]+1))
    partition = binstr(partition, lenght)
    state = binstr(state, lenght)
    
    # 2.33 µs ± 26.3 ns per loop
    
    parts = [[pos for pos, num in enumerate(partition) if not num],
             [pos for pos, num in enumerate(partition) if num]
              ]
    
    neighbors = [
        sorted(set([cell1 for cell1 in parts[(i+1)%2] if np.any(np.array([neighbor(cell0, cell1, lenght)
                                                                          for cell0 in parts[i]]))]))
        for i in range(2)
        ]
    
    fixed = [
        sorted(set(range(lenght)) - set(neighbors[i]))
        for i in range(2)
        ]
    
    #print(parts, neighbors, fixed)
    
    # 3.21 µs ± 93.7 ns per loop
    
    states = [[binpart(state, num, fixed[0], neighbors[0]) for num in range(2**len(neighbors[0]))],
              [binpart(state, num, fixed[1], neighbors[1]) for num in range(2**len(neighbors[1]))]]
    
    #print(parts, state, binpart(state,2,parts[0]))
    # 874 µs ± 9.06 µs per loop
    
    effects = [[evolve(phi,lenght) for phi in states[0]],
               [evolve(phi,lenght) for phi in states[1]]]
    
    # 899 µs ± 23.5 µs per loop
    
    ## does a faster calculation for probs exist. skipping potens? ##
    potens = [[
                binint(
                    np.array([val for pos, val in enumerate(binstr(phi,lenght)) if pos in parts[i]]
                    )) for phi in effects[i]
                ]
                for i in range(2)
                ]
    
    
    # 925 µs ± 14.3 µs per loop
    
    probs = [
        {eff:potens[i].count(eff)/len(potens[i]) for eff in set(potens[i])            
            }
        for i in range(2)
        ]
    
    probs_eps = [
        {eff_eps:sum(
            [probs[i][eff]*(1-eps)**(lenght-distance(eff_eps, eff, lenght))*eps**distance(eff_eps, eff, lenght)
             for eff in probs[i].keys()]
            ) for eff_eps in probs[i].keys()
            } for i in range(2)]
    
    #print( probs == probs_eps)
    # 925 µs ± 18.9 µs per loop
    
    entropies = [-sum([p*np.log2(p) for p in list(probs_eps[i].values())]) for i in range(2)]
    
    #print(partition, sum(entropies))
    
    return sum(entropies)

@njit
def MIP(dim, State, eps=0):
    lenght = 2*dim**2#int(np.log2(Transition[-1]+1))
    candidate = (np.inf,-1)
    
    for part in range(1, 2**(lenght-1)):
        effinf = ei(dim, State, part, eps=eps)
        
        if effinf < candidate[0]:
            candidate = (effinf, part)
        #yield
    
    print(candidate)

from alive_progress import alive_bar

def MIP_bar(dim,state,eps=0):
    lenght= 2*dim**2
    with alive_bar(2**(lenght-1)-1) as bar:
        for i in MIP(dim,state,eps=eps):
            bar()


update2 = update(2)
#print(ei(update2,11,1))

def print_MIPs(dim):
    candidate = ((0,0),-1)
    for state in range(2**(2*dim**2)):
        
        MIP_c = MIP(dim, state)
            
        if MIP_c[0] > candidate[0][0]:
            candidate = (MIP_c, state)
        print("MIP({}): {}".format(state,MIP_c))
        
    return "MIP({}): {}".format(candidate[1],candidate[0])
        
#print_MIPs(update2)

path = 'C:/Users/timof/OneDrive/Documents/Codes/MIP/'

space = np.linspace(0,0.32,33)


def plot_execute():
    if sys.flags.interactive != 1 or not hasattr(qt.QtCore, 'PYQT_VERSION'):
        qt.QtGui.QApplication.exec_()

def __generate_data(dim, state):
    MIP_val = space.copy()
    part_num = space.copy()
    transition = update(dim)
    lenght = 2*dim**2
    statecode = "".join(list(map(str,map(int,binstr(state,lenght)))))
    
    for pos, eps in enumerate(space):
        MIP_val[pos], part_num[pos] = MIP(transition, state, eps=eps)
        print(MIP_val[pos], part_num[pos])
        
    if not os.path.exists(path+"dim={}".format(dim)):
        os.makedirs(path+"dim={}".format(dim))
        
    with open(path+"dim={}/MIP_{}.txt".format(dim, statecode), 'w', encoding='utf-8') as f:
        json.dump(list(zip(MIP_val,part_num)), f, ensure_ascii=False, indent=4)
        


def __plot_data(dim, state):
    lenght = 2*dim**2
    statecode = "".join(list(map(str,map(int,binstr(state,lenght)))))
    
    
    linewidth = 4
    with open(path+"dim={}/MIP_{}.txt".format(dim, statecode), 'r', encoding='utf-8') as f:
        data, partition = list(zip(*json.load(f)))
    # with open(phi_path+"dim={}/Phi_val2.txt".format(dim), 'r', encoding='utf-8') as f:
    #     Phi_val2= np.array(json.load(f))
  
    #Phi_val = (Phi_val1 + Phi_val2)/2
    
    sp_Title = 'ei evolution for dim={}'.format(dim)
    
    ph = qt.plot()
    ph.showGrid(x=True,y=True)
    ph.setXRange(0.0, 0.32)
    #ph.setYRange(0.0, 6)
    
    #ph.setTitle(title='Phi evolution for dim={}'.format(dim))
    ph.setLabel('left', 'ei')
    ph.setLabel('bottom', 'epsilon')
    #ph.setXRange(0, 0.15)
    
    ph.addLegend(offset=(450, 30))
    
    
    #s = ph.plot(np.linspace(0.01,0.32,32), eps_max_freq0, title=sp_Title, pen='w')
    #s = ph.plot(np.linspace(0.01,0.32,32), eps_max_freq1, title=sp_Title, pen='w')
    s = ph.plot(space, data, 
                          name='Effect Information MIP for partition {}'.format(set(partition)), pen=qt.mkPen('c', width=linewidth))
    
    
    #s = ph.plot()
    exporter = qt.exporters.ImageExporter(ph.plotItem)
    
    # set export parameters if needed
    exporter.parameters()['width'] = 1200   # (note this also affects height parameter)
    
    # save to file
    exporter.export(path+"dim={}/MIP_{}.png".format(dim, statecode))
    
    
    plot_execute()
    
def analyse(dim, state):
    __generate_data(dim, state)
    __plot_data(dim, state)
    
#analyse(2,15)