#! /usr/bin/env python
"""
Plot diagonalization results of given data file, which are generated by
either cpp or python version of speedtest.
"""
import sys

import numpy as np
import matplotlib.pyplot as plt

import tbplas as tb


def get_prefix(file_name: str, separator: str) -> str:
    """
    Extract prefix from file name according to separator.

    For example, extracting from 'sample_spin_sigma_x.dat' with separator of
    'spin_sigma_x' yields prefix of 'sample'

    :param file_name: file name to parse
    :param separator: separator.
    """
    return file_name[:file_name.find(separator) - 1]


def plot_bands(prefix: str):
    k_len, bands = tb.load_bands(prefix)
    num_bands = bands.shape[1]
    for i in range(num_bands):
        plt.plot(k_len, bands[:, i], color="r", linewidth=1.0)
    plt.xlim((0, np.amax(k_len)))
    plt.xlabel("k / (1/nm)")
    plt.ylabel("Energy (eV)")
    plt.tight_layout()
    plt.show()


def plot_dos(prefix: str):
    energy, dos = tb.load_dos(prefix)
    plt.plot(energy, dos)
    plt.grid()
    plt.show()


def plot_dyn_pol(prefix: str):
    omegas, dyn_pol = tb.load_dyn_pol(prefix)
    t = 3.0
    a = 0.142
    plt.plot(omegas/t, -dyn_pol[0].imag*t*a**2, "b")
    plt.grid()
    plt.show()


def plot_epsilon(prefix):
    omegas, epsilon = tb.load_epsilon(prefix)
    plt.plot(omegas, epsilon[1].real, "r")
    plt.grid()
    plt.show()


def plot_ac_cond(prefix: str):
    omegas, ac_cond = tb.load_ac_cond(prefix)
    t = 3.0
    plt.plot(omegas/t, ac_cond.real*4, "r")
    plt.grid()
    plt.show()


def plot_epsilon_q0(prefix):
    omegas, epsilon = tb.load_epsilon(prefix, "epsilon_q0")
    plt.plot(omegas, epsilon[0].real, "r")
    plt.grid()
    plt.show()


def plot_z2(prefix: str, vis: tb.Visualizer):
    kb_array, phases = tb.load_z2(prefix)
    vis.plot_phases(kb_array, phases)
    vis.plot_phases(kb_array, phases, polar=True)


def plot_sigma_band(spin_data: tb.SpinData,
                    vis: tb.Visualizer,
                    ib: int = 0) -> None:
    kpt_cart = spin_data.kpt_cart
    sigma_x = spin_data.sigma_x[:, ib]
    sigma_y = spin_data.sigma_y[:, ib]
    sigma_z = spin_data.sigma_z[:, ib]
    vis.plot_scalar(x=kpt_cart[:, 0], y=kpt_cart[:, 1], z=sigma_z, scatter=True,
                    num_grid=(480, 480), cmap="jet", with_colorbar=True)
    vis.plot_vector(x=kpt_cart[:, 0], y=kpt_cart[:, 1], u=sigma_x, v=sigma_y,
                    cmap="jet", with_colorbar=True)


def plot_sigma_eng(spin_data: tb.SpinData,
                   vis: tb.Visualizer) -> None:
    bands = spin_data.bands
    kpt_cart = spin_data.kpt_cart
    sigma_x = spin_data.sigma_x
    sigma_y = spin_data.sigma_y
    e_min, e_max = -2, -1.9
    ik_selected, sigma_x_selected, sigma_y_selected = [], [], []
    for i_k in range(bands.shape[0]):
        idx = np.where((bands[i_k] >= e_min)&(bands[i_k] <= e_max))[0]
        ik_selected.extend([i_k for _ in idx])
        sigma_x_selected.extend(sigma_x[i_k, idx])
        sigma_y_selected.extend(sigma_y[i_k, idx])
    vis.plot_vector(x=kpt_cart[ik_selected, 0], y=kpt_cart[ik_selected, 1],
                    u=sigma_x_selected, v=sigma_y_selected, cmap="jet",
                    with_colorbar=True)


def plot_omega_xy(berry_data: tb.BerryData,
                  vis: tb.Visualizer,
                  ib: int = 0) -> None:
    kpt_cart = berry_data.kpt_cart
    omega_xy = berry_data.omega_xy[:, ib]
    vis.plot_scalar(x=kpt_cart[:, 0], y=kpt_cart[:, 1], z=omega_xy, scatter=True,
                    num_grid=(480, 480), cmap="jet", with_colorbar=True)


def main(data_file: str, ib: int = 0) -> None:
    """
    Driver of plotting functions.
    :param data_file: file name of data file
    :param ib: band index, for plotting band-specific properties
    """
    # spin and berry should go first since their outputs also have the
    # 'bands' tag.
    vis = tb.Visualizer()
    if data_file.find("spin") != -1:
        prefix = get_prefix(data_file, "spin")
        spin_data = tb.SpinData(prefix)
        plot_sigma_band(spin_data, vis, ib)
        plot_sigma_eng(spin_data, vis)
    elif data_file.find("berry") != -1:
        prefix = get_prefix(data_file, "berry")
        berry_data = tb.BerryData(prefix)
        plot_omega_xy(berry_data, vis, ib)
    elif data_file.find("bands") != -1:
        prefix = get_prefix(data_file, "bands")
        plot_bands(prefix)
    elif data_file.find("dos") != -1:
        prefix = get_prefix(data_file, "dos")
        plot_dos(prefix)
    elif data_file.find("dyn_pol") != -1:
        prefix = get_prefix(data_file, "dyn_pol")
        plot_dyn_pol(prefix)
    elif data_file.find("epsilon_q0") != -1:
        prefix = get_prefix(data_file, "epsilon_q0")
        plot_epsilon_q0(prefix)
    elif data_file.find("epsilon") != -1:
        prefix = get_prefix(data_file, "epsilon")
        plot_epsilon(prefix)
    elif data_file.find("ac_cond") != -1:
        prefix = get_prefix(data_file, "ac_cond")
        plot_ac_cond(prefix)
    elif data_file.find("z2") != -1:
        prefix = get_prefix(data_file, "z2")
        plot_z2(prefix, vis)
    else:
        pass


if __name__ == "__main__":
    if len(sys.argv) == 2:
        main(sys.argv[1])
    else:
        main(sys.argv[1], int(sys.argv[2]))