#! /usr/bin/env python
"""
Example for calculating spin texture for bismuth and graphene.
"""
from typing import List

import numpy as np
import matplotlib.pyplot as plt

import tbplas as tb


def plot_sigma_z_band(spin_data: tb.SpinData, ib: int = 0) -> None:
    """
    Plot expectation of sigma_z of specific band as scalar field of kx and ky.
    :param spin_data: spin texture produced of SpinTexture calculator
    :param ib: band index
    :return: None
    """
    # Data aliases
    kpt_cart = spin_data.kpt_cart
    sigma_z = spin_data.sigma_z[:, ib]

    # Plot expectation of sigma_z.
    vis = tb.Visualizer()
    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)


def plot_sigma_xy_band(spin_data: tb.SpinData, ib: int = 0) -> None:
    """
    Plot expectation of sigma_x and sigma_y of specific band as vector field of
    kx and ky.
    :param spin_data: spin texture produced of SpinTexture calculator
    :param ib: band index
    :return: None
    """
    # Data aliases
    kpt_cart = spin_data.kpt_cart
    sigma_x = spin_data.sigma_x[:, ib]
    sigma_y = spin_data.sigma_y[:, ib]

    # Plot expectation of sigma_z.
    vis = tb.Visualizer()
    # Plot expectation of sigma_x and sigma_y.
    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_xy_eng(spin_data: tb.SpinData,
                      e_min: float = -2,
                      e_max: float = -1.9) -> None:
    """
    Plot expectation of sigma_x and sigma_y of specific energy range as vector
    field of kx and ky.
    :param spin_data: spin texture produced of SpinTexture calculator
    :param e_min: lower bound of energy range in eV
    :param e_max: upper bound of energy range in eV
    :return: None
    """
    # Data aliases
    bands = spin_data.bands
    kpt_cart = spin_data.kpt_cart
    sigma_x = spin_data.sigma_x
    sigma_y = spin_data.sigma_y

    # Extract spin texture of specific energy range
    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])

    # Plot expectation of sigma_x and sigma_y.
    vis = tb.Visualizer()
    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_fat_band(spin_data: tb.SpinData,
                  k_idx: np.ndarray,
                  k_label: List[str],
                  component: str = "z") -> None:
    """
    Plot spin-resolved band structure.
    :param spin_data: spin texture produced of SpinTexture calculator
    :param k_idx: indices of high-symmetric k-points
    :param k_label: labels of high-symmetric k-points
    :param component: component of sigma
    :return: None
    """
    # Data aliases
    bands = spin_data.bands
    kpt_cart = spin_data.kpt_cart
    if component == "x":
        projection = spin_data.sigma_x
    elif component == "y":
        projection = spin_data.sigma_y
    else:
        projection = spin_data.sigma_z

    # Evaluate k_len
    k_len = np.zeros(kpt_cart.shape[0])
    for i in range(1, kpt_cart.shape[0]):
        dk = kpt_cart[i] - kpt_cart[i-1]
        k_len[i] = k_len[i-1] + np.linalg.norm(dk)

    # Plot fat band
    for ib in range(bands.shape[1]):
        plt.scatter(k_len, bands[:, ib], c=projection[:, ib], s=5.0, cmap="jet")
    for x in k_len[k_idx]:
        plt.axvline(x, color="k", linewidth=0.8, linestyle="-")
    for y in (-2.0, -1.9):
        plt.axhline(y, color="k", linewidth=1.0, linestyle=":")
    plt.xlim(k_len[0], k_len[-1])
    plt.xticks(k_len[k_idx], k_label)
    plt.ylabel("Energy (t)")
    plt.colorbar(label=fr"$\langle\sigma_{component}\rangle$")
    plt.show()


def calc_spin_texture(cell: tb.PrimitiveCell,
                      prefix: str = "kane_mele",
                      ib: int = 0,
                      spin_major: bool = True) -> None:
    # Sigma_z on a fine grid
    spin_texture = tb.SpinTexture(cell)
    spin_texture.config.prefix = f"{prefix}_sigma_z"
    spin_texture.config.k_points = 2 * (tb.gen_kmesh((240, 240, 1)) - 0.5)
    spin_texture.config.k_points[:, 2] = 0.0
    spin_texture.config.spin_major = spin_major
    spin_data = spin_texture.calc_spin_texture()
    plot_sigma_z_band(spin_data, ib)

    # Sigma_x and sigma_y on a coarse grid
    spin_texture.config.prefix = f"{prefix}_sigma_xy"
    spin_texture.config.k_points = 2 * (tb.gen_kmesh((48, 48, 1)) - 0.5)
    spin_texture.config.k_points[:, 2] = 0.0
    spin_data = spin_texture.calc_spin_texture()
    plot_sigma_xy_band(spin_data, ib)
    plot_sigma_xy_eng(spin_data)


def calc_fat_band(cell: tb.PrimitiveCell,
                  prefix: str = "kane_mele",
                  spin_major: bool = True) -> None:
    # Generate k-path
    k_points = np.array([
        [0.0, 0.0, 0.0],    # Gamma
        [0.5, 0.0, 0.0],    # M
        [2./3, 1./3, 0.0],  # K
        [0.0, 0.0, 0.0],    # Gamma
        [1./3, 2./3, 0.0],  # K'
        [0.5, 0.5, 0.0],    # M'
        [0.0, 0.0, 0.0],    # Gamma
    ])
    k_label = [r"$\Gamma$", "M", "K", r"$\Gamma$",
               r"$K^{\prime}$", r"$M^{\prime}$", r"$\Gamma$"]
    k_path, k_idx = tb.gen_kpath(k_points, [40, 40, 40, 40, 40, 40])

    # Evaluate spin texture (including energies)
    spin_texture = tb.SpinTexture(cell)
    spin_texture.config.prefix = f"{prefix}_fat_band"
    spin_texture.config.k_points = k_path
    spin_texture.config.spin_major = spin_major
    spin_data = spin_texture.calc_spin_texture()

    # Plot fat band
    plot_fat_band(spin_data, k_idx, k_label, "z")


def main():
    cell = tb.make_graphene_soc(is_qsh=True)
    calc_spin_texture(cell, ib=0, spin_major=False, )
    calc_fat_band(cell, spin_major=False)


if __name__ == "__main__":
    main()
