"""Tetrahedron and GGR Brillouin Zone Density of States Integration Methods

This is based on the MATLAB implementation by Boyuan Liu, re-implemented in
Python by Bradley Dice.
"""

import numpy as np
from tqdm.auto import tqdm


def tetrahedron_integration(
    num_points, num_bands, w_grid, w_min=0, w_max=None, num_bins=1000
):
    """Tetrahedron integration method for computing density of states.

    The tetrahedron integration routine is based on:

    Improved tetrahedron method for Brillouin-zone integrations
    Peter E. Blöchl, O. Jepsen, and O. K. Andersen
    Phys. Rev. B 49, 16223 - Published 15 June 1994

    The inputs to this code are generated by the accompanying file
    ``band_solver.py``, which can be used as an example.

    Note that the variables in this function are named according to Python's
    zero-indexed convention, rather than the one-indexed convention used in
    the corresponding paper and MATLAB code.

    Parameters
    ----------
    num_points : tuple[int]
        A tuple of integers like ``(10, 10, 10)`` indicating the number of
        k-points in each dimension.
    w_grid : numpy.ndarray
        A NumPy array of frequencies for each k-point evaluated. The array
        must have shape ``(num_points[0], num_points[1], num_points[2],
        num_bands)``.
    w_min : float, optional
        Minimum frequency at which to compute the density of states (Default
        value = 0).
    w_max : float, optional
        Maximum frequency at which to compute the density of states. If None,
        the maximum frequency is 1.05 times the largest input frequency
        (Default value = None).
    num_bins : int, optional
        Number of bins to use in the density of states calculation (Default
        value = 1000).

    Returns
    -------
    w_bins : numpy.ndarray
        The frequencies at which the density of states was evaluated. The
        array has shape ``(num_bins,)``.
    DOS : numpy.ndarray
        The density of states corresponding to each evaluated frequency. The
        array has shape ``(num_bins,)``.

    """
    if w_max is None:
        w_max = 1.05 * np.max(w_grid)
    w_bins, w_step = np.linspace(w_min, w_max, num_bins, retstep=True)
    DOS = np.zeros_like(w_bins)
    parallelepiped = np.zeros(8)
    tetrahedra = np.zeros((6, 4))
    v_tetra = 1 / 6 / np.prod(num_points)

    def k_indices(num_points, num_bands):
        for nk1 in range(num_points[0] - 1):
            for nk2 in range(num_points[1] - 1):
                for nk3 in range(num_points[2] - 1):
                    for n_band in range(num_bands):
                        yield nk1, nk2, nk3, n_band

    total = np.prod(np.asarray(num_points) - 1) * num_bands
    for nk1, nk2, nk3, n_band in tqdm(k_indices(num_points, num_bands), total=total):
        parallelepiped[0] = w_grid[nk1 + 1, nk2, nk3, n_band]
        parallelepiped[1] = w_grid[nk1 + 1, nk2 + 1, nk3, n_band]
        parallelepiped[2] = w_grid[nk1, nk2, nk3, n_band]
        parallelepiped[3] = w_grid[nk1, nk2 + 1, nk3, n_band]
        parallelepiped[4] = w_grid[nk1 + 1, nk2, nk3 + 1, n_band]
        parallelepiped[5] = w_grid[nk1 + 1, nk2 + 1, nk3 + 1, n_band]
        parallelepiped[6] = w_grid[nk1, nk2, nk3 + 1, n_band]
        parallelepiped[7] = w_grid[nk1, nk2 + 1, nk3 + 1, n_band]
        tetrahedra[:, [0, 3]] = [parallelepiped[2], parallelepiped[5]]
        tetrahedra[0, [1, 2]] = [parallelepiped[0], parallelepiped[1]]
        tetrahedra[1, [1, 2]] = [parallelepiped[1], parallelepiped[3]]
        tetrahedra[2, [1, 2]] = [parallelepiped[3], parallelepiped[7]]
        tetrahedra[3, [1, 2]] = [parallelepiped[6], parallelepiped[7]]
        tetrahedra[4, [1, 2]] = [parallelepiped[4], parallelepiped[6]]
        tetrahedra[5, [1, 2]] = [parallelepiped[0], parallelepiped[4]]

        for n_tetrahedron in range(6):
            w_corner = sorted(tetrahedra[n_tetrahedron])
            w10 = w_corner[1] - w_corner[0]
            w20 = w_corner[2] - w_corner[0]
            w30 = w_corner[3] - w_corner[0]
            w21 = w_corner[2] - w_corner[1]
            w31 = w_corner[3] - w_corner[1]
            w32 = w_corner[3] - w_corner[2]
            nw_min = round((w_corner[0] - w_min) / w_step)
            nw_max = round((w_corner[3] - w_min) / w_step)
            for nw in range(nw_min, nw_max):
                w_tmpt = w_bins[nw]
                if w30 == 0:
                    DOS[nw] += v_tetra / w_step
                    break
                elif w_tmpt < w_corner[0]:
                    # Zero contribution for w_tmpt < w_corner[0]
                    continue
                elif w_tmpt <= w_corner[1]:
                    if w10 > 0:
                        dos_tmpt = (
                            3
                            * v_tetra
                            * (w_tmpt - w_corner[0]) ** 2
                            / (w10 * w20 * w30)
                        )
                    elif w20 > 0:
                        dos_tmpt = 0
                    else:
                        dos_tmpt = 3 * 2 * v_tetra / w30
                elif w_tmpt <= w_corner[2]:
                    dos_tmpt = (
                        3
                        * v_tetra
                        / (w20 * w30)
                        * (
                            w10
                            + 2 * (w_tmpt - w_corner[1])
                            - (w20 + w31) * (w_tmpt - w_corner[1]) ** 2 / (w21 * w31)
                        )
                    )
                elif w_tmpt <= w_corner[3]:
                    dos_tmpt = (
                        3 * v_tetra * (w_tmpt - w_corner[3]) ** 2 / (w30 * w31 * w32)
                    )
                else:
                    # Zero contribution for w_tmpt > w_corner[3]
                    continue
                # Limit maximum DOS contribution for one tetrahedron
                dos_tmpt = min(dos_tmpt, v_tetra / w_step)
                DOS[nw] += dos_tmpt
    return w_bins, DOS


