"""Parser for output files generated by taskGroup of GW."""

from ast import literal_eval
from typing import IO, Dict, Tuple

import numpy as np
from numpy.typing import NDArray


def __parse_already_opened_file_with_vector(file: IO[str], dtype: type = complex) -> NDArray[np.complex128]:
    """Parser for files containing vectors

    This file looks like:
     1st line: 1 (rank=1)
     2nd line: two integers i,m with the ranges: vector[i:m]
     next lines: vector elements (complex numbers)

    :param file: object with the already opened file
    :dtype: data type of the vector
    :return: vector read from file
    """
    dim = int(file.readline().split()[0])
    m_ini, m_end = (int(x) for x in file.readline().split())
    assert dim == 1
    m = m_end - m_ini + 1
    vector = np.zeros(m, dtype=dtype)
    counter = 0
    for line in file:
        for data in line.split():
            if dtype is complex:
                vector[counter] = complex(*literal_eval(data))
            else:
                vector[counter] = dtype(data)
            counter += 1
        if counter >= m:
            break
    return vector


def __parse_file_with_vector(file_name: str) -> NDArray[np.complex128]:
    """Parser for files containing vectors.

    See documentation in `__parse_already_opened_file_with_vector`

    :param file_name: name of the file
    :return: vector read from file
    """
    with open(file_name) as file:
        vector = __parse_already_opened_file_with_vector(file)
    return vector


def __parse_already_opened_file_with_matrix(file: IO[str]) -> NDArray[np.complex128]:
    """Parser for files containing matrices.

    This file looks like:
     1st line: 2 (rank=2)
     2nd line: fours integers i,j,m,n with the ranges: matrix[i:m,j:n]
     next lines: matrix elements (complex numbers)

    :param file_name: name of the file
    :return: matrix read from file
    """
    dim = int(file.readline().split()[0])
    m_ini, n_ini, m_end, n_end = (int(x) for x in file.readline().split())
    assert dim == 2 and m_ini == 1 and n_ini == 1, "file should contain a full matrix"
    matrix = np.zeros((m_end, n_end), dtype=complex)

    counter = 0
    for line in file:
        for data in line.split():
            matrix[np.unravel_index(counter, (m_end, n_end), order="F")] = complex(*literal_eval(data))
            counter += 1
    return matrix


def __parse_file_with_matrix(file_name: str) -> NDArray[np.complex128]:
    """Parser for files containing matrices.

    See documentation in `__parse_already_opened_file_with_matrix`

    :param file_name: name of the file
    :return: matrix read from file
    """
    with open(file_name) as file:
        matrix = __parse_already_opened_file_with_matrix(file)
    return matrix


def __parse_file_with_vector_and_matrix(file_name: str) -> Tuple[NDArray[np.complex128], NDArray[np.complex128]]:
    """Parser for files containing vectors

    This file looks like:
     1st line: -1 (indicator that this file stores multiple arrays)
     2nd line: 2 (number of arrays stored)
     next lines: same as a file containing a vector (see `__parse_file_with_vector`)
     next (last) lines: same as a file containing a matrix (see `__parse_file_with_matrix`)

    :param file_name: name of the file
    :return: vector read from file
    """
    indicator_expected = -1
    n_arrays_expected = 2
    with open(file_name) as file:
        indicator = int(file.readline().split()[0])
        assert indicator == indicator_expected
        n_arrays = int(file.readline().split()[0])
        assert n_arrays == n_arrays_expected
        vector = __parse_already_opened_file_with_vector(file, float)
        matrix = __parse_already_opened_file_with_matrix(file)
    return (vector, matrix)


def __parse_file_with_array_of_rank_3(file_name: str) -> NDArray[np.complex128]:
    """Parser for files containing arrays of rank 3.

    This file looks like:
     1st line: 3 (rank=3)
     2nd line: six integers i,j,k,m,n,p with the ranges: matrix[i:m,j:n,k:p]
     next lines: array elements (complex numbers)

    :param file_name: name of the file
    :return: array read from file
    """
    with open(file_name) as file:
        dim = int(file.readline().split()[0])
        m_ini, n_ini, p_ini, m_end, n_end, p_end = (int(x) for x in file.readline().split())
        assert dim == 3 and m_ini == 1 and n_ini == 1 and p_ini == 1, "file should contain a full array"
        array = np.zeros((m_end, n_end, p_end), dtype=complex)

        counter = 0
        for line in file:
            for data in line.split():
                array[np.unravel_index(counter, (m_end, n_end, p_end), order="F")] = complex(*literal_eval(data))
                counter += 1
    return array


