#! /usr/bin/env python
"""
Example for calculating quasi-eigenstates and comparison with exact eigenstates.
"""

import numpy as np

import tbplas as tb


def make_prim_cell() -> tb.PrimitiveCell:
    """
    Make graphene primitive cell with a single vacancy.

    :return: primitive cell with vacancy
    """
    prim_cell = tb.make_graphene_diamond()

    # Get the index of vacancy
    dim = (17, 17, 1)
    orb_map = tb.RegularOrbitalMap(dim, prim_cell.num_orb)
    vac_id = orb_map.pc2sc(8, 8, 0, 0)

    # Extend cell and remove orbital
    prim_cell = tb.extend_prim_cell(prim_cell, dim=dim)
    prim_cell.remove_orbitals([vac_id])
    return prim_cell


def make_sample() -> tb.Sample:
    """
    Make graphene sample with a single vacancy.

    :return: sample with vacancy
    """
    prim_cell = tb.make_graphene_diamond()

    # Get the index of vacancy
    dim = (17, 17, 1)
    orb_map = tb.RegularOrbitalMap(dim, prim_cell.num_orb)
    vac_id = orb_map.pc2sc(8, 8, 0, 0)

    # Create supercell and sample
    super_cell = tb.SuperCell(prim_cell, dim=dim, pbc=(True, True, False),
                              vacancies=[vac_id])
    sample = tb.Sample(super_cell)
    return sample


def wfc_diag(prim_cell: tb.PrimitiveCell) -> None:
    """
    Calculate wave function using exact diagonalization.

    :param prim_cell: primitive cell
    :return: None
    """
    solver = tb.DiagSolver(prim_cell)
    solver.config.prefix = "diag"
    solver.config.k_points = np.array([[0.0, 0.0, 0.0]])
    bands, states = solver.calc_states()
    i_b = prim_cell.num_orb // 2
    wfc = np.abs(states[0, i_b])**2
    wfc /= wfc.max()
    vis = tb.Visualizer()
    vis.plot_wfc(prim_cell, wfc, scatter=True, site_size=wfc * 100, site_color="r",
                 with_model=True, model_style={"alpha": 0.25, "color": "gray"},
                 fig_name="diag.png", fig_size=(6, 4.5), fig_dpi=100)


def wfc_diag_feast(prim_cell: tb.PrimitiveCell) -> None:
    """
    Calculate wave function using exact diagonalization using feast solver.

    :param prim_cell: primitive cell
    :return: None
    """
    solver = tb.DiagSolver(prim_cell)
    solver.config.prefix = "diag_feast"
    solver.config.k_points = np.array([[0.0, 0.0, 0.0]])
    solver.config.use_feast_dense = True
    # There are 11 states in [-0.6, 0.6] eV, with 5th (counted from 0) state
    # located at 0 eV.
    solver.config.feast_num_states = 11
    solver.config.feast_e_min = -0.6
    solver.config.feast_e_max = 0.6
    bands, states = solver.calc_states()
    i_b = bands.shape[1] // 2
    wfc = np.abs(states[0, i_b])**2
    wfc /= wfc.max()
    vis = tb.Visualizer()
    vis.plot_wfc(prim_cell, wfc, scatter=True, site_size=wfc * 100, site_color="r",
                 with_model=True, model_style={"alpha": 0.25, "color": "gray"},
                 fig_name="diag.png", fig_size=(6, 4.5), fig_dpi=100)


def wfc_tbpm(sample: tb.Sample) -> None:
    """
    Calculate wave function using TBPM.

    :param sample: graphene sample
    :return: None
    """
    solver = tb.TBPMSolver(sample)
    solver.config.prefix = "tbpm"
    solver.config.rescale = 9.0
    solver.config.num_random_samples = 1
    solver.config.num_time_steps = 1024
    solver.config.qe_energies = np.array([0.0])
    qs = solver.calc_quasi_eigenstates()
    wfc = qs[0]  # qs is already the squared norm
    wfc /= wfc.max()
    vis = tb.Visualizer()
    vis.plot_wfc(sample, wfc, scatter=True, site_size=wfc*100, site_color="b",
                 with_model=True, model_style={"alpha": 0.25, "color": "gray"},
                 fig_name="tbpm.png", fig_size=(6, 4.5), fig_dpi=100)

    # Save as reference for C++ version
    np.save("qs_ref.npy", qs)


def main():
    wfc_diag(make_prim_cell())
    wfc_tbpm(make_sample())


if __name__ == "__main__":
    main()
