#! /usr/bin/env python

from typing import Tuple

import numpy as np

import tbplas as tb


def make_graphene_prim_cell() -> tb.PrimitiveCell:
    """Make graphene primitive cell."""
    # Model parameters
    t = -2.7
    lat = 0.246
    f = 1.0 / 3
    onsite = 0.0

    # Generate lattice vectors
    lat_vec = tb.gen_lattice_vectors(a=lat, b=lat, c=1.0, gamma=60)

    # Create primitive cell
    prim_cell = tb.PrimitiveCell(lat_vec, unit=tb.NM)

    # Add orbitals
    prim_cell.add_orbital((0.0, 0.0, 0.0), energy=onsite, label="C_pz")
    prim_cell.add_orbital((f, f, 0.0), energy=onsite, label="C_pz")

    # Add hopping terms
    prim_cell.add_hopping((0, 0, 0), 0, 1, t)
    prim_cell.add_hopping((1, 0, 0), 1, 0, t)
    prim_cell.add_hopping((0, 1, 0), 1, 0, t)
    return prim_cell


def make_graphene_sample(dim: Tuple[int, int, int]) -> tb.Sample:
    """
    Make graphene sample with specific dimension.
    :param dim: dimension along a, b and c directions
    """
    # Create primitive cell
    prim_cell = make_graphene_prim_cell()

    # Create supercell with specific dimension
    super_cell = tb.SuperCell(prim_cell, dim=dim, pbc=(True, True, False))

    # Assemble supercell to a sample
    sample = tb.Sample(super_cell)
    return sample


def test_diag_bands() -> None:
    """Calculate band structure using diagonalization."""
    # Build the model
    model = make_graphene_prim_cell()

    # Create a solver for the model
    solver = tb.DiagSolver(model)

    # Set up parameters of the solver
    k_points = np.array([
        [0.0, 0.0, 0.0],
        [2./3, 1./3, 0.0],
        [0.5, 0.0, 0.0],
        [0.0, 0.0, 0.0]
    ])
    k_path, k_idx = tb.gen_kpath(k_points, (100, 100, 100))
    solver.config.prefix = "graphene"
    solver.config.k_points = k_path

    # Call 'calc_bands' method of solver to evaluate band structure
    # Data files will be saved automatically.
    timer = tb.Timer()
    timer.tic("bands")
    k_len, bands = solver.calc_bands()
    timer.toc("bands")

    # Report time usage and visualization
    if solver.is_master:
        timer.report_total_time()
        vis = tb.Visualizer()
        vis.plot_bands(k_len, bands, k_idx, ["G", "K", "M", "G"])


def test_tbpm_dos():
    """Calculate DOS using TBPM."""
    # Build the model
    model = make_graphene_sample(dim=(512, 512, 1))

    # Create a solver for the model
    solver = tb.TBPMSolver(model)

    # Set up parameters of the solver
    solver.config.prefix = "graphene"
    solver.config.num_random_samples = 1
    solver.config.rescale = 9.0
    solver.config.num_time_steps = 1024

    # Call 'calc_corr_dos' method of solver to evaluate DOS correlation function
    # Data files will be saved automatically.
    timer = tb.Timer()
    timer.tic("corr_dos")
    corr_dos = solver.calc_corr_dos()
    timer.toc("corr_dos")

    # Report time usage and visualization
    if solver.is_master:
        timer.report_total_time()
        analyzer = tb.Analyzer(f"{solver.config.prefix}_info.dat")
        eng, dos = analyzer.calc_dos(corr_dos)
        vis = tb.Visualizer()
        vis.plot_dos(eng, dos)


if __name__ == "__main__":
    test_diag_bands()
    test_tbpm_dos()