def __parse_file_with_array_of_rank_4(file_name: str) -> NDArray[np.complex128]:
    """Parser for files containing arrays of rank 4.

    The file to be parsed should look like:
     1st line: 4 (rank=4)
     2nd line: eight integers i,j,k,l,m,n,p,q with the ranges: tensor[i:m,j:n,k:p,l:q]
     next lines: array elements (complex numbers)

    :param file_name: name of the file
    :return: array read from file
    """
    with open(file_name) as file:
        dim = int(file.readline().split()[0])
        m_ini, n_ini, p_ini, q_ini, m_end, n_end, p_end, q_end = (int(x) for x in file.readline().split())
        assert dim == 4, "file should contain a full array"
        m, n, p, q = (m_end - m_ini + 1, n_end - n_ini + 1, p_end - p_ini + 1, q_end - q_ini + 1)
        array = np.zeros((m, n, p, q), dtype=complex)

        counter = 0
        for line in file:
            for data in line.split():
                array[np.unravel_index(counter, (m, n, p, q), order="F")] = complex(*literal_eval(data))
                counter += 1
    return array


def __square_matrix(a: NDArray) -> NDArray:
    """Square a matrix.

    :param a: input matrix
    :return: the result product
    """
    return np.matmul(a.T.conj(), a)


def parse_barc(file_name: str) -> Dict[str, NDArray[np.complex128]]:
    """Parser for BARC_Q*.OUT, where * is an integer.

    The file contains two arrays:
      - the eigenvalues of the bare Coulomb matrix, and
      - its eigenvectors

    :param file_name: name of the file
    :return: parsed data as dictionary
    """
    eigenvalues, eigenvectors = __parse_file_with_vector_and_matrix(file_name)

    return {"eigenvalues": eigenvalues, "bare_coulomb": eigenvectors @ np.diag(eigenvalues) @ eigenvectors.T.conj()}


def parse_sgi(file_name: str) -> Dict[str, NDArray[np.complex128]]:
    """Parser for SGI_Q*.OUT, where * is an integer.

    This file contains the vectors $\tilde{S_{Gi}}$, defined in Eq. (41) of
    Computer Phys. Comm. 184, 348 (2013).
    By reading the matrix A as stored, and performing A^H*A, one gets the overlap
    matrix between basis elements defined for the interstitial region.

    :param file_name: name of the file
    :return: parsed data as dictionary
    """
    return {"OverlapMatrix": __square_matrix(__parse_file_with_matrix(file_name))}


def parse_epsilon(file_name: str) -> Dict[str, NDArray[np.complex128]]:
    """Parser for the gw-epsilon files, which contain the dielectric function, its head or wings.:
      - EPSH.OUT,
      - EPSW1.OUT,
      - EPSW2.OUT,
      - EPSILON-GW_Q*.OUT, where * is an integer

    :param file_name: name of the file
    :return: parsed data as dictionary
    """
    return {"epsilon_tensor": __parse_file_with_array_of_rank_3(file_name)}


def parse_polarizability_factor(file_name: str) -> Dict[str, NDArray[np.complex128]]:
    """Parser for the files POLARIZABILITY_FACTOR_Q*.OUT, where * is an integer
    These files contain the polarizability factors needed to obtain the dielectric matrix.

    :param file_name: name of the file
    :return: parsed data as dictionary
    """
    return {"polarizability_factor": __parse_file_with_array_of_rank_4(file_name)}


def parse_inverse_epsilon(file_name: str) -> Dict[str, NDArray[np.complex128]]:
    """Parser for the gw-inverse-epsilon files, which contain the inverse of the
    dielectric function, its head or wings:
      - INVERSE-EPSH.OUT,
      - INVERSE-EPSW1.OUT,
      - INVERSE-EPSW2.OUT,
      - INVERSE-EPSILON_Q*.OUT, where * is an integer

    :param file_name: name of the file
    :return: parsed data as dictionary
    """
    return {"inverse_epsilon_tensor": __parse_file_with_array_of_rank_3(file_name)}


def parse_sigmac(file_name: str) -> Dict[str, NDArray[np.complex128]]:
    """Parser for the files SIGMAC_K*.OUT, where * is an integer
    These files contain the correlation part of the self-energy.

    :param file_name: name of the file
    :return: parsed data as dictionary
    """
    return {"sigmac_matrix": __parse_file_with_matrix(file_name)}


def parse_sigmax(file_name: str) -> Dict[str, NDArray[np.complex128]]:
    """Parser for the files SIGMAX_K*.OUT, where * is an integer
    These files contain the exchange part of the self-energy.

    :param file_name: name of the file
    :return: parsed data as dictionary
    """
    return {"sigmax": __parse_file_with_vector(file_name)}
