#! /usr/bin/env python
"""
Post-process and plot TBPM results of given data file, which are generated by
either cpp or python version of speedtest.

NOTE:
1. This script saves post-processed results to *_cpp.dat
2. The python version of speedtest also calls Python TBPMSolver and Analyzer
   automatically and produces *_py.dat.
"""
import sys
from typing import Tuple

import numpy as np

import tbplas as tb

from utils import save_xy, complex2real


def get_prefix_analyzer(file_name: str, separator: str) -> Tuple[str, tb.Analyzer]:
    prefix = file_name[:file_name.find(separator) - 1]
    analyzer = tb.Analyzer(f"{prefix}_info.dat")
    return prefix, analyzer


def plot_haydock(prefix: str, vis: tb.Visualizer) -> None:
    energies, dos = tb.load_ldos_haydock(prefix)
    vis.plot_xy(energies, dos)
    save_xy("ldos_cpp.dat", energies, dos)


def plot_dos(prefix: str, analyzer: tb.Analyzer, vis: tb.Visualizer) -> None:
    corr = tb.load_corr_dos(prefix)
    energies, dos = analyzer.calc_dos(corr)
    vis.plot_dos(energies, dos)
    save_xy("dos_cpp.dat", energies, dos)


def plot_ac_cond(prefix: str, analyzer: tb.Analyzer, vis: tb.Visualizer) -> None:
    # Reference: tbplas/tests/test_unit/plot.py
    t, a = 3.0, 0.142
    corr = tb.load_corr_ac_cond(prefix)
    energies, ac_cond = analyzer.calc_ac_cond(corr)
    vis.plot_xy(energies/t, ac_cond[0].real*4)
    save_xy("ac_cond_cpp.dat", energies/t, ac_cond[0].real*4)


def plot_dyn_pol(prefix: str, analyzer: tb.Analyzer, vis: tb.Visualizer) -> None:
    corr = tb.load_corr_dyn_pol(prefix)
    q, omegas, dp = analyzer.calc_dyn_pol(corr)
    t = 3.0
    a = 0.142
    vis.plot_xy(omegas/t, -dp[0].imag*t*a**2)
    epsilon = analyzer.calc_epsilon(dp)
    vis.plot_xy(omegas, epsilon[1].real)
    save_xy("dyn_pol_cpp.dat", omegas/t, -dp[0].imag*t*a**2)
    save_xy("epsilon_cpp.dat", omegas, epsilon[1].real)


def plot_dc_cond(prefix: str, analyzer: tb.Analyzer, vis: tb.Visualizer) -> None:
    corr_dos, corr_dc = tb.load_corr_dc_cond(prefix)
    energies, dc = analyzer.calc_dc_cond(corr_dos, corr_dc)
    vis.plot_xy(energies, dc[0])
    save_xy("dc_cond_xx_cpp.dat", energies, dc[0])


def plot_hall_cond(prefix: str, vis: tb.Visualizer) -> None:
    energies, cond = tb.load_hall_cond(prefix)
    vis.plot_xy(energies, cond)
    save_xy("hall_cond_cpp.dat", energies, cond)


def plot_quasi_eigen(prefix: str, vis: tb.Visualizer) -> None:
    qs = tb.load_quasi_eigen(prefix)
    for i in range(len(qs)):
        wfc = qs[i]
        x = np.linspace(0, len(wfc)-1, len(wfc))
        vis.plot_xy(x, wfc)
    np.savetxt("qe_cpp.dat", qs.T)


def plot_wft(prefix: str, vis: tb.Visualizer) -> None:
    wft = tb.load_wft(prefix)
    for i in range(len(wft)):
        wfc = np.abs(wft[i])**2
        x = np.linspace(0, len(wfc)-1, len(wfc))
        vis.plot_xy(x, wfc)

    # output
    wft = wft.T
    np.savetxt("wft_cpp.dat", complex2real(wft))


def plot_bands(prefix: str, analyzer: tb.Analyzer, vis: tb.Visualizer) -> None:
    # Assemble data
    k_len, corr_bands = tb.load_corr_bands(prefix)
    energy, bands = analyzer.calc_bands(corr_bands)

    # Plot
    e_min, e_max = -10, 10
    num_grid = (1000, 1000)
    k_idx = [0, 40, 80, 120]
    k_label = ["G", "K", "M", "G"]
    vis.plot_tbpm_bands(energy, bands, k_idx=k_idx, k_len=k_len, k_label=k_label,
                        e_min=e_min, e_max=e_max, num_grid=num_grid, cmap="jet")

    # output
    np.savetxt("energy", energy)
    np.savetxt("bands", bands)
    np.savetxt("k_len", k_len)
    np.savetxt("k_idx", k_idx)


def main(data_file: str) -> None:
    vis = tb.Visualizer()
    if data_file.find("ldos") != -1:
        prefix, analyzer = get_prefix_analyzer(data_file, "ldos")
        plot_haydock(prefix, vis)
    if data_file.find("corr_dos") != -1:
        prefix, analyzer = get_prefix_analyzer(data_file, "corr_dos")
        plot_dos(prefix, analyzer, vis)
    elif data_file.find("corr_ac_cond") != -1:
        prefix, analyzer = get_prefix_analyzer(data_file, "corr_ac_cond")
        plot_ac_cond(prefix, analyzer, vis)
    elif data_file.find("corr_dyn_pol") != -1:
        prefix, analyzer = get_prefix_analyzer(data_file, "corr_dyn_pol")
        plot_dyn_pol(prefix, analyzer, vis)
    elif data_file.find("corr_dc_cond") != -1:
        prefix, analyzer = get_prefix_analyzer(data_file, "corr_dc_cond")
        plot_dc_cond(prefix, analyzer, vis)
    elif data_file.find("hall_cond") != -1:
        prefix, analyzer = get_prefix_analyzer(data_file, "hall_cond")
        plot_hall_cond(prefix, vis)
    elif data_file.find("quasi_eigen") != -1:
        prefix, analyzer = get_prefix_analyzer(data_file, "quasi_eigen")
        plot_quasi_eigen(prefix, vis)
    elif data_file.find("wft") != -1:
        prefix, analyzer = get_prefix_analyzer(data_file, "wft")
        plot_wft(prefix, vis)
    elif data_file.find("corr_bands") != -1:
        prefix, analyzer = get_prefix_analyzer(data_file, "corr_bands")
        plot_bands(prefix, analyzer, vis)
    else:
        pass


if __name__ == "__main__":
    main(sys.argv[1])
