#! /usr/bin/env python
"""
Example for calculating Chern number and Berry curvature using the 'Berry' class.
"""

import numpy as np

import tbplas as tb


def make_haldane() -> tb.PrimitiveCell:
    """
    Make Haldane model for testing Berry utilities.
    Reference: examples/haldane.py of pythtb
    """
    lat = 1.0
    delta = 0.2
    t = -1.0
    t2 = 0.15 * np.exp(1j * np.pi / 2.)
    t2c = t2.conjugate()

    vectors = tb.gen_lattice_vectors(a=lat, b=lat, c=1.0, gamma=60)
    cell = tb.PrimitiveCell(vectors, unit=tb.NM)
    cell.add_orbital((1 / 3., 1 / 3.), energy=-delta)
    cell.add_orbital((2 / 3., 2 / 3.), energy=delta)
    cell.add_hopping((0, 0), 0, 1, t)
    cell.add_hopping((1, 0), 1, 0, t)
    cell.add_hopping((0, 1), 1, 0, t)
    cell.add_hopping((1, 0), 0, 0, t2)
    cell.add_hopping((1, -1), 1, 1, t2)
    cell.add_hopping((0, 1), 1, 1, t2)
    cell.add_hopping((1, 0), 1, 1, t2c)
    cell.add_hopping((1, -1), 0, 0, t2c)
    cell.add_hopping((0, 1), 0, 0, t2c)
    return cell


def plot_omega_xy(berry_data: tb.BerryData, ib: int = 0) -> None:
    """
    Visualize Berry curvature.
    :param berry_data: Berry curvature to plot
    :param ib: band index of Berry curvature to plot
    :return: None
    """
    kpt_cart = berry_data.kpt_cart
    omega_xy = berry_data.omega_xy[:, ib]
    vis = tb.Visualizer()
    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 calc_chern_number() -> None:
    """Calculate Chern number of Haldane model."""
    # Create model and solver.
    cell = make_haldane()
    solver = tb.Berry(cell)

    # Set up parameters
    solver.config.k_grid_size = (100, 100, 1)
    solver.config.num_occ = 1

    # Calculate Berry curvature and Chern number using Kubo formula
    solver.config.prefix = "chern_kubo"
    solver.calc_berry_curvature_kubo()

    # Calculate Berry curvature and Chern number using Wilson loop
    solver.config.prefix = "chern_wilson"
    solver.calc_berry_curvature_wilson()


def calc_berry_curvature() -> None:
    """Calculate and visualize Berry curvature of Haldane model."""
    # Create model and solver.
    cell = make_haldane()
    solver = tb.Berry(cell)

    # Set up parameters
    solver.config.k_grid_size = (100, 100, 1)
    solver.config.bz_size = (2, 2, 1)
    solver.config.num_occ = 1

    # Calculate Berry curvature and Chern number using Kubo formula
    solver.config.prefix = "berry_kubo"
    data_kubo = solver.calc_berry_curvature_kubo()

    # Calculate Berry curvature and Chern number using Wilson loop
    solver.config.prefix = "berry_wilson"
    data_wilson = solver.calc_berry_curvature_wilson()

    # Visualization
    if solver.is_master:
        plot_omega_xy(data_kubo, ib=0)
        plot_omega_xy(data_kubo, ib=1)
        plot_omega_xy(data_wilson)


if __name__ == "__main__":
    calc_chern_number()
    calc_berry_curvature()
