#! /usr/bin/env python
"""
This example demonstrates how to read the Hamiltonian and overlap matrices
produced by DeepH and calculate the band structure of carbon nanotube with
100 atoms. The data files are not shipped with tbplas. They can be generated
following the instructions at
https://github.com/deepmodeling/DeepH-pack/files/9526304/demo_abacus.zip

CAUTION: The Hamiltonian matrix elements produced by DeepH do not strictly
follow the conjugate relation <i,0|H|j,R> = <j,0|H|i,-R>*. However, this
relation is enthusiastically utilized by TBPLaS to reduce memory and CPU cost.
This consistency may introduce errors in the Hamiltonian matrix in Bloch basis
as well as the eigenvalues and eigen-vectors. Unfortunately, the errors are
strongly case-dependent and not predictable. So, test carefully before
production calculations.
"""
from typing import Tuple, List

import h5py
import numpy as np
import matplotlib.pyplot as plt

import tbplas as tb


class CachedPrimitiveCell(tb.PrimitiveCell):
    def __init__(self, prefix: str) -> None:
        super().__init__()
        self._lat_vec = np.load(f"{prefix}.lat_vec.npy")
        self._origin = np.load(f"{prefix}.origin.npy")
        self._cached_data.orb_pos = np.load(f"{prefix}.orb_pos.npy")
        self._cached_data.orb_eng = np.load(f"{prefix}.orb_eng.npy")
        self._cached_data.hop_ind = np.load(f"{prefix}.hop_ind.npy")
        self._cached_data.hop_eng = np.load(f"{prefix}.hop_eng.npy")

    def _sync_array(self, force_sync: bool = False):
        pass

    @property
    def num_orb(self) -> int:
        return self._cached_data.orb_pos.shape[0]

    @property
    def num_hop(self) -> int:
        return self._cached_data.hop_ind.shape[0]


def not_on_site(rn: Tuple[int, int, int], orb_i: int, orb_j: int) -> bool:
    """
    Check if the hopping term is not an on-site term.
    :param rn: cell index
    :param orb_i: orbital index of i
    :param orb_j: orbital index of j
    :return:
    """
    return rn != (0, 0, 0) or orb_i != orb_j


def duplicate_overlap(value_single: np.ndarray) -> np.ndarray:
    """
    Duplicate overlap matrix in SOC case.

    :param value_single: (num_orb, num_orb) complex128 array
        overlap matrix to duplicate
    :return: (2*num_orb, 2*num_orb)  complex128 array
        duplicated overlap matrix
    """
    num_orb = value_single.shape[0]
    value_double = np.zeros((2*num_orb, 2*num_orb), dtype=np.complex128)
    # See read_deep_h for the order of orbitals.
    # overlap between different spins are 0.
    for i in range(num_orb):
        for j in range(num_orb):
            t = value_single[i, j]
            value_double[i, j] = t
            value_double[i+num_orb, j+num_orb] = t
    return value_double


def read_h5(target: tb.PrimitiveCell,
            h5_file: str,
            num_orb: List[int],
            dup_overlap: bool = False,
            cutoff: float = 1.0e-5) -> None:
    """
    Read Hamiltonian or overlap from h5 files.
    :param target: primitive cell or overlap instance
    :param h5_file: name of h5 file
    :param num_orb: number of orbitals on each atom
    :param dup_overlap: whether to duplicate the overlap matrix, enabled in
        SOC case
    :param cutoff: cutoff for elements
    """
    ind_start = [sum(num_orb[:_]) for _ in range(len(num_orb))]
    fid = h5py.File(h5_file, "r")
    for key, value in fid.items():
        value = np.array(value)
        if dup_overlap:
            value = duplicate_overlap(value)

        # Parse rn and atom indices
        key = key.lstrip("[").rstrip("]").split(",")
        rn = (int(key[0]), int(key[1]), int(key[2]))
        atom_i = int(key[3]) - 1
        atom_j = int(key[4]) - 1

        # Loop over orbitals on atoms to add hopping terms
        for i in range(num_orb[atom_i]):
            orb_i = ind_start[atom_i] + i
            for j in range(i, num_orb[atom_j]):
                orb_j = ind_start[atom_j] + j
                energy = value.item(i, j)
                if abs(energy) > cutoff:
                    if not_on_site(rn, orb_i, orb_j):
                        target.add_hopping(rn, orb_i, orb_j, energy)
                    else:
                        target.set_orbital(orb_i, energy=energy.real)
    fid.close()


