# -*- coding: utf-8 -*-
"""
Created on Tue Jun 21 14:39:34 2022

@author: timof
"""

import math as m
import numpy as np
from numba import njit

log2 = np.log2


update2x2 = { "0000" : "0000",
           "0001" : "0000",
           "0010" : "0000",
           "0011" : "0011",
           "0100" : "1100",
           "0101" : "1001",
           "0110" : "0101",
           "0111" : "0011",
           "1000" : "1100",
           "1001" : "1010",
           "1010" : "0110",
           "1011" : "0011",
           "1100" : "1100",
           "1101" : "1111",
           "1110" : "1111",
           "1111" : "1111"
           }

update2x3 = { "000000" : "000000",
              "000001" : "000000",
              "000010" : "000000",
              "000011" : "000100",
              "000100" : "000000",
              "000101" : "000010",
              "000110" : "000001",
              "000111" : "000111",
              "001000" : "110000",
              "001001" : "110000",
              "001010" : "100001",
              "001011" : "100101",
              "001100" : "010001",
              "001101" : "010011",
              "001110" : "000001",
              "001111" : "000111",
              
              "010000" : "101000",
              "010001" : "100010",
              "010010" : "101000",
              "010011" : "100110",
              "010100" : "001010",
              "010101" : "000010",
              "010110" : "001011",
              "010111" : "000111",
              "011000" : "111000",
              "011001" : "110010",
              "011010" : "101001",
              "011011" : "100111",
              "011100" : "111011",
              "011101" : "110011", #circle2
              "011110" : "101011",
              "011111" : "100111",
              
              "100000" : "011000",
              "100001" : "010100",
              "100010" : "001100",
              "100011" : "000100",
              "100100" : "011000",
              "100101" : "010110",
              "100110" : "001101",
              "100111" : "000111",
              "101000" : "111000",
              "101001" : "110100", #circle1
              "101010" : "111101",
              "101011" : "110101", #circle2
              "101100" : "011001", #circle3
              "101101" : "010111",
              "101110" : "011101", #circle
              "101111" : "010111",
              
              "110000" : "111000",
              "110001" : "111110",
              "110010" : "101100",
              "110011" : "101110",
              "110100" : "011010",
              "110101" : "011110",
              "110110" : "001111",
              "110111" : "001111",
              "111000" : "111000",
              "111001" : "111110",
              "111010" : "111101",
              "111011" : "111111",
              "111100" : "111011",
              "111101" : "111111",
              "111110" : "111111",
              "111111" : "111111"
              }

def binstr(num, size):
    return bin(num)[2:].zfill(size)

def evolve2(state):
    
    lenght = len(state)
    layer = int(lenght/2)
    dim = int(m.sqrt(layer))
    
    state = np.array(list(map(int,list(state))))
    
    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 = np.zeros_like(lay0)
    new_lay1 = np.zeros_like(lay1)
    
    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, pos1-1)]])
            alive += 1-lay1[pos0][pos1]
            
            new_lay0[pos0][pos1] = int(alive > 1.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, pos1-1)]])
            alive += lay0[pos0][pos1]
            
            new_lay1[pos0][pos1] = int(alive > 1.5)
                
    new_state = np.hstack((np.reshape(new_lay0,layer),np.reshape(new_lay1,layer)))
    return "".join(list(map(str,new_state)))

def integer(num):
    if num=="0":
        return np.bool_(0)
    else:
        return np.bool_(1)

