#! /usr/bin/env python
import numpy as np

import tbplas as tb


def make_cell() -> tb.PrimitiveCell:
    # Ref:
    # Imura et al. Nanoscale Research Letters 2011, 6:358
    a, b = 1.0, 1.0
    delta = 2 * b
    sigma_x = np.array([[0, 1], [1, 0]])
    sigma_y = np.array([[0, -1j], [1j, 0]])
    sigma_z = np.array([[1, 0], [0, -1]])
    gamma_x = -1j * 0.5 * a * sigma_x + b * sigma_z
    gamma_y = -1j * 0.5 * a * sigma_y + b * sigma_z
    energy = delta - 4 * b

    # Spin-up part
    cell_up = tb.PrimitiveCell()
    cell_up.add_orbital((0.0, 0.0), energy, "A+")
    cell_up.add_orbital((0.0, 0.0), -energy, "B+")
    for ii in range(2):
        for jj in range(2):
            cell_up.add_hopping((1, 0), ii, jj, gamma_x.item(ii, jj))
            cell_up.add_hopping((0, 1), ii, jj, gamma_y.item(ii, jj))

    # Kramers-equivalent part
    gamma_x = gamma_x.conjugate()
    gamma_y = gamma_y.conjugate()
    cell_down = tb.PrimitiveCell()
    cell_down.add_orbital((0.0, 0.0), energy, "A-")
    cell_down.add_orbital((0.0, 0.0), -energy, "B-")
    for ii in range(2):
        for jj in range(2):
            cell_down.add_hopping((-1, 0), ii, jj, gamma_x.item(ii, jj))
            cell_down.add_hopping((0, -1), ii, jj, gamma_y.item(ii, jj))

    # Merge-cells
    merged_cell = tb.merge_prim_cell(cell_up, cell_down)
    return merged_cell


def make_cell2() -> tb.PrimitiveCell:
    # Ref:
    # https://www.guanjihuan.com/archives/2327
    lat0 = 1.0
    a = 0.3645 / 5
    b = -0.686 / 25
    c = 0
    d = -0.512 / 25
    m = -0.01
    e_s = c + m - 4 * (d + b) / lat0**2
    e_p = c - m - 4 * (d - b) / lat0**2
    v_ss = (d + b) / lat0**2
    v_pp = (d - b) / lat0**2
    v_sp = -1j * a / (2*lat0)

    cell = tb.PrimitiveCell()
    cell.add_orbital((0.0, 0.0), e_s, "A+")
    cell.add_orbital((0.0, 0.0), e_p, "B+")
    cell.add_orbital((0.0, 0.0), e_s, "A-")
    cell.add_orbital((0.0, 0.0), e_p, "B-")

    hx = np.array([[v_ss, v_sp], [-v_sp.conjugate(), v_pp]])
    hy = np.array([[v_ss, 1j*v_sp], [1j*v_sp.conjugate(), v_pp]])
    for ii in range(2):
        for jj in range(2):
            cell.add_hopping((1, 0), ii, jj, hx.item(ii, jj))
            cell.add_hopping((1, 0), ii+2, jj+2, hx.item(ii, jj).conjugate())
            cell.add_hopping((0, 1), ii, jj, hy.item(ii, jj))
            cell.add_hopping((0, 1), ii+2, jj+2, hy.item(ii, jj).conjugate())
    return cell


def run_bands(cell):
    cell = tb.extend_prim_cell(cell, (1, 50, 1))
    cell.apply_pbc((True, False, False))

    # Evaluate band structure
    k_points = np.array([
        [-0.5, 0.0, 0.0],
        [0.0, 0.0, 0.0],
        [0.5, 0.0, 0.0],
    ])
    k_label = ["X", r"$\Gamma$", "X"]
    k_path, k_idx = tb.gen_kpath(k_points, [40, 40])
    k_len, bands = tb.calc_bands(cell, k_path)
    vis = tb.Visualizer()
    vis.plot_bands(k_len, bands, k_idx, k_label)


def run_z2(cell):
    # Get phases
    z2 = tb.Z2(cell)
    z2.config.ka_array = np.linspace(-0.5, 0.5, 200)
    z2.config.kb_array = np.linspace(-0.5, 0.0, 200)
    z2.config.kc = 0.0
    z2.config.num_occ = 2
    kb_array, phases = z2.calc_phases()
    vis = tb.Visualizer()
    vis.plot_phases(kb_array, phases, scatter=True, polar=False)


if __name__ == "__main__":
    run_bands(make_cell())
    run_z2(make_cell())