def read_deep_h(data_dir: str = ".",
                with_soc: bool = False,
                hop_eng_cutoff: float = 1.0e-9,
                overlap_cutoff: float = 1.0e-9) -> Tuple[tb.PrimitiveCell, tb.PrimitiveCell]:
    """
    Read primitive cell and overlap from output of DeepH.

    :param data_dir: directory name containing the data files
    :param with_soc: whether SOC has been enabled in the model
    :param hop_eng_cutoff: cutoff for hopping terms in eV
    :param overlap_cutoff: cutoff for overlap terms
    :return: (prim_cell, overlap)
    """
    # Read lattice vectors in Angstrom and create the primitive cell/overlap
    # copy to avoid ValueError: ndarray is not C-contiguous
    lat_vec = np.loadtxt(f"{data_dir}/lat.dat").T.copy()
    prim_cell = tb.PrimitiveCell(lat_vec, unit=tb.ANG)
    overlap = tb.PrimitiveCell(lat_vec, unit=tb.ANG)

    # Read atom positions in Angstrom and convert to fractional coordinates
    atom_pos = np.loadtxt(f"{data_dir}/site_positions.dat").T
    atom_pos = tb.cart2frac(lat_vec, atom_pos)

    # Get the number of orbitals on each atom and starting orbital indices
    # Double the orbitals in SOC case.
    num_orb = []
    with open(f"{data_dir}/orbital_types.dat", "r") as infile:
        content = infile.readlines()
        for line in content:
            s = 0
            for v in line.split():
                s += (2 * int(v) + 1)
            num_orb.append(s)
    if with_soc:
        num_orb = [_ * 2 for _ in num_orb]

    # Add orbitals to the primitive cell
    # Order of orbitals:
    #   atom #0, spin up: orbital #0, orbital #1, ..., orbital #N-1
    #   atom #0, spin dn: orbital #0, orbital #1, ..., orbital #N-1
    #   atom #1, spin up: orbital #0, orbital #1, ..., orbital #N-1
    #   atom #1, spin dn: orbital #0, orbital #1, ..., orbital #N-1
    # ... ...
    for i, pos in enumerate(atom_pos):
        for j in range(num_orb[i]):
            prim_cell.add_orbital(tuple(pos))
            overlap.add_orbital(tuple(pos))

    # Read Hamiltonian
    timer = tb.Timer()
    timer.tic("read_ham")
    read_h5(prim_cell, f"{data_dir}/hamiltonians_pred.h5", num_orb,
            dup_overlap=False, cutoff=hop_eng_cutoff)
    timer.toc("read_ham")
    timer.tic("read_overlap")
    read_h5(overlap, f"{data_dir}/overlaps.h5", num_orb,
            dup_overlap=with_soc, cutoff=overlap_cutoff)
    timer.toc("read_overlap")
    timer.report_total_time()
    return prim_cell, overlap


def save_array(target: tb.PrimitiveCell, prefix: str) -> None:
    np.save(f"{prefix}.lat_vec", target.lat_vec)
    np.save(f"{prefix}.origin", target.origin)
    np.save(f"{prefix}.orb_pos", target.orb_pos)
    np.save(f"{prefix}.orb_eng", target.orb_eng)
    np.save(f"{prefix}.hop_ind", target.hop_ind)
    np.save(f"{prefix}.hop_eng", target.hop_eng)


def save_model() -> None:
    prim_cell, overlap = read_deep_h(with_soc=True)

    # Save arrays of prim_cell and overlap
    save_array(prim_cell, "prim_cell")
    save_array(overlap, "overlap")


def main():
    prim_cell = CachedPrimitiveCell("prim_cell")
    overlap = CachedPrimitiveCell("overlap")

    # Evaluate band structure
    fermi_energy = -4.32
    k_points = np.array([
        [0.0, 0.0, 0.0],
        [0.5, 0.0, 0.0],
        [0.333333, 0.333333, 0.0],
        [0.0, 0.0, 0.0],
    ])
    k_path, k_idx = tb.gen_kpath(k_points, [20, 20, 20])
    solver = tb.DiagSolver(prim_cell, overlap)
    solver.config.k_points = k_path
    solver.config.diag_algo = "lapacke"
    # solver.config.feast_num_states = 20
    # solver.config.feast_m0_ratio = 3.0
    # solver.config.feast_e_min = -1.5
    # solver.config.feast_e_max = 1.5
    # solver.config.feast_reuse_results = False
    timer = tb.Timer()
    timer.tic("calc_bands")
    k_len, bands = solver.calc_bands()
    timer.toc("calc_bands")
    timer.report_total_time()

    # Subtract the Fermi energy
    bands -= fermi_energy

    # Plot and save results
    def _set_energy_range():
        plt.ylim(-1.5, 1.5)

    if solver.is_master:
        vis = tb.Visualizer()
        vis.plot_bands(k_len, bands, k_idx, ["G", "M", "K", "G"], beautifier=_set_energy_range)


if __name__ == "__main__":
    # save_model()
    main()