def evolve(state):
    
    lenght = len(state)
    layer = int(lenght/2)
    dim = int(m.sqrt(layer))
    
    state = np.array(list(map(int,list(state))))
    
    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 = np.zeros_like(lay0)
    new_lay1 = np.zeros_like(lay1)
    
    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-lay1[pos0][pos1]
            
            new_lay0[pos0][pos1] = int(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 += lay0[pos0][pos1]
            
            new_lay1[pos0][pos1] = int(alive > 2.5)
                
    new_state = np.hstack((np.reshape(new_lay0,layer),np.reshape(new_lay1,layer)))
    return "".join(list(map(str,new_state)))

def update_trans(dim):
    size = 2*dim**2
    sig = 2**size
    trans = ["0"*size]*sig
    
    if dim==2:
        ev = evolve2
    elif dim>=3:
        ev = evolve
    
    for phi in range(sig):
        trans[phi] = ev(bin(phi)[2:].zfill(size))
        #print(binstr(phi,size),trans[phi])
        
    return trans
    

def distance(b0, b1):
    lenght = len(b0)
    return sum([abs(int(b1[pos])-int(b0[pos])) for pos in range(lenght)])


def ei(transition, state, partition):
    
    if type(transition)==dict:
    
        lenght = len(list(transition.keys())[0])
        
        parts = [{pos for pos, num in enumerate(partition) if num=="0"},
                 {pos for pos, num in enumerate(partition) if num=="1"}]
        
        states = [{phi for phi in transition.keys() if all(phi[pos]==state[pos] for pos in parts[0])},
                  {phi for phi in transition.keys() if all(phi[pos]==state[pos] for pos in parts[1])}]
        
        effects = [{transition[phi] for phi in states[0]},
                   {transition[phi] for phi in states[1]}]
        
        potens = [{"".join([phi[pos] if pos in parts[i] else "X" for pos in range(lenght)]) for phi in effects[i]}
                  for i in range(2)]
        
        probs = [{phi:len([1 for psi in states[i] if all(phi[pos]==transition[psi][pos] for pos in parts[abs(i)])])/len(states[i])
                  for phi in potens[i]} for i in range(2)]
        
        entropies = [-sum(list(probs[i].values())*np.log2(list(probs[i].values()))) for i in range(2)]
        
    else:
    
        lenght = len(transition[0])
        
        parts = [{pos for pos, num in enumerate(partition) if num=="0"},
                 {pos for pos, num in enumerate(partition) if num=="1"}]
        
        states = [{binstr(phi,lenght) for phi in range(len(transition)) if all(binstr(phi,lenght)[pos]==state[pos] for pos in parts[0])},
                  {binstr(phi,lenght) for phi in range(len(transition)) if all(binstr(phi,lenght)[pos]==state[pos] for pos in parts[1])}]
        
        effects = [{transition[int(phi,2)] for phi in states[0]},
                   {transition[int(phi,2)] for phi in states[1]}]
        
        potens = [{"".join([phi[pos] if pos in parts[i] else "X" for pos in range(lenght)]) for phi in effects[i]}
                  for i in range(2)]
        
        probs = [{phi:len([1 for psi in states[i] if all(phi[pos]==transition[int(psi,2)][pos] for pos in parts[abs(i)])])/len(states[i])
                  for phi in potens[i]} for i in range(2)]
        
        entropies = [-sum(list(probs[i].values())*np.log2(list(probs[i].values()))) for i in range(2)]
    
    return sum(entropies)
    
def MIP(Transition, State):
    
    if type(Transition)==dict:
        lenght = len(list(Transition.keys())[0])
    else:
        lenght = len(Transition[0])
    
    candidate = (np.inf,"empty")
    
    for part in range(1, 2**(lenght-1)):
        part = bin(part)[2:].zfill(lenght)
        
        effinf = ei(Transition, State, part)
        
        if effinf < candidate[0]:
            candidate = (effinf, part)
    
    return candidate

def print_MIPs(Transition):
    if type(Transition)==dict:
        for state in Transition.keys():
            print("MIP("+state+"): {}".format(MIP(Transition, state)))
    else:
        for state in range(len(Transition)):
            state = binstr(state,len(Transition[0]))
            print("MIP("+state+"): {}".format(MIP(Transition, state)))
            
def print_parts(Transition, state):
    if type(Transition)==dict:
        
        lenght = len(list(Transition.keys())[0])
        
        for part in range(1, 2**(lenght-1)):
            part = bin(part)[2:].zfill(lenght)
            
            print("State {}, Partition {}, EI: {}".format(state,part,ei(Transition, state, part)))
    else:
        lenght = len(Transition[0])
        
        for part in range(1, 2**(lenght-1)):
            part = bin(part)[2:].zfill(lenght)
            
            print("State {}, Partition {}, EI: {}".format(state,part,ei(Transition, state, part)))
        
update2x4 = update_trans(2)
print(update2x4)
print(ei(update2x4,"11111111","00010001"))
#print_parts(update2x4,"11111111")
#print_MIPs(update2x4)