def ggr_integration(
    num_points, all_freqs, velocities, w_min=0, w_max=None, num_bins=1000
):
    """Generalized Gilat-Raubenheimer method for computing density of states.

    This code is based on:

    Generalized Gilat–Raubenheimer method for density-of-states calculation
    in photonic crystals.
    Boyuan Liu, Steven G Johnson, John D Joannopoulos, and Ling Lu
    J. Opt. 20 044005 - Published 5 March 2018

    The inputs to this code are generated by the accompanying file
    ``band_solver.py``, which can be used as an example.

    Parameters
    ----------
    num_points : tuple[int]
        A tuple of integers like ``(10, 10, 10)`` indicating the number of
        k-points in each dimension.
    all_freqs : numpy.ndarray
        A NumPy array of frequencies for each k-point evaluated. The array
        will be flattened. A typical input shape is ``(num_k_points,
        num_bands)``.
    velocities : numpy.ndarray
        A NumPy array of velocities for each k-point evaluated. The array
        will be reshaped to ``(-1, 3)``. A typical input shape is
        ``(num_k_points, num_bands, 3)``.
    w_min : float, optional
        Minimum frequency at which to compute the density of states (Default
        value = 0).
    w_max : float, optional
        Maximum frequency at which to compute the density of states. If None,
        the maximum frequency is 1.05 times the largest input frequency
        (Default value = None).
    num_bins : int, optional
        Number of bins to use in the density of states calculation (Default
        value = 1000).

    Returns
    -------
    w_bins : numpy.ndarray
        The frequencies at which the density of states was evaluated. The
        array has shape ``(num_bins,)``.
    DOS : numpy.ndarray
        The density of states corresponding to each evaluated frequency. The
        array has shape ``(num_bins,)``.

    """
    if w_max is None:
        w_max = 1.05 * np.max(all_freqs)
    w_bins, w_step = np.linspace(w_min, w_max, num_bins, retstep=True)
    DOS = np.zeros_like(w_bins)

    h_side = 0.5 / np.max(num_points)
    sorted_velocities = -np.sort(-np.abs(velocities), axis=-1)
    flat_freqs = all_freqs.flatten()
    flat_velocities = sorted_velocities.reshape(-1, 3)
    w0 = h_side * np.abs(
        flat_velocities[:, 0] - flat_velocities[:, 1] - flat_velocities[:, 2]
    )
    w1 = h_side * (
        flat_velocities[:, 0] - flat_velocities[:, 1] + flat_velocities[:, 2]
    )
    w2 = h_side * (
        flat_velocities[:, 0] + flat_velocities[:, 1] - flat_velocities[:, 2]
    )
    w3 = h_side * (
        flat_velocities[:, 0] + flat_velocities[:, 1] + flat_velocities[:, 2]
    )

    for w_k, v_k, w0_k, w1_k, w2_k, w3_k in zip(
        tqdm(flat_freqs), flat_velocities, w0, w1, w2, w3
    ):
        n_w_kcenter = np.rint((w_k - w_min) / w_step).astype(int)
        v_mag = np.linalg.norm(v_k)
        flag_delta_n_w = False

        for v_direction in [0, 1]:
            for delta_n_w in range(num_bins):
                n_tmpt = np.rint(n_w_kcenter + (-1) ** v_direction * delta_n_w).astype(
                    int
                )
                delta_w = np.abs(w_k - (n_tmpt * w_step + w_min))
                if delta_w <= w0_k:
                    if v_k[0] >= (v_k[1] + v_k[2]):
                        DOS_contribution = 4 * h_side ** 2 / v_k[0]
                    else:
                        DOS_contribution = (
                            2
                            * h_side ** 2
                            * (v_k[0] * v_k[1] + v_k[1] * v_k[2] + v_k[2] * v_k[0])
                            - (delta_w ** 2 + (h_side * v_mag) ** 2)
                        ) / np.prod(v_k)
                elif delta_w < w1_k:
                    DOS_contribution = (
                        h_side ** 2
                        * (v_k[0] * v_k[1] + 3 * v_k[1] * v_k[2] + v_k[2] * v_k[0])
                        - h_side * delta_w * (-v_k[0] + v_k[1] + v_k[2])
                        - (delta_w ** 2 + h_side ** 2 * v_mag ** 2) / 2
                    ) / np.prod(v_k)
                elif delta_w < w2_k:
                    DOS_contribution = (
                        2
                        * (h_side ** 2 * (v_k[0] + v_k[1]) - h_side * delta_w)
                        / (v_k[0] * v_k[1])
                    )
                elif delta_w < w3_k:
                    DOS_contribution = (
                        (h_side * np.sum(v_k) - delta_w) ** 2 / np.prod(v_k) / 2
                    )
                else:
                    break

                # Limit maximum DOS contribution
                DOS_contribution = min(DOS_contribution, 8 * h_side ** 3 / w_step)

                if delta_n_w == 1:
                    if not flag_delta_n_w:
                        DOS[n_tmpt] += DOS_contribution
                        flag_delta_n_w = True
                    continue
                elif (n_tmpt >= 0) and (n_tmpt < num_bins):
                    DOS[n_tmpt] += DOS_contribution
    return w_bins, DOS
