# Copyright (C) 2020 Atsushi Togo
# All rights reserved.
#
# This file is part of phono3py.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""SSCHA calculation

Formulae implemented are based on these papers:

    Ref. 1: Bianco et al. https://doi.org/10.1103/PhysRevB.96.014111
    Ref. 2: Aseginolaza et al. https://doi.org/10.1103/PhysRevB.100.214307

"""

import numpy as np
from phonopy.units import VaspToTHz
from phonopy.harmonic.dynmat_to_fc import DynmatToForceConstants
from phono3py.phonon.func import mode_length


class SupercellPhonon(object):
    """Supercell phonon class

    Dynamical matrix is created for supercell atoms and solved in real.
    All phonons at commensurate points are folded to those at Gamma point.
    Phonon eigenvectors are represented in real type.

    Attributes
    ----------
    eigenvalues : ndarray
        Phonon eigenvalues of supercell dynamical matrix.
        shape=(3 * num_satom, ), dtype='double', order='C'
    eigenvectors : ndarray
        Phonon eigenvectors of supercell dynamical matrix.
        shape=(3 * num_satom, 3 * num_satom), dtype='double', order='C'
    frequencies : ndarray
        Phonon frequencies of supercell dynamical matrix. Frequency conversion
        factor to THz is multiplied.
        shape=(3 * num_satom, ), dtype='double', order='C'
    force_constants : ndarray
        Supercell force constants. The array shape is different from
        the input force constants.
        shape=(3 * num_satom, 3 * num_satom), dtype='double', order='C'
    supercell : PhonopyAtoms or its derived class
        Supercell.

    """

    def __init__(self,
                 supercell,
                 force_constants,
                 frequency_factor_to_THz=VaspToTHz):
        """

        Parameters
        ----------
        supercell : PhonopyAtoms
            Supercell.
        force_constants : array_like
            Second order force constants.
            shape=(num_satom, num_satom, 3, 3), dtype='double', order='C'
        frequency_factor_to_THz : float
            Frequency conversion factor to THz.

        """

        self._supercell = supercell
        N = len(supercell)
        _fc2 = np.array(np.transpose(force_constants, axes=[0, 2, 1, 3]),
                        dtype='double', order='C')
        _fc2 = _fc2.reshape((3 * N, 3 * N))
        _fc2 = np.array(_fc2, dtype='double', order='C')
        inv_sqrt_masses = 1.0 / np.repeat(np.sqrt(supercell.masses), 3)
        dynmat = np.array(inv_sqrt_masses * (inv_sqrt_masses * _fc2).T,
                          dtype='double', order='C')
        eigvals, eigvecs = np.linalg.eigh(dynmat)
        freqs = np.sqrt(np.abs(eigvals)) * np.sign(eigvals)
        freqs *= frequency_factor_to_THz
        self._eigenvalues = np.array(eigvals, dtype='double', order='C')
        self._eigenvectors = np.array(eigvecs, dtype='double', order='C')
        self._frequencies = np.array(freqs, dtype='double', order='C')
        self._force_constants = _fc2

    @property
    def eigenvalues(self):
        return self._eigenvalues

    @property
    def eigenvectors(self):
        return self._eigenvectors

    @property
    def frequencies(self):
        return self._frequencies

    @property
    def force_constants(self):
        return self._force_constants

    @property
    def supercell(self):
        return self._supercell


class DispCorrMatrix(object):
    """Calculate displacement correlation matrix <u_a u_b>

    Attributes
    ----------
    psi_matrix : ndarray
        Displacement-displacement correlation matrix at temperature.
        Physical unit is [Angstrom^2].
        shape=(3 * num_satom, 3 * num_satom), dtype='double', order='C'
    upsilon_matrix : ndarray
        Inverse displacement-displacement correlation matrix at temperature.
        Physical unit is [1/Angstrom^2].
        shape=(3 * num_satom, 3 * num_satom), dtype='double', order='C'
    supercell_phonon : SupercellPhonon
        Supercell phonon object. Phonons at Gamma point, where
        eigenvectors are not complex type but real type.

    """

    def __init__(self,