import numpy as np
import sys
import math
from numpy.core.arrayprint import format_float_positional

from numpy.core.fromnumeric import trace

from joblib import delayed, Parallel
import numpy as np
import scipy.special
import pickle

import numba
from numba import uint64, float32, float64
import sys
import os

import lattice_symmetries
from lattice_symmetries import *

@numba.jit("uint64(uint64)", nogil=True, nopython=True)
def _hamming_weight(n: int):
    """The Hamming weight counts the non zero-values in a bit string. 
    In the case of spin-chains this is the amount of up-spins.
    Example: 10 -> 1010 -> hamming weight = 2
    See https://stackoverflow.com/a/9830282
    """
    n = (n & 0x5555555555555555) + ((n & 0xAAAAAAAAAAAAAAAA) >> 1)
    n = (n & 0x3333333333333333) + ((n & 0xCCCCCCCCCCCCCCCC) >> 2)
    n = (n & 0x0F0F0F0F0F0F0F0F) + ((n & 0xF0F0F0F0F0F0F0F0) >> 4)
    n = (n & 0x00FF00FF00FF00FF) + ((n & 0xFF00FF00FF00FF00) >> 8)
    n = (n & 0x0000FFFF0000FFFF) + ((n & 0xFFFF0000FFFF0000) >> 16)
    n = (n & 0x00000000FFFFFFFF) + ((n & 0xFFFFFFFF00000000) >> 32)
    return n

def sector_dm(trace_states, sub_states, amplitudes, spin_inversion):
    distinct_sub_states = list(set(sub_states))
    distinct_trace_states = list(set(trace_states))
    n = len(distinct_sub_states)
    # print('Trace states:', trace_states)
    # print('Sub_states:', sub_states)
    matrix = np.zeros((n,n))
    for state in distinct_trace_states:
        
        single_trace_result = [(amplitudes[i], distinct_sub_states.index(sub_states[i])) 
                                for i, x in enumerate(trace_states) if x == state]
        
        psi = np.zeros(n)
        for x in single_trace_result:  
            psi[x[1]] += x[0]
        matrix_add = np.outer(psi,psi)
        if spin_inversion is not None:
            matrix += 0.5 * (matrix_add + matrix_add[::-1,::-1])
        else:
            matrix += matrix_add
    return matrix

def reduced_dm(sub_dim, dim, hamming, amplitudes, states_of_amplitudes, spin_inversion=None, first_trace_spin=None):
    """
    :sub_dim: amount of spins in the subsystem (uint)
    :dim: amount of spins of total system (uint)
    :hamming: amount of upspins of total system(uint or None)
    :amplitudes: amplitudes of input state (1d-array of floats)
    :states_of_amplitudes: state which belongs to amplitude (1d-array of uints)
    :first_spin: first spin that's traced out; start counting at 0 (rightmost spin)!
    :return: List of sector density matrices
    """
    if spin_inversion is not None: 
        dim -= 1
    if first_trace_spin is None:
        sub_states = states_of_amplitudes >> (dim - sub_dim)
        trace_states = states_of_amplitudes & (2**(dim-sub_dim) - 1)
    else:
        right_sub_states = states_of_amplitudes & (2**first_trace_spin - 1)
        left_sub_states = states_of_amplitudes >> (dim - (sub_dim - first_trace_spin))
        sub_states = right_sub_states + (left_sub_states << first_trace_spin)
        trace_states = states_of_amplitudes & (2**(dim - sub_dim + first_trace_spin) - 2**first_trace_spin )
        trace_states >>= first_trace_spin
    
    hamming_of_sub_states = np.array([_hamming_weight(n) for n in sub_states])
    # print(sub_states, trace_states, hamming_of_sub_states, amplitudes)

    rho = []
    if hamming is None:
        rho.append(sector_dm(trace_states, sub_states, amplitudes, spin_inversion))
    else:
        for sub_hamming in range(max(0, hamming - (dim - sub_dim)), min(hamming, sub_dim) + 1):
            # print('Sub hamming is', sub_hamming)
            indices = np.where(hamming_of_sub_states == sub_hamming)[0].tolist()
            sector_sub_states = [sub_states[i] for i in indices]
            sector_trace_states = [trace_states[i] for i in indices]
            # print(sector_trace_states)
            sector_amplitudes = [amplitudes[i] for i in indices]
            rho.append(sector_dm(sector_trace_states, sector_sub_states, sector_amplitudes, spin_inversion))
    return rho

    
def lambda_log_lambda(x):
    """
    Von Neumann Entropy (log with basis 2)
    :x: array of eigenvalues of a sector density matrix
    :return: negative Bipartite Entanglement entropy of a sector density_matrix
    """
    y = np.where(x > 1e-7, np.log(x) / np.log(2), 1e-17)
    y *= x
    return y

def compute_entropy(rho):
    """
    :rho: list of sector density matrices
    :return: Bipartite Entanglement Entropy (von Neumann)
    """
    entropy = 0
    for i in range(len(rho)):
            spectrum, _ = np.linalg.eigh(rho[i])
            entropy -= lambda_log_lambda(spectrum).sum()
    return entropy

def compute_MI(sub_dim, dim, hamming, amplitudes, states_of_amplitudes, spin_inversion=None, first_trace_spin=None):
    # so far left and right sub-dimension is the same and equal to sub_dim
    assert (first_trace_spin + sub_dim) <= dim and first_trace_spin >= sub_dim
    first_trace_spin_left = first_trace_spin - sub_dim
    #S_AB:
    rhos_AB = reduced_dm(sub_dim*2, dim, hamming, amplitudes, states_of_amplitudes, 
                            spin_inversion=spin_inversion, first_trace_spin=first_trace_spin)
    S_AB = compute_entropy(rhos_AB)
    #S_B (right)
    rhos_B = reduced_dm(sub_dim, dim, hamming, amplitudes, states_of_amplitudes, 
                            spin_inversion=spin_inversion, first_trace_spin=first_trace_spin)
    S_B = compute_entropy(rhos_B)
    #S_A (left)
    rhos_A = reduced_dm(sub_dim, dim, hamming, amplitudes, states_of_amplitudes, 
                            spin_inversion=spin_inversion, first_trace_spin=first_trace_spin_left)
    S_A = compute_entropy(rhos_A)
    return [S_A, S_B, S_AB] #S_A + S_B -S_AB