"""Functions for k-points operations."""

from typing import Tuple, List, Union

import numpy as np

from tbplas.base.lattice import gen_reciprocal_vectors, frac2cart


__all__ = ["gen_kpath", "gen_kdist", "gen_kmesh"]


def gen_kpath(hs_kpoints: np.ndarray,
              num_interp: Union[Tuple[int, ...], List[int], np.ndarray]) -> Tuple[np.ndarray, np.ndarray]:
    """
    Generate path in the reciprocal space connecting highly symmetric k-points.
    :param hs_kpoints: (nk, 3) float64 array
        fractional coordinates of highly symmetric k-points
    :param num_interp: (nk-1, ) tuple of integers
        numbers of intermediate k-points between two highly symmetric k-points
    :return: (kpath, hs_index)
        kpath: (sum(num_interp)+1, 3) float64 array
        fractional coordinates of k-points along the path
        hs_index: (nk, ) int32 array
        indices of highly symmetric k-points in kpath
    """
    if not isinstance(hs_kpoints, np.ndarray):
        hs_kpoints = np.array(hs_kpoints)
    assert hs_kpoints.shape[0] == (len(num_interp) + 1)
    kpath = []
    for ik, nk in enumerate(num_interp):
        k0 = hs_kpoints[ik]
        k1 = hs_kpoints[ik+1]
        kpath.append(np.linspace(k0, k1, nk, endpoint=False))
    kpath.append([hs_kpoints[-1]])
    kpath = np.concatenate(kpath)
    hs_index = [sum(num_interp[:_]) for _ in range(len(num_interp) + 1)]
    hs_index = np.array(hs_index)
    return kpath, hs_index


def gen_kdist(lat_vec: np.ndarray, kpoints: np.ndarray) -> np.ndarray:
    """
    Convert k_path generated by gen_path into distances in reciprocal space.
    :param lat_vec: (3, 3) float64 array
        Cartesian coordinates of lattice vectors, with each ROW being a vector
    :param kpoints: (nk, 3) float64 array
        fractional coordinates of kpoints
    :return: (nk, ) float64 array
        distance in reciprocal space in unit of reciprocal lattice vectors
    """
    recip_vec = gen_reciprocal_vectors(lat_vec)
    kpt_cart = frac2cart(recip_vec, kpoints)
    kdist = np.zeros(kpoints.shape[0])
    for i in range(1, kpoints.shape[0]):
        dk = kpt_cart[i] - kpt_cart[i-1]
        kdist[i] = kdist[i-1] + np.linalg.norm(dk)
    return kdist


def gen_kmesh(grid_size: Tuple[int, int, int]) -> np.ndarray:
    """
    Generate uniform mesh grid in the first Brillouin zone.
    :param grid_size: (na, nb, nc)
        dimension of mesh grid along three directions
    :return: (na*nb*nc, 3) float64 array
        fractional coordinates of kpoints in the grid
    """
    assert len(grid_size) == 3
    kmesh = np.array([(kx, ky, kz)
                     for kx in np.linspace(0, 1, grid_size[0], endpoint=False)
                     for ky in np.linspace(0, 1, grid_size[1], endpoint=False)
                     for kz in np.linspace(0, 1, grid_size[2], endpoint=False)])
    return kmesh
