""" 
Model declarations, from the Slater-Koster tight-binding model
to toy models such as the BHZ model or Wilson fermions.
"""

from __future__ import annotations
from typing import List, Tuple, Union

from .system import System, FrozenClass
from .crystal import Crystal
import numpy as np
import scipy.sparse as sp
import sys
import math
import cmath
from .utils import overrides

# Module-level variables
sigma_x = np.array([[0, 1],
                    [1, 0]], dtype=np.complex_)
sigma_y = np.array([[0, -1j],
                    [1j, 0]], dtype=np.complex_)
sigma_z = np.array([[1, 0],
                    [0, -1]], dtype=np.complex_)
"""
Definition of Pauli matrix x, y, z. Used by several predefined models.
"""

# ------------------ Constants and auxiliary methods ------------------
PI = 3.14159265359

# ------------------------------ Models ------------------------------


class SlaterKoster(System):
    """ 
    Implementation of Slater-Koster tight-binding model.
    
    :param configuration: Dictionary generated by fileparse.
    :param mode: Used to specify how to find the bonds. Either 'minimal' or
        'radius'. Minimal mode looks for first, second, ..., neighbours depending on
        the maximum distance specified on the configuration file. Radius mode defines
        bonds for all atoms within the specified radius. Defaults to 'minimal'.
    :param r: Value of the radius. Relevant only in 'radius' mode. Defaults to None.
    :param boundary: Boundary conditions, either 'PBC' or 'OBC'. Defaults to 'PBC'.
    """

    def __init__(self, configuration: dict = None, mode: str = 'minimal', r: float = None, boundary: str = "PBC") -> None:
        if configuration is not None:
            super().__init__(system_name=configuration["System name"],
                             bravais_lattice=configuration["Bravais lattice"],
                             motif=configuration["Motif"])
        else:
            super().__init__()

        self.configuration = configuration
        self.species = configuration['Species']
        self.boundary = boundary
        self._mode = mode
        self._r = r

        self.hamiltonian = None
        self.neighbours = None
        self.spin_orbit_hamiltonian = None        
        self.__zeeman = None
        self.__ordering = None
        self._unit_cell_list = None

        if mode not in ['minimal', 'radius']:
            print('Error: Incorrect mode')
            sys.exit(1)
        
        self.neighbours = len(configuration['SK amplitudes'].keys())

    @property
    def ordering(self):
        return self.__ordering

    @ordering.setter
    def ordering(self, ordering: str):
        if ordering != "atomic" and ordering != "spin":
            raise ValueError("ordering must be either atomic or spin")
        self.__ordering = ordering

    @property
    def nspecies(self):
        return self.configuration["Species"]

    @property
    def filling(self) -> int:
        if 'Filling' in self.configuration:
            nelectrons = 0
            for atom in self.motif:
                species = int(atom[3])
                nelectrons += self.configuration['Filling'][species]

            return int(nelectrons)


    # --------------- Methods ---------------
    def _hopping_amplitude(self, position_diff: np.ndarray, orbitals: list) -> float:
        """ 
        Routine to compute the hopping amplitude from one atom to another depending on the
        participating orbitals, following the Slater-Koster approximation. 

        :param position_diff: Array with the vector connecting the initial atom and the final atom.
        :param orbitals: List with [initial_orbital, initial_species, final_orbital, final_species, neighbour_index].
            neighbour_index denotes if they are first nn, second, etc.
        :return: Hopping amplitude.
        """

        initial_orbital = orbitals[0]
        initial_species = int(orbitals[1])
        final_orbital   = orbitals[2]
        final_species   = int(orbitals[3])
        neigh           = orbitals[4]

        initial_orbital_type = initial_orbital[0]
        final_orbital_type = final_orbital[0]

        amplitudes = self.configuration['SK amplitudes'][neigh]

        possible_orbitals = {'s': 0, 'p': 1, 'd': 2}
        if possible_orbitals[initial_orbital_type] > possible_orbitals[final_orbital_type]:
            position_diff = np.array(position_diff) * (-1)

            initial_orbital, final_orbital = final_orbital, initial_orbital
            initial_species, final_species = final_species, initial_species

            initial_orbital_type = initial_orbital[0]
            final_orbital_type = final_orbital[0]
            
        d_orbitals = {'dxy': 0, 'dyz': 0, 'dzx': 0, 'dx2-y2': 1, 'd3z2-r2':2}
        if (initial_orbital_type == "d" and final_orbital_type == "d"
                and
            d_orbitals[initial_orbital] > d_orbitals[final_orbital]):
            position_diff = np.array(position_diff) * (-1)
           
            initial_orbital, final_orbital = final_orbital, initial_orbital
            initial_species, final_species = final_species, initial_species

            initial_orbital_type = initial_orbital[0]
            final_orbital_type = final_orbital[0]
        
        species_pair = str(initial_species) + str(final_species)
        if species_pair not in amplitudes.keys():
            species_pair = str(final_species) + str(initial_species)

        effective_amplitudes = amplitudes[species_pair]
        direction_cosines = position_diff / np.linalg.norm(position_diff)
        direction_cosines = {'x': direction_cosines[0], 'y': direction_cosines[1], 'z': direction_cosines[2]}
        [l, m, n] = direction_cosines.values()
        (Vsss, Vsps, Vpps, Vppp, Vsds, Vpds, Vpdp, Vdds, Vddp, Vddd) = effective_amplitudes
        special_orbital = True if final_orbital == 'dx2-y2' or final_orbital == 'd3z2-r2' else False

        # Start checking the different possibilities
        if initial_orbital_type == "s":
            if final_orbital == "s":
                hopping = Vsss
            elif final_orbital_type == "p":
                coordinate_final = final_orbital[1]
                hopping = direction_cosines[coordinate_final] * Vsps
            elif final_orbital_type == "d" and not special_orbital:
                first_coordinate = final_orbital[1]
                second_coordinate = final_orbital[2]
                hopping = math.sqrt(3) * direction_cosines[first_coordinate] * direction_cosines[second_coordinate] * Vsds
            elif final_orbital == "dx2-y2":
                hopping = math.sqrt(3)/2.*(l*l - m*m)*Vsds
            elif final_orbital == "d3z2 - r2":
                hopping = (n*n - (l*l + m*m)/2.)*Vsds

        elif initial_orbital_type == "p":
            coordinate_initial = initial_orbital[1]
            if final_orbital_type == "p":
                coordinate_final = final_orbital[1]
                hopping = direction_cosines[coordinate_initial]*direction_cosines[coordinate_final]*(Vpps - Vppp)
                if coordinate_final == coordinate_initial:
                    hopping += Vppp
            elif final_orbital_type == "d" and not special_orbital:
                first_coordinate = final_orbital[1]
                second_coordinate = final_orbital[2]
                hopping = (direction_cosines[coordinate_initial]*direction_cosines[first_coordinate]
                           * direction_cosines[second_coordinate]*(math.sqrt(3)*Vpds - 2*Vpdp))
                if coordinate_initial == first_coordinate:
                    hopping += direction_cosines[second_coordinate]*Vpdp
                elif coordinate_initial == second_coordinate:
                    hopping += direction_cosines[first_coordinate]*Vpdp
            elif final_orbital == "dx2-y2":
                hopping = direction_cosines[coordinate_initial]*(l*l - m*m)*(math.sqrt(3)/2.*Vpds - Vpdp)
                if coordinate_initial == 'x':
                    hopping += l*Vpdp
                elif coordinate_initial == 'y':
                    hopping += -m*Vpdp
            elif final_orbital == "d3z2-r2":
                hopping = direction_cosines[coordinate_initial]*(n*n - (l*l + m*m)/2.)*Vpds
                if coordinate_initial == 'z':
                    hopping += math.sqrt(3)*n*(l*l + m*m)*Vpdp
                else:
                    hopping -= math.sqrt(3)*direction_cosines[coordinate_initial]*n*n*Vpdp

        elif initial_orbital_type == "d" and initial_orbital not in ["dx2-y2", "d3z2-r2"]:
            coordinate_initial_first = initial_orbital[1]
            coordinate_initial_second = initial_orbital[2]
            for coordinate in ["x", "y", "z"]:  # Determine non present coordinate
                if coordinate not in [coordinate_initial_first, coordinate_initial_second]:
                    coordinate_initial_third = coordinate
            if not special_orbital:
                coordinate_final_first = final_orbital[1]
                coordinate_final_second = final_orbital[2]
                [l_aux, m_aux, n_aux] = ([direction_cosines[coordinate_initial_first],
                                          direction_cosines[coordinate_initial_second],
                                          direction_cosines[coordinate_initial_third]])
                hopping = (3.*l_aux*m_aux*
                           direction_cosines[coordinate_final_first]*direction_cosines[coordinate_final_second]*Vdds)
                if initial_orbital == final_orbital:
                    hopping += (l_aux**2 + m_aux**2 - 4*l_aux**2*m_aux**2)*Vddp + (n_aux**2 + l_aux**2*m_aux**2)*Vddd
                else:
                    non_repeated_coordinates = list({coordinate_initial_first, coordinate_initial_second} ^
                                                    {coordinate_final_first, coordinate_final_second})
                    repeated_coordinate = list({coordinate_initial_first, coordinate_initial_second} &
                                               {coordinate_final_first, coordinate_final_second})[0]
                    m_aux = direction_cosines[repeated_coordinate]
                    [l_aux, n_aux] = [direction_cosines[non_repeated_coordinates[0]],
                                      direction_cosines[non_repeated_coordinates[1]]]
                    hopping += l_aux*n_aux*(1 - 4*m_aux**2)*Vddp + l_aux*n_aux*(m_aux**2 - 1)*Vddd
            elif final_orbital == "dx2-y2":
                hopping = (direction_cosines[coordinate_initial_first]*direction_cosines[coordinate_initial_second] *
                           (l*l - m*m)*(3./2*Vdds - 2*Vddp + Vddd/2.))
                if initial_orbital == "dyz":
                    hopping += (direction_cosines[coordinate_initial_first]*direction_cosines[coordinate_initial_second] *
                                (-Vddp + Vddd))
                elif initial_orbital == "dzx":
                    hopping += (direction_cosines[coordinate_initial_first]*direction_cosines[coordinate_initial_second] *
                                (Vddp - Vddd))
            elif final_orbital == "d3z2-r2":
                hopping = math.sqrt(3)*(direction_cosines[coordinate_initial_first] *
                                        direction_cosines[coordinate_initial_second]*(n*n - (l*l + m*m)/2))*Vdds
                if initial_orbital == 'dxy':
                    hopping += math.sqrt(3)*(-2*l*m*n*n*Vddp + (l*m*(1 + n*n)/2)*Vddd)
                else:
                    hopping += math.sqrt(3)*(direction_cosines[coordinate_initial_first]*
                    direction_cosines[coordinate_initial_second] * ((l*l + m*m - n*n)*Vddp - (l*l + m*m)/2.*Vddd))

        elif initial_orbital == "dx2-y2":
            hopping = 0.0
            if final_orbital == "dx2-y2":
                hopping += 3./4*(l*l - m*m)**2*Vdds + (l*l + m*m - (l*l - m*m)**2)*Vddp + (n*n + (l*l - m*m)**2/4)*Vddd
            elif final_orbital == "d3z2-r2":
                hopping += math.sqrt(3)*((l*l - m*m)*(n*n - (l*l + m*m)/2)*Vdds/2 + n*n*(m*m - l*l)*Vddp + ((1 + n*n) *
                              (l*l - m*m)/4)*Vddd)
        elif initial_orbital == "d3z2-r2":
            hopping = (n*n - (l*l + m*m)/2)**2*Vdds + 3*n*n*(l*l + m*m)*Vddp + 3./4*(l*l + m*m)**2*Vddd

        else:
            raise ValueError("Error: Shouldn't get called (some bug in orbital assertion in code)")

        return hopping


    def __initialize_spin_orbit_coupling(self) -> None:
        """ 
        Method to initialize the whole spin-orbit coupling matrix corresponding to the
        orbitals that participate in the tight-binding model.
        """

        dimension_block = 1 + 3 + 5
        # We hardcode the whole spin-orbit hamilonian up to d orbitals
        spin_orbit_hamiltonian = np.zeros([dimension_block*2, dimension_block*2], dtype=np.complex_)
        p_orbital_beginning = 1
        d_orbital_beginning = 4
        # p orbitals
        spin_orbit_hamiltonian[p_orbital_beginning, p_orbital_beginning + 1] = -1j
        spin_orbit_hamiltonian[p_orbital_beginning, dimension_block + p_orbital_beginning + 2] = 1
        spin_orbit_hamiltonian[p_orbital_beginning + 1, dimension_block + p_orbital_beginning + 2] = -1j
        spin_orbit_hamiltonian[p_orbital_beginning + 2, dimension_block + p_orbital_beginning] = -1
        spin_orbit_hamiltonian[p_orbital_beginning + 2, dimension_block + p_orbital_beginning + 1] = 1j
        spin_orbit_hamiltonian[dimension_block + p_orbital_beginning, dimension_block + p_orbital_beginning + 1] = 1j

        # d orbitals
        spin_orbit_hamiltonian[d_orbital_beginning, d_orbital_beginning + 3] = 2j
        spin_orbit_hamiltonian[d_orbital_beginning, dimension_block + d_orbital_beginning + 1] = 1
        spin_orbit_hamiltonian[d_orbital_beginning, dimension_block + d_orbital_beginning + 2] = -1j
        spin_orbit_hamiltonian[d_orbital_beginning + 1, d_orbital_beginning + 2] = 1j
        spin_orbit_hamiltonian[d_orbital_beginning + 1, dimension_block + d_orbital_beginning] = -1
        spin_orbit_hamiltonian[d_orbital_beginning + 1, dimension_block + d_orbital_beginning + 3] = -1j
        spin_orbit_hamiltonian[d_orbital_beginning + 1, dimension_block + d_orbital_beginning + 4] = -math.sqrt(3)*1j
        spin_orbit_hamiltonian[d_orbital_beginning + 2, dimension_block + d_orbital_beginning] = 1j
        spin_orbit_hamiltonian[d_orbital_beginning + 2, dimension_block + d_orbital_beginning + 3] = -1
        spin_orbit_hamiltonian[d_orbital_beginning + 2, dimension_block + d_orbital_beginning + 4] = math.sqrt(3)
        spin_orbit_hamiltonian[d_orbital_beginning + 3, dimension_block + d_orbital_beginning + 1] = 1j
        spin_orbit_hamiltonian[d_orbital_beginning + 3, dimension_block + d_orbital_beginning + 2] = 1
        spin_orbit_hamiltonian[d_orbital_beginning + 4, dimension_block + d_orbital_beginning + 1] = 1j*math.sqrt(3)
        spin_orbit_hamiltonian[d_orbital_beginning + 4, dimension_block + d_orbital_beginning + 2] = -math.sqrt(3)
        spin_orbit_hamiltonian[dimension_block + d_orbital_beginning, dimension_block + d_orbital_beginning + 3] = -2j
        spin_orbit_hamiltonian[dimension_block + d_orbital_beginning + 1,
                            dimension_block + d_orbital_beginning + 2] = -1j

        spin_orbit_hamiltonian += np.conj(spin_orbit_hamiltonian.T)

        self.spin_orbit_hamiltonian = spin_orbit_hamiltonian/3

    def __spin_orbit_h(self) -> None:
        """ 
        Method to obtain the actual spin-orbit hamiltonian that corresponds to the orbitals
        that participate in the model.
        """

        possible_orbitals = ['s', 'px', 'py', 'pz', 'dxy', 'dyz', 'dzx', 'dx2-y2', 'd3z2-r2']
        npossible_orbitals = len(possible_orbitals)
        spin_orbit_hamiltonian = sp.lil_matrix((self.basisdim, self.basisdim), dtype=np.complex_)
        matrix_index = 0
        for n, atom in enumerate(self.motif):
            species = int(atom[3])
            orbitals = self.configuration["Orbitals"][species]
            norbitals = self.norbitals[species]
            orbitals_indices = [index + i for index, orbital in enumerate(possible_orbitals) for i in [0, npossible_orbitals]
                                if orbital in orbitals]
            
            soc = (self.spin_orbit_hamiltonian[np.ix_(orbitals_indices, orbitals_indices)] *
                   self.configuration['Spin-orbit coupling'][species])
            spin_orbit_hamiltonian[matrix_index : matrix_index + norbitals, matrix_index : matrix_index + norbitals] = soc
            matrix_index += norbitals

        self.spin_orbit_hamiltonian = spin_orbit_hamiltonian

    def initialize_hamiltonian(self, find_bonds: bool = True, verbose: bool = True) -> None:
        """ 
        Routine to initialize the hamiltonian matrices which describe the system. 

        :find_bonds: Toggle on/or bond search before initializing the Hamiltonian. Useful
            when the bond search algorithm fails or we want to specify some bonds. Defaults to True.
        :verbose: Toggle print info. Defaults to True.
        """

        self.norbitals = [len(orbitals) for orbitals in self.configuration["Orbitals"]]
        self.basisdim = np.sum([self.norbitals[int(atom[3])] for atom in self.motif])

        if find_bonds:
            print('Computing first neighbours...\n')
            self.find_neighbours(mode=self._mode, nn=self.neighbours, r=self._r)
        self._determine_connected_unit_cells()

        hamiltonian = []
        for _ in self._unit_cell_list:
            hamiltonian.append(sp.lil_matrix((self.basisdim, self.basisdim), dtype=np.complex_))

        matrix_index = 0
        for n, atom in enumerate(self.motif):
            species = int(atom[3])  # To match list beginning on zero
            hamiltonian_atom_block = sp.diags(np.array(self.configuration['Onsite energy'][species]), format='lil')
            hamiltonian[0][matrix_index : matrix_index + self.norbitals[species],
                           matrix_index : matrix_index + self.norbitals[species]] = hamiltonian_atom_block
            matrix_index += self.norbitals[species]

        atom_index = {0: 0}
        for n, atom in enumerate(self.motif[1:]):
            atom_index[n + 1] = atom_index[n] + self.norbitals[int(self.motif[n][3])]

        for bond in self.bonds:
            initial_atom_index, final_atom_index, cell, neigh = bond
            initial_atom = self.motif[initial_atom_index][:3]
            initial_atom_species = int(self.motif[initial_atom_index][3])
            final_atom = self.motif[final_atom_index][:3]
            final_atom_species = int(self.motif[final_atom_index][3])
            iorbitals = self.configuration['Orbitals'][initial_atom_species]
            forbitals = self.configuration['Orbitals'][final_atom_species]
            for i, initial_orbital in enumerate(iorbitals):
                for j, final_orbital in enumerate(forbitals):
                    position_difference = np.array(final_atom) + np.array(cell) - np.array(initial_atom)
                    orbital_config = [initial_orbital, initial_atom_species, final_orbital, final_atom_species, neigh]
                    h_cell = self._unit_cell_list.index(list(cell))
                    hopping_amplitude = self._hopping_amplitude(position_difference, orbital_config)
                    hamiltonian[h_cell][atom_index[initial_atom_index] + i,
                                        atom_index[final_atom_index] + j] = hopping_amplitude
        if verbose:
            print("Hoppings computed")

        # Substract half-diagonal from all hamiltonian matrices to compensate transposing later
        #for h in hamiltonian:
        #    h -= sp.diags(h.diagonal()/2)

        # Check spinless or spinful model and initialize spin-orbit coupling
        if self.configuration['Spin']:
            self.norbitals = [2*norbitals for norbitals in self.norbitals]
            self.basisdim = self.basisdim * 2

            for index, cell in enumerate(self._unit_cell_list):
                hamiltonian[index] = sp.kron(hamiltonian[index], sp.eye(2, 2))

            # Add Zeeman term
            self.__zeeman_term(0)
            hamiltonian[0] += self.__zeeman

            if self.configuration['Spin-orbit coupling'] != 0:
                self.__initialize_spin_orbit_coupling()
                self.__spin_orbit_h()
                hamiltonian[0] += self.spin_orbit_hamiltonian
        
        if self.matrix_type == "dense":
            self.hamiltonian = [h.toarray() for h in hamiltonian]
        else:
            self.hamiltonian = [h.tocsr() for h in hamiltonian]

    def __zeeman_term(self, intensity: float) -> None:
        """ 
        Routine to incorporate a Zeeman term to the Hamiltonian. 
        
        :param intensity: Amplitude of the Zeeman term.
        """

        zeeman_h = sp.lil_matrix((self.basisdim, self.basisdim))
        matrix_index = 0
        for atom in self.motif:
            species = int(atom[3])
            zeeman_h[matrix_index: matrix_index + self.norbitals[species], 
                     matrix_index: matrix_index + self.norbitals[species]] = sp.kron(
                     np.array([[1, 0], [0, -1]]), 
                     sp.eye(self.norbitals[species]//2))
            matrix_index += self.norbitals[species]

        zeeman_h *= intensity
        self.__zeeman = zeeman_h

    def hamiltonian_k(self, k: Union[List, np.ndarray]) -> np.ndarray:
        """ 
        Routine to compute the Bloch Hamiltonian at a given k point. It can be given
        as a sparse matrix or a dense matrix depending on the matrix_type attribute.
        
        :param k: kpoint.
        :return: Bloch Hamiltonian evaluated at k.
        """

        if self.matrix_type == "dense":
            hamiltonian_k = np.zeros((self.basisdim, self.basisdim), dtype=np.complex_)
        else:
            hamiltonian_k = sp.csr_matrix((self.basisdim, self.basisdim), dtype=np.complex_)

        for cell_index, cell in enumerate(self._unit_cell_list):
            hamiltonian_k += self.hamiltonian[cell_index] * cmath.exp(1j*np.dot(k, cell))

        return hamiltonian_k
    
    def add_flat_band(self, energy: float, atom_index: int, spin: bool = False) -> SlaterKoster:
        """
        Adds a new orbital to one atom with given onsite energy and no hoppings to other atoms,
        resulting in a flat band in the Bloch Hamiltonian (localized state).
        Only works with crystalline systems.

        :param energy: Onsite energy of new orbital.
        :param atom_index: Index of atom to which we add the new orbital.
        :param spin: If set to True, two orbitals are added (one per spin).
        """

        if atom_index < 0:
            raise ValueError("Atom index must be a positive number.")
        elif atom_index > (self.natoms - 1):
            raise ValueError("Index of atom higher than number of chemical species (counting starts at 0).")
                
        index = 0
        for i in range(atom_index + 1):
            index += self.norbitals[i]

        nbasisdim = self.basisdim + 1 + spin
        for n, fock_matrix in enumerate(self.hamiltonian):
            new_fock_matrix = np.zeros([nbasisdim, nbasisdim], dtype=np.complex_)

            new_fock_matrix[0:index, 0:index] = fock_matrix[0:index, 0:index]
            new_fock_matrix[0:index, (index + 1 + spin):] = fock_matrix[0:index, index:]
            new_fock_matrix[(index + 1 + spin):, 0:index] = fock_matrix[index:, 0:index]
            new_fock_matrix[(index + 1 + spin):, (index + 1 + spin):] = fock_matrix[index:, index:]

            self.hamiltonian[n] = new_fock_matrix
        
        self.hamiltonian[0][index, index] = energy
        if spin:
            self.hamiltonian[0][index + 1, index + 1] = energy

        # Modify model attributes to take into account new orbital
        self.norbitals[atom_index] += 1 + spin
        self.basisdim += 1 + spin
        self.configuration["Filling"][atom_index] += 1 + spin

        return self
    
    def insert_SK_model(self, config: dict) -> SlaterKoster:
        """
        Takes a configuration dictionary of a system with only one atom and appends the corresponding
        model as a block to the current Hamiltonian. Note that it must describe the
        same lattice to make sense. Works only with crystalline systems.

        :param config: Configuration dictionary.
        """

        core_model = SlaterKoster(config)
        if len(core_model.motif) != 1:
            ValueError("New subsystem must have only one atom.")

        core_model.initialize_hamiltonian()
        nbasisdim = self.basisdim

        atom = core_model.motif[0][:3]
        self.motif = np.array(self.motif)
        atom_index = np.where(np.linalg.norm(self.motif[:, :3] - atom, axis=1) < 1E-4)[0]
        if len(atom_index) != 1:
            raise IndexError("Invalid atom positions in block SK model.")
        
        atom_index = atom_index[0]
        species = int(self.motif[atom_index, 3])
        
        index = 0
        for i in range(atom_index + 1):
            index += self.norbitals[i]
        
        norb = core_model.norbitals[0]
        nbasisdim += norb
        for n, fock_matrix in enumerate(self.hamiltonian):
            new_fock_matrix = np.zeros([nbasisdim, nbasisdim], dtype=np.complex_)

            new_fock_matrix[0:index, 0:index] = fock_matrix[0:index, 0:index]
            new_fock_matrix[0:index, (index + norb):] = fock_matrix[0:index, index:]
            new_fock_matrix[(index + norb):, 0:index] = fock_matrix[index:, 0:index]
            new_fock_matrix[(index + norb):, (index + norb):] = fock_matrix[index:, index:]

            new_fock_matrix[index:index + norb, index:index + norb] = core_model.hamiltonian[n]

            self.hamiltonian[n] = new_fock_matrix
        
        self.norbitals[atom_index] += core_model.norbitals[0]
        self.configuration["Filling"][species] += core_model.configuration["Filling"][0]
        self.basisdim = nbasisdim
        
        return self


    # ---------------------------- IO routines ----------------------------
    def export_model(self, filename: str, fmt: str = "%20.16f") -> None:
        """ 
        Routine to write the Hamiltonian matrices calculated to a file. 
        
        :param filename: Name of the file where we write the model.
        """

        h_fmt = [fmt+"%+"+fmt[1:]+"j" for _ in range(self.hamiltonian[0].shape[0])]
        with open(filename, "w") as file:
            # Write ndim, natoms, norbitals, ncells and bravais lattice basis vectors
            file.write("# dimension\n" + str(self.ndim) + "\n")
            file.write("# norbitals\n")
            np.savetxt(file, [self.norbitals], fmt="%i")
            file.write("# bravaislattice\n")
            np.savetxt(file, self.bravais_lattice, fmt=fmt)
            file.write("# motif\n")
            for atom in self.motif:
                np.savetxt(file, [atom], fmt=fmt)
            file.write("# bravaisvectors\n")
            for vector in self._unit_cell_list:
                np.savetxt(file, [vector], fmt=fmt)
            file.write("# hamiltonian\n")
            for h in self.hamiltonian:
                np.savetxt(file, h, fmt=h_fmt)
                file.write("&\n")
            if self.configuration["Filling"]:
                file.write("# filling\n" + str(self.filling) + "\n")
            file.write("#")


    @classmethod
    def import_model(cls, filename: str) -> 'SlaterKoster':
        """
        Routine to read the Hamiltonian matrix from a file.
        TODO: Has to be adapted to new file format. 

        :param filename: File with the configuration to load. Note that this
            file follows the format of export_model, not the one usually with fileparse.
        :return: Instance of SlaterKoster.
        """

        model = cls()
        file = open(filename, "r")
        file_contents = {}
        content = []
        for line in file.readlines():
            line = line.split()
            if line[0] == "#":
                arg = line[1]
                if not content:
                    continue
                file_contents[arg] = content
                content = []
            else:
                content.append(line)
        
        for key in file_contents.keys():
            if key == "dimension":
                model.dimension = int(file_contents[key][0])
            elif key == "norbitals":
                model.norbitals = [int(value) for value in file_contents[key][0]]
            elif key == "bravaislattice":
                bravais_lattice = []
                for vector in file_contents[key]:
                    bravais_lattice.append([float(value) for value in vector])


            # dimension
            line = file.readline()
            dimension = int(file.readline())

            # norbitals
            line = [int(value) for value in file.readline().split()]

            # bravaislattice
            

            ndim, natoms, norbitals, ncells = [int(num) for num in line]
            basisdim = norbitals * natoms
            bravais_lattice = []
            for i in range(ndim):
                line = file.readline().split()
                if len(line) != 3:
                    print("Unexpected line found, exiting...")
                    sys.exit(1)
                bravais_lattice.append([float(num) for num in line])

            motif = []
            for i in range(natoms):
                line = file.readline().split()
                if len(line) != 3:
                    print("Unexpected line found, exiting...")
                    sys.exit(1)
                motif.append([float(num) for num in line])

            unit_cell_list = []
            hamiltonian = []
            hamiltonian_matrix = np.zeros([basisdim, basisdim], dtype=np.complex_)
            for line in file.readlines():
                line = line.split()
                if len(line) == 3:
                    unit_cell_list.append([float(num) for num in line])
                elif line[0] == "#":
                    it = 0
                    hamiltonian.append(hamiltonian_matrix)
                    hamiltonian_matrix = np.zeros([basisdim, basisdim], dtype=np.complex_)
                else:
                    hamiltonian_matrix[it, :] = [complex(line[2*i], line[2*i + 1]) for i in range(len(line))]
                    it += 1

        if ncells != len(unit_cell_list):
            print("Mismatch between number of Bloch matrices provided and declared, exiting...")
            sys.exit(1)

        model = cls()
        model.system_name = filename
        model.ndim = ndim
        model.norbitals = norbitals
        model.natoms = natoms
        model.basisdim = basisdim
        model._unit_cell_list = unit_cell_list
        model.bravais_lattice = bravais_lattice
        model.hamiltonian = hamiltonian

        return model


class AmorphousSlaterKoster(SlaterKoster):
    """ 
    Extension of the Slater-Koster tight binding model to describe amorphous solids,
    taking the crystalline solid as the reference point. This model by default works in radius mode. 

    :param configuration: Configuration dictionary as generated by parsefile.
    :param r: Radius up to which bonds are considered. Defaults to None.
    :param boundary: 'PBC' or 'OBC'. Defaults to 'PBC'.
    """

    def __init__(self, configuration=None, r=None, boundary="PBC"):
        super().__init__(configuration, mode='radius', r=r, boundary=boundary)

        # Specific attributes of AmorphousSlaterKoster
        self._reference_lengths = {}
        self._decay_amplitude   = None
        self._decay_mode        = 'exponential'
        self._mode              = 'radius'

    # Properties
    @property
    def reference_lengths(self):
        return self._reference_lengths

    @reference_lengths.setter
    def reference_lengths(self, *values):
        print('reference_lengths must be mutated using the set_reference_length method')

    @property
    def decay_amplitude(self):
        return self._decay_amplitude
    
    @decay_amplitude.setter
    def decay_amplitude(self, amplitude):
        if amplitude < 0:
            raise ValueError('decay_amplitude: amplitude must be positive number')
        self._decay_amplitude = amplitude
    
    @property
    def decay_mode(self):
        return self._decay_mode
    
    @decay_mode.setter
    def decay_mode(self, mode):
        if mode != "exponential" and mode != "polynomial":
            raise ValueError('decay_mode: mode must be either exponential or polynomial')
        self._decay_mode = mode

    def set_reference_length(self, first_species: int, second_species: int, length: float) -> None:
        """ 
        Method to set manually the reference bond length between two species. 
        
        :param first_species: Index of first species.
        :param second_species: Index of second species.
        :param length: Reference length value.
        """

        if first_species > self.species or second_species > self.species:
            raise ValueError('set_reference_length: provided species must be numbers between 0 and nspecies - 1')
        if first_species > second_species:
            species_pair = str(second_species) + str(first_species)
        else:
            species_pair = str(first_species) + str(second_species)
        self._reference_lengths[species_pair] = length
        print(f"Reference lengths for species: {self._reference_lengths}")

    def __compute_reference_lengths(self):
        """ 
        Private method to compute the length between different chemical species in the crystalline solid. 
        To be used only when the model is initialized with a configuration file corresponding to a crystalline solid. 
        """

        # First compute neighbours in crystalline solid
        print("Computing reference lengths...")
        self.find_neighbours(mode=self._mode, r = self._r)
        length_dict = {}
        motif = np.array(self.motif)
        for i, j, cell, _ in self.bonds:
            initial_species = int(self.motif[i][3])
            final_species   = int(self.motif[j][3])
            distance = np.linalg.norm(motif[i, :3] - motif[j, :3] - cell)
            species_pair = str(final_species) + str(initial_species) if (i > j) else str(initial_species) + str(final_species)
            if species_pair in length_dict:
                if distance < length_dict[species_pair]:
                    length_dict[species_pair] = distance
            else:
                length_dict[species_pair] = distance

        print(f"Reference lengths for species: {length_dict}")
        self._reference_lengths = length_dict

    @overrides(SlaterKoster)
    def _hopping_amplitude(self, position_diff: np.ndarray, orbitals: list) -> float:
        """ 
        Method to incorporate a decay to the hoppings as computed within a Slater-Koster model to describe
        variable length bonds 
        
        :param position_diff: Vector connecting the initial and the final atoms.
        :param orbitals: List [initial_orbital, initial_species, final_orbital, final_species, neigh_index]
        :return: Hopping amplitude.
        """

        initial_species = int(orbitals[1])
        final_species   = int(orbitals[3])
        species_pair = str(initial_species) + str(final_species)
        if species_pair not in self.reference_lengths.keys():
            species_pair = str(final_species) + str(initial_species)
        
        reference_bond_length = self.reference_lengths[species_pair]
        r = np.linalg.norm(position_diff)
        displacement = round(r - reference_bond_length, 4)
        
        hopping = super()._hopping_amplitude(position_diff, orbitals)
        
        if self.decay_mode == "exponential":
            hopping *= np.exp(-self.decay_amplitude*displacement)
        else:
            # hopping *= 1./(1 + self.decay_amplitude*(r - reference_bond_length))
            hopping *= (r/reference_bond_length)**(-self.decay_amplitude)

        return hopping

    @overrides(SlaterKoster)
    def initialize_hamiltonian(self, find_bonds: bool = True, override_bond_lengths: bool = False) -> None:
        """
        Method to initialize the 

        :param find_bonds: Toggle on/off bond search. Defaults to True.
        :param override_bond_length: Overwrite previous reference bond lengths.
        """

        if override_bond_lengths:
            self.__compute_reference_lengths()
        super().initialize_hamiltonian(find_bonds=find_bonds)


class BHZ(System, FrozenClass):
    """
    Implementation of generalized BHZ model. The model is based on a
    2D square lattice of side a=1, and is a four band model. The model takes three parameters:
    
    :param g: Mixing term
    :param u: Sublattice potential
    :param c: Coupling operator "amplitude"
    """

    def __init__(self, g, u, c):
        super().__init__(system_name="BHZ model",
                         crystal=Crystal([[1, 0, 0], [0, 1, 0]],
                                         motif=[[0, 0, 0, 0]]))
        self.norbitals = [4]
        self.filling = 2
        self.basisdim = self.norbitals[0]
        self.g = g
        self.u = u
        self.c = c

        self._freeze()

    def hamiltonian_k(self, k: Union[List, np.ndarray]) -> np.ndarray:
        """
        Method to compute the Bloch Hamiltonian.
        
        :param k: Array used to compute the Bloch Hamiltonian.
        :return: H(k).
        """
        
        global sigma_x, sigma_y, sigma_z
        coupling = self.c * sigma_y

        id2 = np.eye(2)
        kx, ky = k[0], k[1]

        hamiltonian = (np.kron(id2, (self.u + np.cos(kx) + np.cos(ky)) * sigma_z + np.sin(ky) * sigma_y) +
                       np.kron(sigma_z, np.sin(kx) * sigma_x) + np.kron(sigma_x, coupling) +
                       self.g * np.kron(sigma_z, sigma_y) * (np.cos(kx) + np.cos(7 * ky) - 2))

        return hamiltonian


class WilsonFermions2D(System, FrozenClass):
    """
    Implementation of Wilson-fermions model. This model takes a 2D square lattice of side a, it is
    a four band model.

    :param side: Length of lattice (square) side. Defaults to 1.
    :param t: Hopping amplitude. Defaults tot 1.
    :param m: Mass term. Defaults to 1.
    """

    def __init__(self, side=1, t=1, m=1):
        super().__init__(system_name="Wilson-fermions 2D",
                         crystal=Crystal([[side, 0, 0], [0, side, 0]],
                                         motif=[[0, 0, 0, 0]]))
        self.system_name = "Wilson-fermions model"
        self.norbitals = 4
        self.basisdim = self.norbitals * 1
        self.filling = 0.5

        self.a = side
        self.t = t
        self.m = m

        self._freeze()

    def hamiltonian_k(self, k: Union[List, np.ndarray]) -> np.ndarray:
        """
        Method to compute the Bloch Hamiltonian.
        
        :param k: Array used to compute the Bloch Hamiltonian.
        :return: H(k).
        """

        global sigma_x, sigma_y, sigma_z

        alpha_x = np.kron(sigma_x, sigma_x)
        alpha_y = np.kron(sigma_x, sigma_y)
        beta = np.kron(sigma_z, np.eye(2, dtype=np.complex_))

        hamiltonian = self.t * (np.sin(k[0] * self.a) * alpha_x + np.sin(k[1] * self.a) * alpha_y) + (
                np.cos(k[0] * self.a) + np.cos(k[1] * self.a) + self.m - 3) * beta

        return hamiltonian


class WilsonFermions3D(System, FrozenClass):
    """
    Implementation of Wilson-fermions model. This model takes a 3D square lattice of side a, it is
    a four band model.

    :param side: Length of lattice (square) side. Defaults to 1.
    :param t: Hopping amplitude. Defaults tot 1.
    :param m: Mass term. Defaults to 1.
    """

    def __init__(self, side=1, t=1, m=1):
        super().__init__(system_name="Wilson-fermions 3D",
                         crystal=Crystal([[side, 0, 0], [0, side, 0], [0, 0, side]],
                                         motif=[[0, 0, 0, 0]]))
        self.num_orbitals = 4
        self.basisdim = self.num_orbitals * self.natoms

        self.a = side
        self.t = t
        self.m = m

        self._freeze()

    def hamiltonian_k(self, k: Union[List, np.ndarray]) -> np.ndarray:
        """
        Method to compute the Bloch Hamiltonian.
        
        :param k: Array used to compute the Bloch Hamiltonian.
        :return: H(k).
        """

        global sigma_x, sigma_y, sigma_z

        alpha_x = np.kron(sigma_x, sigma_x)
        alpha_y = np.kron(sigma_x, sigma_y)
        alpha_z = np.kron(sigma_x, sigma_z)
        beta = np.kron(sigma_z, np.eye(2, dtype=np.complex_))

        hamiltonian = (self.t * (np.sin(k[0] * self.a) * alpha_x + np.sin(k[1] * self.a) * alpha_y +
                                 np.sin(k[2] * self.a) * alpha_z) +
                       (np.cos(k[0] * self.a) + np.cos(k[1] * self.a) + np.cos(k[2] * self.a) +
                        self.m - 3) * beta)

        return hamiltonian


class WilsonAmorphous(System):
    """
    Implementation of 3D Wilson-fermion model in real-space. This models is a generalization of the
    standard Wilson-fermion model, which uses a cubic unit cell. Instead, this allows for any spatial
    distribution of the atoms, and since the Hamiltonian is built in real-space any number of atoms
    can be considered.

    :param side: First-neighbour distance in the crystalline system. Defaults to 1.
    :param t: Hopping parameter. Defaults to 1.
    :param m: Mass term. Defaults to 1.
    :param r: Cutoff distance for neighbour interaction. Defaults to 1.
    """

    def __init__(self, side=1, t=1, m=1, r=1.1):
        super().__init__(system_name="Wilson Amorphous model",
                         crystal=Crystal([[side, 0, 0], [0, side, 0], [0, 0, side]],
                                         motif=[[0, 0, 0, 0]]))

        self.filling = 2
        self.norbitals = [4]
        self.basisdim = self.norbitals[0] * len(self.motif)
        self.boundary = "PBC"

        self.a = side
        self.t = t
        self.m = m
        self.r = r
        self.parameters = {"C": self.t, "a": self.a, "M": self.m}

    @staticmethod
    def _hopping_matrix(initial_position: np.ndarray, final_position: np.ndarray, parameters: dict) -> np.ndarray:
        """ 
        Computes hopping matrix according to Wilson-fermion model for two
        given atomic positions.

        :param initial_position: Array 1x3 initial position.
        :param final_position: Array 1x3 final position.
        :param parameters: Dictionary{C, a, M}
        :return: Hopping matrix.
        """

        x, y, z = final_position - initial_position
        r = math.sqrt(x ** 2 + y ** 2 + z ** 2)
        phi = math.atan2(y, x)
        theta = math.acos(z / r)

        hopping = np.zeros([4, 4], dtype=np.complex_)
        np.fill_diagonal(hopping, [1, 1, -1, -1])
        offdiag_block = np.array([[-1j * math.cos(theta), -1j * cmath.exp(-1j * phi) * math.sin(theta)],
                                  [-1j * cmath.exp(1j * phi) * math.sin(theta), 1j * math.cos(theta)]],
                                 dtype=np.complex_)
        hopping[0:2, 2:4] = offdiag_block
        hopping[2:4, 0:2] = offdiag_block
        hopping *= 0.5
        hopping *= parameters["C"] * math.exp(1 - r / parameters["a"])

        return hopping

    def initialize_hamiltonian(self, find_bonds: bool = True) -> None:
        """ 
        Routine to initialize the matrices that compose the Bloch Hamiltonian
        
        :param find_bonds: Toggle to find bonds in the system before initializating
            the Hamiltonian.
        """

        self.basisdim = self.natoms * self.norbitals[0]
        norbitals = 4
        if find_bonds:
            print("Computing neighbours...")
            self.find_neighbours(mode="radius", r=self.r)
        if not find_bonds and self.bonds is None:
            raise ArithmeticError("No bonds found to initialize Hamiltonian, exiting...")
        self._determine_connected_unit_cells()

        hamiltonian = []
        for _ in self._unit_cell_list:
            hamiltonian.append(np.zeros([self.basisdim, self.basisdim], dtype=np.complex_))

        hamiltonian_atom_block = np.diag(np.array([-3 + self.m, -3 + self.m,
                                                   3 - self.m, 3 - self.m])*0.5)
        for n, atom in enumerate(self.motif):
            hamiltonian[0][norbitals * n:norbitals * (n + 1),
                           norbitals * n:norbitals * (n + 1)] = hamiltonian_atom_block

        for bond in self.bonds:
            initial_atom_index, final_atom_index, cell, _ = bond
            initial_atom = self.motif[initial_atom_index][:3]
            final_atom = self.motif[final_atom_index][:3]
            neigh_position = np.array(final_atom) + np.array(cell)
            h_cell = self._unit_cell_list.index(list(cell))
            hamiltonian[h_cell][4 * initial_atom_index:4 * (initial_atom_index + 1),
                                4 * final_atom_index:4 * (final_atom_index + 1)] = \
                self._hopping_matrix(initial_atom, neigh_position, self.parameters)

        self.hamiltonian = hamiltonian

    def hamiltonian_k(self, k: Union[list, np.ndarray]) -> np.ndarray:
        """
        Routine to evaluate the Bloch Hamiltonian at a given k point. It adds the k dependency of the Bloch Hamiltonian
        through the complex exponentials.

        :param k: k vector (Array 1x3)
        :return: Bloch Hamiltonian matrix
        """

        dimension = len(self.hamiltonian[0])
        hamiltonian_k = np.zeros((dimension, dimension), dtype=np.complex_)
        for cell_index, cell in enumerate(self._unit_cell_list):
            hamiltonian_k += (self.hamiltonian[cell_index] * cmath.exp(1j * np.dot(k, cell)))

        hamiltonian_k = (hamiltonian_k + np.transpose(np.conjugate(hamiltonian_k)))
        return hamiltonian_k
    

class HaldaneModel(System, FrozenClass):
    """
    Implementation of Haldane model in reciprocal space.

    :param t1: Strength of first-neighbour hoppings.
    :param t2: Strength of second-neighbour hoppings.
    :param phi: Phase of time-reversal breaking term.
    :param m: Mass term.
    """

    def __init__(self, t1: float = -1, t2: float = 0, phi: float = 0, m: float = 0):
        
        lattice_vectors = np.array([[np.sqrt(3)/2,  1./2, 0], 
                                    [np.sqrt(3)/2, -1./2, 0]])
        motif = np.array([[            0, 0, 0, 0],
                          [1./np.sqrt(3), 0, 0, 0]])
        super().__init__(system_name="Haldane model", crystal=Crystal(lattice_vectors, motif))

        self.t1  = t1
        self.t2  = t2
        self.phi = phi
        self.m   = m 

        

        self.norbitals = [1]
        self.basisdim  = self.norbitals[0] * len(self.motif)
        self.boundary  = "PBC"
        self.filling   = 1

        self._freeze()

    def hamiltonian_k(self, k: Union[List, np.ndarray]) -> np.ndarray:
        """
        Routine to evaluate the Bloch Hamiltonian at a given k point. It adds the k dependency of the Bloch Hamiltonian
        through the complex exponentials.

        :param k: k vector (Array 1x3)
        :return: Bloch Hamiltonian matrix
        """

        global sigma_x, sigma_y, sigma_z
        identity = np.eye(2)

        a = 1./np.sqrt(3)
        angle = 60 * np.pi / 180
        a1 = np.array([a, 0., 0.]) 
        a2 = np.array([-a * np.cos(angle), a * np.sin(angle), 0]) 
        a3 = np.array([-a * np.cos(angle), -a * np.sin(angle), 0])

        first_neighbours = [a1, a2, a3]
        second_neighbours = [-a2 + a3, -a3 + a1, -a1 + a2]

        h = np.zeros([2, 2], dtype=np.complex_)

        h += 2 * self.t2 * np.cos(self.phi) * np.sum([np.cos(np.dot(k, b)) for b in second_neighbours]) * identity
        h += self.t1 * np.sum([np.cos(np.dot(k, a)) for a in first_neighbours]) * sigma_x
        h += self.t1 * np.sum([np.sin(np.dot(k, a)) for a in first_neighbours]) * sigma_y
        h += (self.m - 2 * self.t2 * np.sin(self.phi) * np.sum([np.sin(np.dot(k, b)) for b in second_neighbours])) * sigma_z
        
        return h
    

class AgarwalaChern(System):
    """
    Implementation of the Agarwala model of amorphous Chern insulators in 2D, which is a generalization 
    of the BHZ model for arbitrary atomic positions, with additional terms to break particle-hole and inversion symmetry.

    :param side: First-neighbour distance in the crystalline system. Defaults to 1.
    :param t: Hopping parameter. Defaults to 1.
    :param m: Mass term. Defaults to 1.
    :param t2: Particle-hole symmetry breaking parameter. Defaults to 0.
    :param lda: Inversion symmetry breaking parameters. Defaults to 0
    :param r: Cutoff distance for neighbour interaction. Defaults to side*1.1.
    """

    def __init__(self, side=1, t=1, m=1, t2=0, lda=0,  r=1.1):
        super().__init__(system_name="Wilson Amorphous model",
                         crystal=Crystal([[side, 0, 0], [0, side, 0]],
                                         motif=[[0, 0, 0, 0]]))

        self.filling = 1
        self.norbitals = [2]
        self.basisdim = self.norbitals[0] * len(self.motif)
        self.boundary = "PBC"

        self.a = side
        self.t = t
        self.m = m
        self.t2 = t2
        self.lda = lda
        self.r = r * self.a
        self.parameters = {"C": self.t, "a": self.a, "M": self.m, "t2": self.t2, "lambda": lda}

    def _hopping_matrix(self, initial_position: np.ndarray, final_position: np.ndarray) -> np.ndarray:
        """ 
        Computes hopping matrix according to Wilson-fermion model for two
        given atomic positions.

        :param initial_position: Array 1x3 initial position.
        :param final_position: Array 1x3 final position.
        :return: Hopping matrix.
        """

        x, y, z = final_position - initial_position
        r = math.sqrt(x ** 2 + y ** 2 + z ** 2)
        theta = math.atan2(y, x)

        hopping = np.zeros([2, 2], dtype=np.complex_)
        hopping[0, 0] = -1 + self.t2
        hopping[1, 1] =  1 + self.t2
        hopping[0, 1] = -1j*cmath.exp(-1j*theta) + self.lda * ((1 + 1j)*math.sin(theta)**2 - 1)
        hopping[1, 0] = -1j*cmath.exp(1j*theta) + self.lda * ((1 - 1j)*math.sin(theta)**2 - 1)
        
        hopping *= 0.5 * math.exp(-(r - self.a))

        return hopping

    def initialize_hamiltonian(self, find_bonds: bool = True) -> None:
        """ 
        Routine to initialize the matrices that compose the Bloch Hamiltonian
        
        :param find_bonds: Toggle to find bonds in the system before initializating
            the Hamiltonian.
        """

        self.basisdim = self.natoms * self.norbitals[0]
        if find_bonds:
            print("Computing neighbours...")
            self.find_neighbours(mode="radius", r=self.r)
        if not find_bonds and self.bonds is None:
            raise ArithmeticError("No bonds found to initialize Hamiltonian, exiting...")
        self._determine_connected_unit_cells()

        hamiltonian = []
        for _ in self._unit_cell_list:
            hamiltonian.append(np.zeros([self.basisdim, self.basisdim], dtype=np.complex_))

        hamiltonian_atom_block = np.array([[       2 + self.m, (1 - 1j)*self.lda], 
                                           [(1 + 1j)*self.lda,       -2 - self.m]])

        for n, atom in enumerate(self.motif):
            hamiltonian[0][self.norbitals[0] * n:self.norbitals[0] * (n + 1),
                           self.norbitals[0] * n:self.norbitals[0] * (n + 1)] = hamiltonian_atom_block

        for bond in self.bonds:
            initial_atom_index, final_atom_index, cell, _ = bond
            initial_atom = self.motif[initial_atom_index][:3]
            final_atom = self.motif[final_atom_index][:3]
            neigh_position = np.array(final_atom) + np.array(cell)
            h_cell = self._unit_cell_list.index(list(cell))
            hamiltonian[h_cell][2 * initial_atom_index:2 * (initial_atom_index + 1),
                                2 * final_atom_index:2 * (final_atom_index + 1)] = \
                self._hopping_matrix(initial_atom, neigh_position)

        self.hamiltonian = hamiltonian

    def hamiltonian_k(self, k: Union[list, np.ndarray]) -> np.ndarray:
        """
        Routine to evaluate the Bloch Hamiltonian at a given k point. It adds the k dependency of the Bloch Hamiltonian
        through the complex exponentials.

        :param k: k vector (Array 1x3)
        :param conditions: defaults to PBC. Can be either PBC or OBC
        :return: Bloch Hamiltonian matrix
        """

        dimension = len(self.hamiltonian[0])
        hamiltonian_k = np.zeros((dimension, dimension), dtype=np.complex_)
        for cell_index, cell in enumerate(self._unit_cell_list):
            hamiltonian_k += (self.hamiltonian[cell_index] * cmath.exp(1j * np.dot(k, cell)))

        return hamiltonian_k


class RealSpace(System):
    """ 
    Class to construct toy models, in which one sets the hoppings manually. This models
    are by default OBC; by setting one hopping between different unit cells it automatically becomes
    PBC. 
    
    :param system_name: Name of system. Defaults to None.
    :param bravais_lattice: Bravais basis.
    :param motif: List of atoms of the solid.
    """
    
    def __init__(self, system_name: str = None, bravais_lattice: np.ndarray = None, motif: np.ndarray = None):
        super().__init__(system_name=system_name, bravais_lattice=bravais_lattice, motif=motif)
        self.onsite_energy = None
        self.hoppings = {}
        self.boundary = "OBC"
        self.orbitals = None
        self.bond_graphs = None

    def add_onsite_energy(self, index: int, energy: float) -> None:
        """ 
        Method to specify the onsite energy of a given atom. 
        
        :param index: Index of atom.
        :param energy: Value of onsite energy.
        """
        
        self.hoppings[index] = energy

    def add_onsite_energies(self, indices: List[float], energies: List[float]) -> None:
        """
        Same as add_onsite_energy. Used to specify the onsite energies of a collection
        of atom indices. 
        
        :param indices: List with indices of atoms.
        :param energies: List with onsite energies.
        """

        for index, energy in zip(indices, energies):
            self.add_onsite_energy(index, energy)

    def add_hopping(self, hopping: complex, initial: int, final: int, cell: list = [0., 0., 0.]) -> None:
        """ 
        Method to add a hopping between two atoms of the motif.
        NB: The hopping has a specified direction, from initial to final. Since the final
        Hamiltonian is computed taking into account hermiticity, it is not necessary to specify the hopping
        in the other direction.
        Parameters:
        :param hopping: Value of hopping.
        :param initial: Index of first atom of bond.
        :param final: Index of final atom of bond.
        :param cell: Bravais vector connecting the cells of the two atoms. Defaults to zero. 
        """

        assert type(initial) == int, "initial must be an integer"
        assert type(final)   == int, "initial must be an integer"
        assert type(hopping) != str, "hopping must be a complex number"

        self.add_bond(initial, final, cell)
        bond_index = len(self.bonds) - 1
        self.hoppings.append([bond_index, hopping])

    def add_hoppings(self, hoppings: List[complex], initial: List[int], final: List[int], cells: List[list] = None) -> None:
        """ 
        Same method as add_hopping but we input a list of hoppings at once.
        Parameters:
        :param hoppings: List with values of hoppings.
        :param initial: List of indices of initial atom of the hoppings.
        :param final: List of indices of final atoms of the motif.
        :param cells: Each row denotes the Bravais vector connecting two cells. Defaults to None. 
        """
        
        if len(hoppings) != len(initial) or len(initial) != len(final):
            raise ValueError("Provided list must have the same length")

        if cells is None:
            cells = np.zeros([len(hoppings), 3])
        else:
            self.boundary = "PBC"
            cells = np.array(cells)
        for n, hopping in enumerate(hoppings):
            self.add_hopping(hopping, initial[n], final[n], cells[n, :])

    def specify_hopping_amplitude(self, amplitude: float, bond_index: int) -> None:
        """ 
        Instead of adding the hopping directly as a bond with an amplitude, we
        can also specify the hopping amplitude of a preexisting bond. 
        
        :param amplitude: Hopping amplitude.
        :param bond_index: Index of bond.
        """

        self.hoppings.append([bond_index, amplitude])

    def specify_hopping_amplitudes(self, amplitudes: List[float], bond_indices: List[int]) -> None:
        """ 
        Same as specify_hopping_amplitude but for a list of amplitudes-bond indices. 
        
        :param amplitudes: List of amplitudes of bonds.
        :param bond_indices: List of indices of bonds.
        """
        
        
        for amplitude, bond_index in zip(amplitudes, bond_indices):
            self.specify_hopping_amplitude(amplitude, bond_index)

    def __construct_bond_graph(self) -> None:
        """ 
        Private method to create the graph of bonds between all atoms 
        """
        
        bond_graphs = []
        for _ in self._unit_cell_list:
            bond_graphs.append(np.zeros([self.natoms, self.natoms]))

        for bond, _ in self.hoppings:
            initial_atom, final_atom, cell = bond
            cell_index = self._unit_cell_list.index(cell)
            bond_graphs[cell_index][initial_atom, final_atom] += 1

        for i in range(self._unit_cell_list):
            bond_graphs[i] += bond_graphs[i].T

        self.bond_graphs = bond_graphs

    def __calculate_orbitals(self) -> None:
        """ 
        Private method to compute the dimension of the basis for the model, taking into account
        the number of bonds present for each atom of the motif. 
        """
        
        orbitals = []
        for i in range(self.natoms):
            max_orbitals = 0
            for j in range(len(self._unit_cell_list)):
                max_bonds_in_unit_cell = np.max(self.bond_graphs[j][i, :])
                if max_bonds_in_unit_cell > max_orbitals:
                    max_orbitals = max_bonds_in_unit_cell
            orbitals.append(max_orbitals)

        orbitals[orbitals == 0] = 1
        self.orbitals = orbitals

    def __calculate_basisdim(self) -> None:
        """ 
        Method to compute the dimension of the basis based on the orbitals of each atom 
        """
        
        self.basisdim = np.sum(self.orbitals)

    def initialize_hamiltonian(self) -> None:
        """
        Method to set up the matrices that compose the Hamiltonian, either the Bloch Hamiltonian
        or the real-space one. 
        TODO initialize_hamiltonian RSmodel
        """

        # First call private methods to compute orbitals per atom and basis dimension
        self._determine_connected_unit_cells()
        self.__calculate_orbitals()
        self.__calculate_basisdim()

        # Build the matrices that compose the Bloch Hamiltonian
        assert len(self.bonds) == len(self.hoppings)
        hamiltonian = []
        for _ in range(self._unit_cell_list):
            hamiltonian.append(np.zeros((self.basisdim, self.basisdim), dtype=np.complex))

        # Onsite energies
        for n in range(len(self.natoms)):
            hamiltonian[0][n:n + self.orbitals[n], n:n+self.orbitals[n]] = self.onsite_energy[n]/2.

        # Hoppings
        for bond_index, amplitude in self.hoppings:
            bond = self.bonds[bond_index]
            initial_atom_index, final_atom_index, cell = bond

    def hamiltonian_k(self, k):
        hamiltonian_k = np.zeros((self.basisdim, self.basisdim), dtype=np.complex_)
        for cell_index, cell in enumerate(self._unit_cell_list):
            hamiltonian_k += (self.hamiltonian[cell_index] * cmath.exp(1j * np.dot(k, cell)))

        return hamiltonian_k


class Stack(System):
    """ 
    Class that implements stacking of bidimensional systems. Its main task is stacking
    correctly the different layers and adding the hoppings between layers. The hamiltonian within
    each layer is delegated upon the class responsible of creating given system. 

    :param height: Distance between layers.
    :param bottom_layer: System corresponding to the bottom layer.
    :param top_layer: System corresponding to the top layer. Defaults to None.
    """

    def __init__(self, height, bottom_layer: System, top_layer: System = None) -> None:
        if bottom_layer.ndim != 2 or (top_layer and top_layer.ndim != 2):
            raise ValueError("Systems must be two-dimensional")
        if height < 0:
            raise ValueError("Interlayer height must be a positive number")

        if top_layer:
            for vector_bottom, vector_top in zip(bottom_layer.bravais_lattice, top_layer.bravais_lattice):
                if np.linalg.norm(vector_bottom - vector_top) > 1E-2:
                    raise ValueError('Both systems must have the same Bravais lattice')
        else:
            top_layer = bottom_layer

        super().__init__()
        self.ndim         = 2
        self.bottom_layer = bottom_layer
        print(self.bottom_layer)
        self.top_layer    = top_layer
        self.height       = height
        self.vertical = False if top_layer else True
        self.__interlayer_hopping = 0
        self.initialize_system_attributes()  

    @property
    def filling(self):
        nelectrons = self.bottom_layer.filling + self.top_layer.filling
        return nelectrons

    @property
    def interlayer_hopping(self):
        return self.__interlayer_hopping

    @interlayer_hopping.setter
    def interlayer_hopping(self, hopping: int):
        """ 
        Setter for interlayer hopping. Note that it is a single value instead
        of multiple values for different orbitals as in SK models. 
        """
        
        self.__interlayer_hopping = hopping

    def initialize_system_attributes(self) -> None:
        """ 
        Routine to initialize all System atttributes from the bottom layer.
        """

        # Bravais lattice
        self.bravais_lattice = self.bottom_layer.bravais_lattice

        # Motif
        top_layer_motif = np.copy(self.top_layer.motif)
        top_layer_motif[:, 2] += np.max(np.array(self.bottom_layer.motif)[:, 2]) + self.height
        self.motif  = np.concatenate((self.bottom_layer.motif, top_layer_motif), axis=0)
        self.natoms = len(self.motif)

    def __find_closest_atom_from_next_layer(self, atom_index: int) -> int:
        """ 
        Routine to find closest atom from next layer. 
        
        :param atom_index: Index of atom from bottom layer.
        :return: Closest atom from top layer.
        """

        atom_position = self.bottom_layer.motif[atom_index][:3]
        distances = np.linalg.norm(np.array(self.top_layer.motif)[:, :3] + self.height - atom_position, axis=1)
        closest_atom_index = np.argmin(distances)

        return closest_atom_index

    def find_interlayer_bonds(self, vertical: bool = False) -> list:
        """ 
        Routine to find list of all bonds between layers. If vertical, bonds are
        created directly without search. 

        :param vertical: If True, the bond is taken simply as the same atom but from the
            top layer, i.e. it is for two equal layers. Defaults to False.
        :return: List of interlayer bonds.
        """

        interlayer_bonds = []
        origin = np.array([0., 0., 0.])
        for atom_index in range(self.bottom_layer.natoms):
            if vertical:
                next_layer_atom_index = atom_index
            else:
                next_layer_atom_index = self.__find_closest_atom_from_next_layer(atom_index)
            interlayer_bonds.append([atom_index, next_layer_atom_index + self.bottom_layer.natoms, origin, 1])
        
        return interlayer_bonds
    
    def initialize_hamiltonian(self):
        """ 
        Init. hamiltonian method for Stack class. 
        """ 

        # Check if interlayer hopping was specified
        if self.interlayer_hopping == 0:
            print('Warning: Interlayer hopping is zero')

        # First initialize hamiltonian of each layer and bonds between layers
        self.bottom_layer.initialize_hamiltonian()
        self.top_layer.initialize_hamiltonian()
        self._unit_cell_list = self.bottom_layer._unit_cell_list
        self.interlayer_bonds = self.find_interlayer_bonds(self.vertical)
        
        # Works only for
        self.norbitals = self.bottom_layer.norbitals

        # Compute basisdim
        self.basisdim = self.bottom_layer.basisdim + self.top_layer.basisdim

        # Now glue the hamiltonians of each layer together
        self.hamiltonian = []
        for cell_index in range(len(self.bottom_layer._unit_cell_list)):
            h_cell = np.zeros([self.basisdim, self.basisdim], dtype=np.complex_)
            h_cell[:self.bottom_layer.basisdim, 
                   :self.bottom_layer.basisdim] = self.bottom_layer.hamiltonian[cell_index]
            h_cell[self.bottom_layer.basisdim:self.basisdim, 
                   self.bottom_layer.basisdim:self.basisdim] = self.top_layer.hamiltonian[cell_index]
            
            self.hamiltonian.append(h_cell)

        # Precompute dictionary with positions of each atom within the hamiltonian
        bottom_atom_index_dictionary, top_atom_index_dictionary = {}, {}
        counter = 0
        for index, atom in enumerate(self.bottom_layer.motif):
            bottom_atom_index_dictionary[index] = counter
            counter += self.bottom_layer.norbitals[int(atom[3])]
        counter = 0
        for index, atom in enumerate(self.top_layer.motif):
            top_atom_index_dictionary[index] = counter
            counter += self.top_layer.norbitals[int(atom[3])]

        # Add hopping between layers
        h = np.zeros([self.basisdim, self.basisdim])
        for initial_atom, final_atom, _, _ in self.interlayer_bonds:
            final_atom -= self.bottom_layer.natoms
            initial_species = int(self.bottom_layer.motif[initial_atom][3])
            final_species   = int(self.top_layer.motif[final_atom][3])
            i = bottom_atom_index_dictionary[initial_atom]
            i_orbitals = self.bottom_layer.norbitals[initial_species]
            f = top_atom_index_dictionary[final_atom] + self.bottom_layer.basisdim
            f_orbitals = self.top_layer.norbitals[final_species]
            h[i : i + i_orbitals, f : f + f_orbitals] = self.interlayer_hopping*np.eye(i_orbitals, f_orbitals)
        
        h += h.T
        self.hamiltonian[0] += h

        # Update bonds in top_layer to match self.motif and glue all
        top_layer_bonds = []
        for initial_atom, final_atom, cell, nn in self.top_layer.bonds:
            new_bond = [initial_atom + self.bottom_layer.natoms, 
                        final_atom + self.bottom_layer.natoms,
                        cell, nn]
            top_layer_bonds.append(new_bond)

        self.bonds = self.bottom_layer.bonds + top_layer_bonds + self.interlayer_bonds


    def hamiltonian_k(self, k):
        hamiltonian_k = np.zeros([self.basisdim, self.basisdim], dtype=np.complex_)
        for cell_index, cell in enumerate(self._unit_cell_list):
            hamiltonian_k += self.hamiltonian[cell_index] * cmath.exp(1j*np.dot(k, cell))

        return hamiltonian_k


def bethe_lattice(z: int = 3, depth: int = 3, length: float = 1) -> Tuple[list, list]:
    """ 
    Routine to generate a Bethe lattice.
    
    :param z: Coordination number. Defaults to 3
    :param depth: Number of shells of the lattice (central atom is depth 0)
    :param length: Float to specify length of bonds
    :return: Motif: list of all atoms' positions 
    """

    motif_previous_shell = [[0., 0., 0., 0]]
    motif = [[0., 0., 0., 0]]
    cell = [0., 0., 0.]
    bonds = []
    previous_angles = [0]
    bond_length = length
    atom_index = 1
    for i in range(1, depth + 1):
        natoms_in_shell = z*(z - 1)**(i - 1)
        angles = np.linspace(0, 2*np.pi, natoms_in_shell + 1)
        angle_between = angles[1] - angles[0]
        angles += angle_between/2 + previous_angles[0]
        for angle in angles[:-1]:
            atom = bond_length*np.array([np.cos(angle), np.sin(angle), 0, 0])
            motif.append(atom)
            distance = np.linalg.norm(atom[:3] - np.array(motif_previous_shell)[:, :3], axis=1)
            neighbour = np.where(distance == np.min(distance))[0][0]
            bonds.append([atom_index, neighbour, cell, 1])
            atom_index += 1

        motif_previous_shell = np.copy(motif)
        bond_length += length
        previous_angles = angles

    return motif, bonds






