# 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.

import numpy as np


def get_displacements_and_forces_fc3(disp_dataset):
    """Returns displacements and forces from disp_dataset

    Note
    ----
    Dipslacements and forces of all atoms in supercells are returned.

    Parameters
    ----------
    disp_dataset : dict
        Displacement dataset that may be obtained by
        file_IO.parse_disp_fc3_yaml.

    Returns
    -------
    displacements : ndarray
        Displacements of all atoms in all supercells.
        shape=(snapshots, supercell atoms, 3), dtype='double', order='C'
    forces : ndarray or None
        Forces of all atoms in all supercells.
        shape=(snapshots, supercell atoms, 3), dtype='double', order='C'
        None is returned when forces don't exist.

    """

    if 'first_atoms' in disp_dataset:
        natom = disp_dataset['natom']
        ndisp = len(disp_dataset['first_atoms'])
        for disp1 in disp_dataset['first_atoms']:
            ndisp += len(disp1['second_atoms'])
        displacements = np.zeros((ndisp, natom, 3), dtype='double', order='C')
        forces = np.zeros_like(displacements)
        indices = []
        count = 0
        for disp1 in disp_dataset['first_atoms']:
            indices.append(count)
            displacements[count, disp1['number']] = disp1['displacement']
            forces[count] = disp1['forces']
            count += 1

        for disp1 in disp_dataset['first_atoms']:
            for disp2 in disp1['second_atoms']:
                if 'included' in disp2:
                    if disp2['included']:
                        indices.append(count)
                else:
                    indices.append(count)
                displacements[count, disp1['number']] = disp1['displacement']
                displacements[count, disp2['number']] = disp2['displacement']
                forces[count] = disp2['forces']
                count += 1
        return (np.array(displacements[indices], dtype='double', order='C'),
                np.array(forces[indices], dtype='double', order='C'))
    elif 'forces' in disp_dataset and 'displacements' in disp_dataset:
        return disp_dataset['displacements'], disp_dataset['forces']
    else:
        raise RuntimeError("disp_dataset doesn't contain correct information.")
