# Copyright 2020 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
r"""
This module contains various utility functions needed to perform quantum chemistry calculations with
Strawberry Fields.
"""
from typing import Tuple

import numpy as np
from scipy.constants import c, h, m_u, pi
from thewalrus import quantum


def duschinsky(
    Li: np.ndarray, Lf: np.ndarray, ri: np.ndarray, rf: np.ndarray, wf: np.ndarray, m: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
    r"""Generate the Duschinsky rotation matrix :math:`U` and displacement vector :math:`\delta`.

    The Duschinsky transformation relates the normal coordinates of the initial and
    final states in a vibronic transition, :math:`q_i` and :math:`q_f` respectively, as:

    .. math:: q_f = U q_i + d,

    where :math:`U` is the Duschinsky rotation matrix and :math:`d` is a vector giving the
    displacement between the equilibrium structures of the two states involved in the vibronic
    transition. The normal coordinates of a molecule can be represented in terms of atomic
    displacements as:

    .. math:: q = L^T \sqrt{m} (r -r_e),

    where :math:`r_e` represents the equilibrium geometry of the molecule, :math:`m` represents
    atomic masses and :math:`L` is a matrix containing the eigenvectors of the mass-weighted
    Hessian. The Duschinsky parameters :math:`U` and :math:`d` can be obtained as:

    .. math:: U = L_f^T L_i,

    .. math:: d = L_f^T \sqrt{m} (r_e^i-r_e^f).

    Note that :math:`i` and :math:`f` refer to the initial and final states, respectively. The
    parameter :math:`d` is usually represented as a dimensionless parameter :math:`\delta` as:

    .. math:: \delta = l^{-1} d,

    where :math:`l` is a diagonal matrix containing the vibrational frequencies :math:`\omega` of
    the final state:

    .. math:: l_{kk} = \left ( \frac{\hbar }{2 \pi \omega_k c} \right )^{1/2},

    where :math:`\hbar` is the reduced Planck constant and :math:`c` is the speed of light.

    The vibrational normal mode matrix for a molecule with :math:`M` vibrational modes and
    :math:`N` atoms is a :math:`3N \times M` matrix where :math:`M = 3N - 6` for nonlinear molecules and
    :math:`M = 3N - 5` for linear molecules. The Duschinsky rotation matrix of a molecule is an
    :math:`M \times M` matrix and the Duschinsky displacement vector has :math:`M` components.

    **Example usage:**

    >>> Li = np.array([[-0.28933191], [0.0], [0.0], [0.95711104], [0.0], [0.0]])
    >>> Lf = np.array([[-0.28933191], [0.0], [0.0], [0.95711104], [0.0], [0.0]])
    >>> ri = np.array([-0.0236, 0.0, 0.0, 1.2236, 0.0, 0.0])
    >>> rf = np.array([0.0, 0.0, 0.0, 1.4397, 0.0, 0.0])
    >>> wf = np.array([1363.2])
    >>> m = np.array([11.0093] * 3 + [1.0078] * 3)
    >>> U, delta = duschinsky(Li, Lf, ri, rf, wf, m)
    >>> U, delta
    (array([[0.99977449]]), array([-1.17623073]))

    Args:
        Li (array): mass-weighted normal modes of the initial electronic state
        Lf (array): mass-weighted normal modes of the final electronic state
        ri (array): equilibrium molecular geometry of the initial electronic state
        rf (array): equilibrium molecular geometry of the final electronic state
        wf (array): normal mode frequencies of the final electronic state in units of :math:`\mbox{cm}^{-1}`
        m (array): atomic masses in unified atomic mass units

    Returns:
        tuple[array, array]: Duschinsky rotation matrix :math:`U`, Duschinsky displacement vector
        :math:`\delta`
    """
    U = Lf.T @ Li

    d = Lf.T * m**0.5 @ (ri - rf)

    l0_inv = np.diag((h / (wf * 100.0 * c)) ** (-0.5) * 2.0 * pi) / 1.0e10 * m_u**0.5

    delta = np.array(d @ l0_inv)

    return U, delta


def read_gamess(
    file, scale_factor: float = 1.0
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    r"""Reads molecular data from the output file generated by the GAMESS quantum chemistry package
    :cite:`schmidt1993general`.

    This function extracts the atomic coordinates (r), atomic masses (m), vibrational frequencies
    (w), and normal modes (n) of a molecule from the output file of a vibrational frequency
    calculation performed with the GAMESS quantum chemistry package. The output file must contain
    the results of a `RUNTYP=HESSIAN` calculation performed with GAMESS. We recommend checking the
    output of this function with the GAMESS results to assure that the GAMESS output file is parsed
    correctly. Please also check if the normal modes are mass-weighted.

    **Example usage:**

    >>> r, m, w, n = read_gamess('../BH_data.out')
    >>> r # atomic coordinates
    array([[0.0000000, 0.0000000, 0.0000000],
           [1.2536039, 0.0000000, 0.0000000]])
    >>> m # atomic masses
    array([11.00931,  1.00782])
    >>> w # vibrational frequencies
    array([19.74, 19.73, 0.00, 0.00, 0.00, 2320.32])
    >>> n # normal modes
    array([[-0.0000000e+00, -7.5322000e-04, -8.7276210e-02,  0.0000000e+00,
         8.2280900e-03,  9.5339055e-01],
       [-0.0000000e+00, -8.7276210e-02,  7.5322000e-04,  0.0000000e+00,
         9.5339055e-01, -8.2280900e-03],
       [ 2.8846925e-01, -2.0000000e-08,  2.0000000e-08,  2.8846925e-01,
        -2.0000000e-08,  2.0000000e-08],
       [ 2.0000000e-08,  2.8846925e-01, -2.0000000e-08,  2.0000000e-08,
         2.8846925e-01, -2.0000000e-08],
       [-2.0000000e-08,  2.0000000e-08,  2.8846925e-01, -2.0000000e-08,
         2.0000000e-08,  2.8846925e-01],
       [-8.7279460e-02,  0.0000000e+00,  0.0000000e+00,  9.5342606e-01,
        -0.0000000e+00, -0.0000000e+00]])

    Args:
        file (str): path to the GAMESS output file
        scale_factor: empirical scale factor for correcting the vibrational frequencies

    Returns:
        tuple[array, array, array, array]: atomic coordinates, atomic masses, normal mode
        frequencies, normal modes
    """
    with open(file, "r", encoding="utf-8") as f:

        r = []
        m = []
        w = []
        n = []

        for line in f:

            if "INPUT CARD> $data" in line or "INPUT CARD> $DATA" in line:
                line = [next(f) for _ in range(3)][-1]
                while "end" not in line and "END" not in line:
                    r.append(np.array(line.rstrip().split()[-3:], float))
                    line = next(f).rstrip()

            if "ATOMIC WEIGHTS" in line:
                next(f)
                for _ in range(len(r)):
                    m.append(np.array(next(f).rstrip().split()[-1:], float))

            if "FREQUENCY" in line:
                line = line.rstrip().split()
                n_mode = len(line) - 1
                w.append(np.array(line[-n_mode:], float))

                while f.readline() != "\n":
                    pass

                d = []
                for _ in range(len(r) * 3):
                    d.append(f.readline().rstrip().split()[-n_mode:])
                n.append(np.array(d, float).T)

    if not r:
        raise ValueError("No atomic coordinates found in the output file")

    if not m:
        raise ValueError("No atomic masses found in the output file")

    if not w:
        raise ValueError("No vibrational frequencies found in the output file")

    return (
        np.concatenate(r).reshape(len(r), 3),
        np.concatenate(m),
        np.concatenate(w) * scale_factor,
        np.concatenate(n),
    )


def prob(samples: list, excited_state: list) -> float:
    r"""Estimate probability of observing an excited state.

    The probability is estimated by calculating the relative frequency of the excited
    state among the samples.

    **Example usage:**

    >>> excited_state = [0, 2]
    >>> samples = [[0, 2], [1, 1], [0, 2], [2, 0], [1, 1], [0, 2], [1, 1], [1, 1], [1, 1]]
    >>> prob(samples, excited_state)
    0.3333333333333333

    Args:
        samples list[list[int]]: a list of samples
        excited_state (list): a Fock state

    Returns:
        float: probability of observing a Fock state in the given samples
    """
    if len(samples) == 0:
        raise ValueError("The samples list must not be empty")

    if len(excited_state) == 0:
        raise ValueError("The excited state list must not be empty")

    if not len(excited_state) == len(samples[0]):
        raise ValueError("The number of modes in the samples and the excited state must be equal")

    if np.any(np.array(excited_state) < 0):
        raise ValueError("The excited state must not contain negative values")

    return samples.count(excited_state) / len(samples)


def marginals(mu: np.ndarray, V: np.ndarray, n_max: int, hbar: float = 2.0) -> np.ndarray:
    r"""Generate single-mode marginal distributions from the displacement vector and covariance
    matrix of a Gaussian state.

    **Example usage:**

    >>> mu = np.array([0.00000000, 2.82842712, 0.00000000,
    ...                0.00000000, 0.00000000, 0.00000000])
    >>> V = np.array([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    ...               [0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
    ...               [0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
    ...               [0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
    ...               [0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
    ...               [0.0, 0.0, 0.0, 0.0, 0.0, 1.0]])
    >>> n_max = 10
    >>> marginals(mu, V, n_max)
    array([[1.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
            0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
            0.00000000e+00, 0.00000000e+00],
           [1.35335284e-01, 2.70670567e-01, 2.70670566e-01, 1.80447044e-01,
            9.02235216e-02, 3.60894085e-02, 1.20298028e-02, 3.43708650e-03,
            8.59271622e-04, 1.90949249e-04],
           [1.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
            0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
            0.00000000e+00, 0.00000000e+00]])

    Args:
        mu (array): displacement vector
        V (array): covariance matrix
        n_max (int): maximum number of vibrational quanta in the distribution
        hbar (float): the value of :math:`\hbar` in the commutation relation :math:`[\x,\p]=i\hbar`.

    Returns:
        array[list[float]]: marginal distributions
    """
    if not V.shape[0] == V.shape[1]:
        raise ValueError("The covariance matrix must be a square matrix")

    if not len(mu) == len(V):
        raise ValueError(
            "The dimension of the displacement vector and the covariance matrix must be equal"
        )

    if n_max <= 0:
        raise ValueError("The number of vibrational states must be larger than zero")

    n_modes = len(mu) // 2

    p = np.zeros((n_modes, n_max))

    for mode in range(n_modes):
        mui, vi = quantum.reduced_gaussian(mu, V, mode)
        for i in range(n_max):
            p[mode, i] = np.real(quantum.density_matrix_element(mui, vi, [i], [i], hbar=hbar))

    return p
