import lattice_symmetries as ls
import numpy as np

# Spin operators
sigma_x = np.array([[0., 1.], [1., 0.]])
sigma_y = np.array([[0., -1.0j], [1.0j, 0.]])
sigma_z = np.diag([1., -1.])
sigma_0 = np.diag([1., 1.])
sigma_p = sigma_x + 1j * sigma_y
sigma_m = sigma_x - 1j * sigma_y

def create_2spin_matrix(model_name, param):
    """
    :model_name: xxz or tfim
    :param: delta_over_J or h_over_J for the respective models
    returns 2 spin hamiltonian
    """
    if model_name == 'xxz':
        #return (0.5 * (np.kron(sigma_p, sigma_m) + np.kron(sigma_m, sigma_p)) +
        #        param * np.kron(sigma_z, sigma_z)).astype('float64').tolist()
        return (0.5 * (np.kron(sigma_p, sigma_m) + np.kron(sigma_m, sigma_p)) +
                param * np.kron(sigma_z, sigma_z) + 1e-6 *(np.kron(sigma_z, sigma_0) + np.kron(sigma_0, sigma_z))).astype('float64')
    elif model_name == 'tfim':
        return ( np.kron(sigma_z, sigma_z) - 0.5*param*(np.kron(sigma_x, sigma_0) + np.kron(sigma_0, sigma_x))).astype('float64').tolist()
    else:
        raise NameError

def generate_symmetries(model_name, number_spins, periodic):
    symmetries = []
    sites = np.arange(number_spins)
    # Momentum in x direction with eigenvalue π
    if periodic: 
        T = (sites + 1) % number_spins
        symmetries.append(ls.Symmetry(T, sector=number_spins // 2))
    # Parity with eigenvalue π
    P = sites[::-1]
    symmetries.append(ls.Symmetry(P, sector=1))
    return symmetries

def create_edges(number_spins, periodic):
    edges = [(i, i+1) for i in range(number_spins-1)]
    if periodic: 
        edges.append((number_spins-1, 0))
    return edges

class spin_model:
    #constructor
    def __init__(self, model_name, number_spins, periodic, param, hamming_weight, use_symmetries=False, spin_inversion=None):
        self.__number_spins = number_spins
        self.__periodic = periodic
        self.__param = param
        self.__symmetries = []
        self.__spin_inversion = spin_inversion
        if use_symmetries:
            self.__symmetries = generate_symmetries(model_name, number_spins, periodic)
        self.__name = model_name
        self.__hamming_weight = hamming_weight
        self.basis = None
        self.create_basis()
        #creating Hamiltonian
        self.__matrix = create_2spin_matrix(model_name, param)
        self.edges = create_edges(number_spins, periodic)
        self.hamiltonian = None
        self.create_hamiltonian()
        self.eigenvalues = None
        self.eigenstates = None
        
    def print_characterisitcs(self):
        print("MODEL:           ", self.__name)
        print("Number of Spins: ", self.__number_spins)
        print("Periodic:        ", self.__periodic)
        print("2-spin-Matrix    ", self.__matrix)

    def create_basis(self):
        # Constructing the basis
        symmetry_group = ls.Group(self.__symmetries)
        self.basis = ls.SpinBasis(symmetry_group, number_spins=self.__number_spins, 
                            hamming_weight=self.__hamming_weight, spin_inversion=self.__spin_inversion)
        self.basis.build()
        print ("Hilbert space dimension is {}".format(self.basis.number_states))

    def create_hamiltonian(self):
        self.hamiltonian = ls.Operator(self.basis, [ls.Interaction(self.__matrix, self.edges)])

    def compute_ew_and_ev(self, print_gs_energy=True):
        if (self.__hamming_weight == 0 or self.__hamming_weight == self.__number_spins) and self.__name == 'xxz':
            self.eigenvalues = np.array([self.__number_spins * self.__param])
            if not self.__periodic:
                self.eigenvalues -= self.__param
            self.eigenstates = np.array([[1.],])

        else:
            self.eigenvalues, self.eigenstates = ls.diagonalize(self.hamiltonian, k=1)

        if print_gs_energy:
            print("Ground state energy is {:.10f}".format(self.eigenvalues[0]))
    
    
    # Getters for private variables 
    def number_spins(self):
        return self.__number_spins
    def periodic(self):
        return self.__periodic
    def hamming_weight(self):
        return self.__hamming_weight
    def matrix(self):
        return self.__matrix
    # destructor
    def __del__(self):
        del self

    """
    def update_basis_and_hamiltonian(self, rearrangement):
        # TODO: find out if I need to update edges as well!!!
        new_states = idxs_rearrangement(self.basis.states, rearrangement)
        self.basis.build(representatives=new_states)
        self.hamiltonian = ls.Operator(self.basis, [ls.Interaction(self.__matrix, self.edges)])
    """