#! /usr/bin/env python

import numpy as np

import tbplas as tb


def h2_molecule():
    """Plot the bonding and anti-boding states of H2 molecule."""
    # 1nm * 1nm * 1nm cubic cell
    lattice = np.eye(3, dtype=np.float64)
    h_bond = 0.074  # H-H bond length in nm
    prim_cell = tb.PrimitiveCell(lattice, unit=tb.NM)
    prim_cell.add_orbital((0.5-0.5*h_bond, 0.5, 0.5))
    prim_cell.add_orbital((0.5+0.5*h_bond, 0.5, 0.5))
    prim_cell.add_hopping((0, 0, 0), 0, 1, -1.0)
    qn = np.array([(1, 1, 0, 0) for _ in range(prim_cell.num_orb)])

    # Calculate wave function
    k_points = np.array([[0.0, 0.0, 0.0]])
    solver = tb.DiagSolver(prim_cell)
    solver.config.prefix = "h2_molecule"
    solver.config.k_points = k_points
    bands, states = solver.calc_states()

    # Define plotting range
    # Cube volume: [0.25, 0.75] * [0.25, 0.75] * [0.25, 0.75] in nm
    cube_origin = np.array([0.25, 0.25, 0.25])
    cube_size = np.array([0.5, 0.5, 0.5])
    rn_max = np.array([0, 0, 0])

    # Plot wave function
    vis = tb.Visualizer()
    vis.plot_wfc3d(prim_cell, wfc=states[0, 0], quantum_numbers=qn,
                   convention=1, k_point=k_points[0], rn_max=rn_max,
                   cube_name="h2.bond.cube", cube_origin=cube_origin,
                   cube_size=cube_size, kind="abs2")
    vis.plot_wfc3d(prim_cell, wfc=states[0, 1], quantum_numbers=qn,
                   convention=1, k_point=k_points[0], rn_max=rn_max,
                   cube_name="h2.anti-bond.cube", cube_origin=cube_origin,
                   cube_size=cube_size, kind="abs2")


def h2_chain():
    """Plot the wave function of a hydrogen chain."""
    # 0.074nm * 0.074nm * 0.074nm cubic cell
    lattice = 0.074 * np.eye(3, dtype=np.float64)
    prim_cell = tb.PrimitiveCell(lattice, unit=tb.NM)
    prim_cell.add_orbital((0.0, 0.0, 0.0))
    prim_cell.add_hopping((1, 0, 0), 0, 0, -1.0)
    qn = np.array([(1, 1, 0, 0) for _ in range(prim_cell.num_orb)])

    # Calculate wave function
    k_points = np.array([[0.0, 0.0, 0.0]])
    solver = tb.DiagSolver(prim_cell)
    solver.config.prefix = "h2_chain"
    solver.config.k_points = k_points
    bands, states = solver.calc_states()

    # Define plotting range
    # Cube volume: [-0.75, 0.75] * [-0.25, 0.25] * [-0.25, 0.25] in nm
    cube_origin = np.array([-0.75, -0.25, -0.25])
    cube_size = np.array([1.5, 0.5, 0.5])
    rn_max = np.array([15, 0, 0])

    # Plot wave function
    vis = tb.Visualizer()
    vis.plot_wfc3d(prim_cell, wfc=states[0, 0], quantum_numbers=qn,
                   convention=1, k_point=k_points[0], rn_max=rn_max,
                   cube_name="h2_chain.cube", cube_origin=cube_origin,
                   cube_size=cube_size, kind="real")


def graphene():
    vectors = tb.gen_lattice_vectors(a=0.246, b=0.246, gamma=60)
    prim_cell = tb.PrimitiveCell(vectors, unit=tb.NM)
    prim_cell.add_orbital((0.0, 0.0), label="C_pz")
    prim_cell.add_orbital((1/3., 1/3.), label="C_pz")
    prim_cell.add_hopping((0, 0), 0, 1, -2.7)
    prim_cell.add_hopping((1, 0), 1, 0, -2.7)
    prim_cell.add_hopping((0, 1), 1, 0, -2.7)
    qn = np.array([(6, 2, 1, 0) for _ in range(prim_cell.num_orb)])

    # Calculate wave function
    k_points = np.array([[0.0, 0.0, 0.0]])
    solver = tb.DiagSolver(prim_cell)
    solver.config.prefix = "graphene"
    solver.config.k_points = k_points
    bands, states = solver.calc_states()

    # Define plotting range
    # Cube volume: [-0.75, 0.75] * [-0.75, 0.75] * [-0.25, 0.25] in nm
    cube_origin = np.array([-0.75, -0.75, -0.25])
    cube_size = np.array([1.5, 1.5, 0.5])
    rn_max = np.array([3, 3, 0])

    # Plot wave function
    vis = tb.Visualizer()
    vis.plot_wfc3d(prim_cell, wfc=states[0, 0], quantum_numbers=qn,
                   convention=1, k_point=k_points[0], rn_max=rn_max,
                   cube_name="graphene.cube", cube_origin=cube_origin,
                   cube_size=cube_size, kind="real")


if __name__ == "__main__":
    h2_molecule()
    h2_chain()
    graphene()
