#! /usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt

import tbplas as tb


def tbpm_dos(solver: tb.TBPMSolver, analyzer: tb.Analyzer) -> None:
    """
    Calculate DOS with TBPM.

    :param solver: TBPM solver
    :param analyzer: TBPM analyzer
    :return: None
    """
    # Get DOS correlation function with solver.
    corr_dos = solver.calc_corr_dos()

    # Analyze DOS correlation function with analyzer.
    energies_dos, dos = analyzer.calc_dos(corr_dos)

    # Plot DOS
    plt.plot(energies_dos, dos)
    plt.xlabel("Energy (eV)")
    plt.ylabel("DOS")
    plt.savefig("DOS.png")
    plt.close()


def tbpm_ac_cond(solver: tb.TBPMSolver, analyzer: tb.Analyzer) -> None:
    """
    Calculate AC conductivity with TBPM.

    :param solver: TBPM solver
    :param analyzer: TBPM analyzer
    :return: None
    """
    corr_ac = solver.calc_corr_ac_cond()
    omegas_ac, ac = analyzer.calc_ac_cond(corr_ac)
    plt.plot(omegas_ac, ac[0].real)
    plt.xlabel("Energy (eV)")
    plt.ylabel("sigma_xx")
    plt.savefig("ACxx.png")
    plt.close()


def main():
    # In this tutorial we show how to calculate different properties using the
    # TBPM algorithms implemented in TBPlaS. First, we create a large graphene
    # sample with periodic boundary conditions.
    prim_cell = tb.make_graphene_diamond()
    super_cell_pbc = tb.SuperCell(prim_cell, dim=(1024, 1024, 1),
                                  pbc=(True, True, True))
    sample_pbc = tb.Sample(super_cell_pbc)

    # Create the solver and analyzer
    solver = tb.TBPMSolver(sample_pbc)
    analyzer = tb.Analyzer(f"{solver.config.prefix}_info.dat")

    # The Hamiltonian of sample must be rescaled for any TBPM calculation.
    # Otherwise, the wave function will diverge during propagation. A scaling
    # factor is need to scale the Hamiltonian. If the factor is not specified,
    # it will be estimated by inspecting the Hamiltonian automatically.
    solver.config.rescale = 9.0

    # We will use 4 random samples. For each sample the propagation will take
    # 1024 steps.
    solver.config.num_random_samples = 1
    solver.config.num_time_steps = 1024

    # Calculate DOS
    tbpm_dos(solver, analyzer)

    # Now we demonstrate more features of TBPM algorithms. First, we need to put
    # more parameters into configuration.
    solver.config.num_spin = 2

    # Then calculate ac_cond
    tbpm_ac_cond(solver, analyzer)


if __name__ == "__main__":
    main()
