<commit_before># Copyright (C) 2020 <NAME>
# All rights reserved.
#
# This file is part of phonopy.
#
# 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 sys
import os
import numpy as np
from phonopy import Phonopy, __version__
from phonopy.harmonic.force_constants import (
    distribute_force_constants_by_translations)
from phonopy.structure.cells import print_cell
from phonopy.structure.atoms import atom_data, symbol_map
from phonopy.structure.dataset import forces_in_dataset
from phonopy.interface.phonopy_yaml import PhonopyYaml
from phonopy.interface.fc_calculator import fc_calculator_names
from phonopy.interface.calculator import (
    write_supercells_with_displacements,
    get_default_physical_units, get_default_displacement_distance)
from phonopy.interface.vasp import create_FORCE_CONSTANTS
from phonopy.phonon.band_structure import (
    get_band_qpoints, get_band_qpoints_by_seekpath)
from phonopy.phonon.dos import get_pdos_indices
from phonopy.file_IO import (
    parse_FORCE_CONSTANTS, parse_FORCE_SETS,
    write_FORCE_CONSTANTS, write_force_constants_to_hdf5,
    get_born_parameters, parse_QPOINTS, is_file_phonopy_yaml)
from phonopy.cui.create_force_sets import create_FORCE_SETS
from phonopy.cui.load_helper import (
    read_force_constants_from_hdf5, get_nac_params,
    set_dataset_and_force_constants)
from phonopy.cui.show_symmetry import check_symmetry
from phonopy.cui.settings import PhonopyConfParser
from phonopy.cui.collect_cell_info import collect_cell_info
from phonopy.cui.phonopy_argparse import (
    get_parser, show_deprecated_option_warnings)


# AA is created at http://www.network-science.de/ascii/.
def print_phonopy():
    print(r"""        _
  _ __ | |__   ___  _ __   ___   _ __  _   _
 | '_ \| '_ \ / _ \| '_ \ / _ \ | '_ \| | | |
 | |_) | | | | (_) | | | | (_) || |_) | |_| |
 | .__/|_| |_|\___/|_| |_|\___(_) .__/ \__, |
 |_|                            |_|    |___/""")


def print_version(version, package_name="phonopy"):
    try:
        version_text = ('%s' % version).rjust(44)
        import pkg_resources
        dist = pkg_resources.get_distribution(package_name)
        if dist.has_version():
            ver = dist.version.split('.')
            if len(ver) > 3:
                rev = ver[3]
                version_text = ('%s-r%s' % (version, rev)).rjust(44)
    except ImportError:
        pass
    except Exception as err:
        if (err.__module__ == 'pkg_resources' and
            err.__class__.__name__ == 'DistributionNotFound'):
            pass
        else:
            raise
    finally:
        print(version_text)
        print('')


def print_end():
    print(r"""                 _
   ___ _ __   __| |
  / _ \ '_ \ / _` |
 |  __/ | | | (_| |
  \___|_| |_|\__,_|
""")


def print_error():
    print(r"""  ___ _ __ _ __ ___  _ __
 / _ \ '__| '__/ _ \| '__|
|  __/ |  | | | (_) | |
 \___|_|  |_|  \___/|_|
""")


def print_attention(attention_text):
    print("*" * 67)
    print(attention_text)
    print("*" * 67)
    print('')


def print_error_message(message):
    print('')
    print(message)


def file_exists(filename, log_level, is_any=False):
    if os.path.exists(filename):
        return True
    else:
        if is_any:
            return False
        else:
            error_text = "\"%s\" was not found." % filename
            print_error_message(error_text)
            if log_level > 0:
                print_error()
            sys.exit(1)


def files_exist(filename_list, log_level, is_any=False):
    filenames = []
    for filename in filename_list:
        if file_exists(filename, log_level, is_any=is_any):
            filenames.append(filename)

    if filenames:
        return filenames
    else:
        if len(filenames) == 2:
            all_filenames = "\"%s\" or \"%s\"" % tuple(filenames)
        else:
            all_filenames = ", ".join(["\"%s\"" %
                                       fn for fn in filename_list[:-1]])
            all_filenames += " or \"%s\"" % filename_list[-1]
        error_text = "Any of %s was not found." % all_filenames
        print_error_message(error_text)
        if log_level > 0:
            print_error()
        sys.exit(1)


def finalize_phonopy(log_level,
                     settings,
                     confs,
                     phonon,
                     filename="phonopy.yaml"):
    units = get_default_physical_units(phonon.calculator)

    if settings.save_params:
        exists_fc_only = (not forces_in_dataset(phonon.dataset) and
                          phonon.force_constants is not None)
        yaml_settings = {
            'displacements': not exists_fc_only,
            'force_sets': not exists_fc_only,
            'force_constants': exists_fc_only,
            'born_effective_charge': True,
            'dielectric_constant': True
        }
        _filename = "phonopy_params.yaml"
    else:
        yaml_settings = {
            'force_sets': settings.include_force_sets,
            'force_constants': settings.include_force_constants,
            'born_effective_charge': settings.include_nac_params,
            'dielectric_constant': settings.include_nac_params,
            'displacements': settings.include_displacements
        }
        _filename = filename

    phpy_yaml = PhonopyYaml(configuration=confs,
                            physical_units=units,
                            settings=yaml_settings)
    phpy_yaml.set_phonon_info(phonon)
    with open(_filename, 'w') as w:
        w.write(str(phpy_yaml))

    if log_level > 0:
        print("")
        if settings.save_params:
            print("Summary of calculation and parameters were written "
                  "in \"%s\"." % _filename)
        else:
            print("Summary of calculation was written in \"%s\"." % _filename)
        print_end()
    sys.exit(0)


def print_cells(phonon, unitcell_filename):
    supercell = phonon.get_supercell()
    unitcell = phonon.get_unitcell()
    primitive = phonon.get_primitive()
    p2p_map = primitive.get_primitive_to_primitive_map()
    mapping = np.array(
        [p2p_map[x] for x in primitive.get_supercell_to_primitive_map()],
        dtype='intc')
    s_indep_atoms = phonon.get_symmetry().get_independent_atoms()
    p_indep_atoms = mapping[s_indep_atoms]
    u2s_map = supercell.get_unitcell_to_supercell_map()
    print("-" * 30 + " primitive cell " + "-" * 30)
    print_cell(primitive, stars=p_indep_atoms)
    print("-" * 32 + " unit cell " + "-" * 33)  # 32 + 11 + 33 = 76
    u2u_map = supercell.get_unitcell_to_unitcell_map()
    u_indep_atoms = [u2u_map[x] for x in s_indep_atoms]
    print_cell(unitcell, mapping=mapping[u2s_map], stars=u_indep_atoms)
    print("-" * 32 + " super cell " + "-" * 32)
    print_cell(supercell, mapping=mapping, stars=s_indep_atoms)
    print("-" * 76)


def print_settings(settings,
                   phonon,
                   is_primitive_axes_auto,
                   unitcell_filename,
                   load_phonopy_yaml):
    primitive_matrix = phonon.primitive_matrix
    supercell_matrix = phonon.supercell_matrix
    interface_mode = phonon.calculator
    run_mode = settings.run_mode
    if interface_mode:
        print("Calculator interface: %s" % interface_mode)
    if (settings.cell_filename is not None and
        settings.cell_filename != unitcell_filename):
        print("\"%s\" was not able to be used."
              % settings.cell_filename)
    print("Crystal structure was read from \"%s\"." % unitcell_filename)
    physical_units = get_default_physical_units(interface_mode)
    print("Unit of length: %s" % physical_units['length_unit'])
    if is_band_auto(settings) and not is_primitive_axes_auto:
        print("Automatic band structure mode forced automatic choice "
              "of primitive axes.")
    if run_mode == 'band':
        if is_band_auto(settings):
            print("Band structure mode (Auto)")
        else:
            print("Band structure mode")
    if run_mode == 'mesh':
        print("Mesh sampling mode")
    if run_mode == 'band_mesh':
        print("Band structure and mesh sampling mode")
    if run_mode == 'anime':
        print("Animation mode")
    if run_mode == 'modulation':
        print("Modulation mode")
    if run_mode == 'irreps':
        print("Ir-representation mode")
    if run_mode == 'qpoints':
        if settings.write_dynamical_matrices:
            print("QPOINTS mode (dynamical matrices written out)")
        else:
            print("QPOINTS mode")
    if ((run_mode == 'band' or run_mode == 'mesh' or run_mode == 'qpoints') and
        settings.is_group_velocity):
        gv_delta_q = settings.group_velocity_delta_q
        if gv_delta_q is not None:
            print("  With group velocity calculation (dq=%3.1e)" % gv_delta_q)
        else:
            print('')
    if settings.create_displacements:
        print("Displacements creation mode")
        if not settings.is_plusminus_displacement == 'auto':
            if settings.is_plusminus_displacement:
                print("  Plus Minus displacement: full plus minus directions")
            else:
                print("  Plus Minus displacement: only one direction")
        if not settings.is_diagonal_displacement:
            print("  Diagonal displacement: off")
        if settings.random_displacements:
            print("  Number of supercells with random displacements: %d"
                  % settings.random_displacements)
            if settings.random_seed is not None:
                print("  Random seed: %d" % settings.random_seed)

    print("Settings:")
    if load_phonopy_yaml:
        if not settings.is_nac:
            print("  Non-analytical term correction (NAC): off")
    else:
        if settings.is_nac:
            print("  Non-analytical term correction (NAC): on")
            if settings.nac_q_direction is not None:
                print("  NAC q-direction: %s" % settings.nac_q_direction)
    if settings.fc_spg_symmetry:
        print("  Enforce space group symmetry to force constants: on")
    if load_phonopy_yaml:
        if not settings.fc_symmetry:
            print("  Force constants symmetrization: off")
    else:
        if settings.fc_symmetry:
            print("  Force constants symmetrization: on")
    if settings.lapack_solver:
        print("  Use Lapack solver via Lapacke: on")
    if settings.symmetry_tolerance is not None:
        print("  Symmetry tolerance: %5.2e"
              % settings.symmetry_tolerance)
    if run_mode == 'mesh' or run_mode == 'band_mesh':
        mesh = settings.mesh_numbers
        if type(mesh) is float:
            print("  Length for sampling mesh: %.1f" % mesh)
        elif mesh is not None:
            print("  Sampling mesh: %s" % np.array(mesh))
        if settings.is_thermal_properties:
            cutoff_freq = settings.cutoff_frequency
            if cutoff_freq is None:
                pass
            elif cutoff_freq < 0:
                print("  - Thermal properties are calculatd with "
                      "absolute phonon frequnecy.")
            else:
                print("  - Phonon frequencies > %5.3f are used to calculate "
                      "thermal properties." % cutoff_freq)
        elif (settings.is_thermal_displacements or
              settings.is_thermal_displacement_matrices):
            fmin = settings.min_frequency
            fmax = settings.max_frequency
            text = None
            if (fmin is not None) and (fmax is not None):
                text = "  - Phonon frequency from %5.3f to %5.3f" % (fmin,
                                                                     fmax)
                text += " are used to calculate\n"
                text += "    thermal displacements."
            elif (fmin is None) and (fmax is not None):
                text = "Phonon frequency < %5.3f" % fmax
                text = ("  - %s are used to calculate thermal displacements." %
                        text)
            elif (fmin is not None) and (fmax is None):
                text = "Phonon frequency > %5.3f" % fmin
                text = ("  - %s are used to calculate thermal displacements." %
                        text)
            if text:
                print(text)
    if (np.diag(np.diag(supercell_matrix)) - supercell_matrix).any():
        print("  Supercell matrix:")
        for v in supercell_matrix:
            print("    %s" % v)
    else:
        print("  Supercell: %s" % np.diag(supercell_matrix))
    if is_primitive_axes_auto or is_band_auto(settings):
        print("  Primitive matrix (Auto):")
        for v in primitive_matrix:
            print("    %s" % v)
    elif primitive_matrix is not None:
        print("  Primitive matrix:")
        for v in primitive_matrix:
            print("    %s" % v)


def write_displacements_files_then_exit(phonon,
                                        settings,
                                        confs,
                                        optional_structure_info,
                                        log_level):
    """Write supercells with displacements and displacement dataset

    Note
    ----
    From phonopy v1.15.0, displacement dataset is written into
    phonopy_disp.yaml.

    """

    cells_with_disps = phonon.supercells_with_displacements
    additional_info = {'supercell_matrix': phonon.supercell_matrix}
    write_supercells_with_displacements(phonon.calculator,
                                        phonon.supercell,
                                        cells_with_disps,
                                        optional_structure_info,
                                        additional_info=additional_info)

    if log_level > 0:
        print("\"phonopy_disp.yaml\" and supercells have been created.")

    settings.set_include_displacements(True)

    finalize_phonopy(log_level,
                     settings,
                     confs,
                     phonon,
                     filename="phonopy_disp.yaml")


def create_FORCE_SETS_from_settings(settings, symprec, log_level):
    if settings.create_force_sets:
        filenames = settings.create_force_sets
        force_sets_zero_mode = False
    elif settings.create_force_sets_zero:
        filenames = settings.create_force_sets_zero
        force_sets_zero_mode = True
    else:
        print_error_message("Something wrong for parsing arguments.")
        sys.exit(0)

    disp_filenames = files_exist(
        ['phonopy_disp.yaml', 'disp.yaml'], log_level, is_any=True)
    interface_mode = settings.calculator

    if disp_filenames[0] == 'phonopy_disp.yaml':
        try:
            phpy_yaml = PhonopyYaml()
            phpy_yaml.read('phonopy_disp.yaml')
            if phpy_yaml.calculator is not None:
                interface_mode = phpy_yaml.calculator  # overwrite
            disp_filename = 'phonopy_disp.yaml'
        except KeyError:
            file_exists('disp.yaml', log_level)
            if log_level > 0:
                print("\"phonopy_disp.yaml\" was found but wasn't used "
                      "because of the old-style format.")
            disp_filename = 'disp.yaml'
    else:
        disp_filename = disp_filenames[0]

    files_exist(filenames, log_level)

    create_FORCE_SETS(
        interface_mode,
        filenames,
        symmetry_tolerance=symprec,
        force_sets_zero_mode=force_sets_zero_mode,
        disp_filename=disp_filename,
        log_level=log_level)


def produce_force_constants(phonon,
                            settings,
                            phpy_yaml,
                            unitcell_filename,
                            log_level):
    num_satom = len(phonon.supercell)
    p2s_map = phonon.primitive.p2s_map
    is_full_fc = settings.fc_spg_symmetry or settings.is_full_fc

    if settings.read_force_constants:
        if settings.is_hdf5 or settings.readfc_format == 'hdf5':
            try:
                import h5py
            except ImportError:
                print_error_message("You need to install python-h5py.")
                if log_level:
                    print_error()
                sys.exit(1)

            file_exists("force_constants.hdf5", log_level)
            fc = read_force_constants_from_hdf5(
                filename="force_constants.hdf5",
                p2s_map=p2s_map,
                calculator=phonon.calculator)
            fc_filename = "force_constants.hdf5"
        else:
            file_exists("FORCE_CONSTANTS", log_level)
            fc = parse_FORCE_CONSTANTS(filename="FORCE_CONSTANTS",
                                       p2s_map=p2s_map)
            fc_filename = "FORCE_CONSTANTS"

        if log_level:
            print("Force constants are read from \"%s\"." % fc_filename)

        if fc.shape[1] != num_satom:
            error_text = ("Number of atoms in supercell is not consistent "
                          "with the matrix shape of\nforce constants read "
                          "from ")
            if settings.is_hdf5 or settings.readfc_format == 'hdf5':
                error_text += "force_constants.hdf5.\n"
            else:
                error_text += "FORCE_CONSTANTS.\n"
            error_text += ("Please carefully check DIM, FORCE_CONSTANTS, "
                           "and %s.") % unitcell_filename
            print_error_message(error_text)
            if log_level:
                print_error()
            sys.exit(1)

        # Compact fc is expanded to full fc when full fc is required.
        if is_full_fc and fc.shape[0] != fc.shape[1]:
            fc = compact_fc_to_full_fc(phonon, fc, log_level=log_level)

        phonon.force_constants = fc
    else:
        def read_force_sets_from_phonopy_yaml(phpy_yaml):
            if (phpy_yaml.dataset is not None and
                ('forces' in phpy_yaml.dataset or
                 ('first_atoms' in phpy_yaml.dataset and
                  'forces' in phpy_yaml.dataset['first_atoms'][0]))):
                return phpy_yaml.dataset
            else:
                return None

        force_sets = None

        if phpy_yaml is not None:
            force_sets = read_force_sets_from_phonopy_yaml(phpy_yaml)
            if log_level:
                if force_sets is None:
                    print("Force sets were not found in \"%s\"."
                          % unitcell_filename)
                else:
                    print("Forces and displacements were read from \"%s\"."
                          % unitcell_filename)

        if force_sets is None:
            file_exists("FORCE_SETS", log_level)
            force_sets = parse_FORCE_SETS(natom=num_satom)
            if log_level:
                print("Forces and displacements were read from \"%s\"."
                      % "FORCE_SETS")

        if (log_level and
            force_sets is not None and
            'displacements' in force_sets):
            print("%d snapshots were found."
                  % len(force_sets['displacements']))

        if 'natom' in force_sets:
            natom = force_sets['natom']
        else:
            natom = force_sets['forces'].shape[1]
        if natom != num_satom:
            error_text = "Number of atoms in supercell is not consistent with "
            error_text += "the data in FORCE_SETS.\n"
            error_text += ("Please carefully check DIM, FORCE_SETS,"
                           " and %s") % unitcell_filename
            print_error_message(error_text)
            if log_level:
                print_error()
            sys.exit(1)

        (fc_calculator,
         fc_calculator_options) = get_fc_calculator_params(settings)

        phonon.dataset = force_sets
        if log_level:
            if fc_calculator is not None:
                print("Force constants calculation by %s starts."
                      % fc_calculator_names[fc_calculator])
            else:
                print("Computing force constants...")

        if is_full_fc:
            # Need to calculate full force constant tensors
            phonon.produce_force_constants(
                fc_calculator=fc_calculator,
                fc_calculator_options=fc_calculator_options)
        else:
            # Only force constants between atoms in primitive cell and
            # supercell
            phonon.produce_force_constants(
                calculate_full_force_constants=False,
                fc_calculator=fc_calculator,
                fc_calculator_options=fc_calculator_options)


def compact_fc_to_full_fc(phonon, compact_fc, log_level=0):
    fc = np.zeros((compact_fc.shape[1], compact_fc.shape[1], 3, 3),
                  dtype='double', order='C')
    fc[phonon.primitive.p2s_map] = compact_fc
    distribute_force_constants_by_translations(
        fc, phonon.primitive, phonon.supercell)
    if log_level:
        print("Force constants were expanded to full format.")

    return fc


def store_force_constants(phonon,
                          settings,
                          phpy_yaml,
                          unitcell_filename,
                          load_phonopy_yaml,
                          log_level):
    physical_units = get_default_physical_units(phonon.calculator)
    p2s_map = phonon.primitive.p2s_map

    if load_phonopy_yaml:
        is_full_fc = settings.fc_spg_symmetry or settings.is_full_fc
        if phpy_yaml.force_constants is not None:
            if log_level:
                print("Force constants were read from \"%s\"."
                      % unitcell_filename)

            fc = phpy_yaml.force_constants

            if fc.shape[1] != len(phonon.supercell):
                error_text = ("Number of atoms in supercell is not consistent "
                              "with the matrix shape of\nforce constants read "
                              "from %s." % unitcell_filename)
                print_error_message(error_text)
                if log_level:
                    print_error()
                sys.exit(1)

            # Compact fc is expanded to full fc when full fc is required.
            if is_full_fc and fc.shape[0] != fc.shape[1]:
                fc = compact_fc_to_full_fc(phonon, fc, log_level=log_level)

            phonon.set_force_constants(fc, show_drift=(log_level > 0))

        if phpy_yaml.dataset is not None:
            if log_level:
                text = "Force sets were read from \"%s\"" % unitcell_filename
                if phpy_yaml.force_constants is not None:
                    text += " but not to be used."
                else:
                    text += "."
                print(text)

        if phpy_yaml.force_constants is None:
            (fc_calculator,
             fc_calculator_options) = get_fc_calculator_params(settings)
            set_dataset_and_force_constants(
                phonon,
                phpy_yaml.dataset,
                None,
                fc_calculator=fc_calculator,<commit_msg>fix of phonopy-load command to accept fc-calc-opt<commit_after># Copyright (C) 2020 <NAME>
# All rights reserved.
#
# This file is part of phonopy.
#
# 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 sys
import os
import numpy as np
from phonopy import Phonopy, __version__
from phonopy.harmonic.force_constants import (
    distribute_force_constants_by_translations)
from phonopy.structure.cells import print_cell
from phonopy.structure.atoms import atom_data, symbol_map
from phonopy.structure.dataset import forces_in_dataset
from phonopy.interface.phonopy_yaml import PhonopyYaml
from phonopy.interface.fc_calculator import fc_calculator_names
from phonopy.interface.calculator import (
    write_supercells_with_displacements,
    get_default_physical_units, get_default_displacement_distance)
from phonopy.interface.vasp import create_FORCE_CONSTANTS
from phonopy.phonon.band_structure import (
    get_band_qpoints, get_band_qpoints_by_seekpath)
from phonopy.phonon.dos import get_pdos_indices
from phonopy.file_IO import (
    parse_FORCE_CONSTANTS, parse_FORCE_SETS,
    write_FORCE_CONSTANTS, write_force_constants_to_hdf5,
    get_born_parameters, parse_QPOINTS, is_file_phonopy_yaml)
from phonopy.cui.create_force_sets import create_FORCE_SETS
from phonopy.cui.load_helper import (
    read_force_constants_from_hdf5, get_nac_params,
    set_dataset_and_force_constants)
from phonopy.cui.show_symmetry import check_symmetry
from phonopy.cui.settings import PhonopyConfParser
from phonopy.cui.collect_cell_info import collect_cell_info
from phonopy.cui.phonopy_argparse import (
    get_parser, show_deprecated_option_warnings)


# AA is created at http://www.network-science.de/ascii/.
def print_phonopy():
    print(r"""        _
  _ __ | |__   ___  _ __   ___   _ __  _   _
 | '_ \| '_ \ / _ \| '_ \ / _ \ | '_ \| | | |
 | |_) | | | | (_) | | | | (_) || |_) | |_| |
 | .__/|_| |_|\___/|_| |_|\___(_) .__/ \__, |
 |_|                            |_|    |___/""")


def print_version(version, package_name="phonopy"):
    try:
        version_text = ('%s' % version).rjust(44)
        import pkg_resources
        dist = pkg_resources.get_distribution(package_name)
        if dist.has_version():
            ver = dist.version.split('.')
            if len(ver) > 3:
                rev = ver[3]
                version_text = ('%s-r%s' % (version, rev)).rjust(44)
    except ImportError:
        pass
    except Exception as err:
        if (err.__module__ == 'pkg_resources' and
            err.__class__.__name__ == 'DistributionNotFound'):
            pass
        else:
            raise
    finally:
        print(version_text)
        print('')


def print_end():
    print(r"""                 _
   ___ _ __   __| |
  / _ \ '_ \ / _` |
 |  __/ | | | (_| |
  \___|_| |_|\__,_|
""")


def print_error():
    print(r"""  ___ _ __ _ __ ___  _ __
 / _ \ '__| '__/ _ \| '__|
|  __/ |  | | | (_) | |
 \___|_|  |_|  \___/|_|
""")


def print_attention(attention_text):
    print("*" * 67)
    print(attention_text)
    print("*" * 67)
    print('')


def print_error_message(message):
    print('')
    print(message)


def file_exists(filename, log_level, is_any=False):
    if os.path.exists(filename):
        return True
    else:
        if is_any:
            return False
        else:
            error_text = "\"%s\" was not found." % filename
            print_error_message(error_text)
            if log_level > 0:
                print_error()
            sys.exit(1)


def files_exist(filename_list, log_level, is_any=False):
    filenames = []
    for filename in filename_list:
        if file_exists(filename, log_level, is_any=is_any):
            filenames.append(filename)

    if filenames:
        return filenames
    else:
        if len(filenames) == 2:
            all_filenames = "\"%s\" or \"%s\"" % tuple(filenames)
        else:
            all_filenames = ", ".join(["\"%s\"" %
                                       fn for fn in filename_list[:-1]])
            all_filenames += " or \"%s\"" % filename_list[-1]
        error_text = "Any of %s was not found." % all_filenames
        print_error_message(error_text)
        if log_level > 0:
            print_error()
        sys.exit(1)


def finalize_phonopy(log_level,
                     settings,
                     confs,
                     phonon,
                     filename="phonopy.yaml"):
    units = get_default_physical_units(phonon.calculator)

    if settings.save_params:
        exists_fc_only = (not forces_in_dataset(phonon.dataset) and
                          phonon.force_constants is not None)
        yaml_settings = {
            'displacements': not exists_fc_only,
            'force_sets': not exists_fc_only,
            'force_constants': exists_fc_only,
            'born_effective_charge': True,
            'dielectric_constant': True
        }
        _filename = "phonopy_params.yaml"
    else:
        yaml_settings = {
            'force_sets': settings.include_force_sets,
            'force_constants': settings.include_force_constants,
            'born_effective_charge': settings.include_nac_params,
            'dielectric_constant': settings.include_nac_params,
            'displacements': settings.include_displacements
        }
        _filename = filename

    phpy_yaml = PhonopyYaml(configuration=confs,
                            physical_units=units,
                            settings=yaml_settings)
    phpy_yaml.set_phonon_info(phonon)
    with open(_filename, 'w') as w:
        w.write(str(phpy_yaml))

    if log_level > 0:
        print("")
        if settings.save_params:
            print("Summary of calculation and parameters were written "
                  "in \"%s\"." % _filename)
        else:
            print("Summary of calculation was written in \"%s\"." % _filename)
        print_end()
    sys.exit(0)


def print_cells(phonon, unitcell_filename):
    supercell = phonon.get_supercell()
    unitcell = phonon.get_unitcell()
    primitive = phonon.get_primitive()
    p2p_map = primitive.get_primitive_to_primitive_map()
    mapping = np.array(
        [p2p_map[x] for x in primitive.get_supercell_to_primitive_map()],
        dtype='intc')
    s_indep_atoms = phonon.get_symmetry().get_independent_atoms()
    p_indep_atoms = mapping[s_indep_atoms]
    u2s_map = supercell.get_unitcell_to_supercell_map()
    print("-" * 30 + " primitive cell " + "-" * 30)
    print_cell(primitive, stars=p_indep_atoms)
    print("-" * 32 + " unit cell " + "-" * 33)  # 32 + 11 + 33 = 76
    u2u_map = supercell.get_unitcell_to_unitcell_map()
    u_indep_atoms = [u2u_map[x] for x in s_indep_atoms]
    print_cell(unitcell, mapping=mapping[u2s_map], stars=u_indep_atoms)
    print("-" * 32 + " super cell " + "-" * 32)
    print_cell(supercell, mapping=mapping, stars=s_indep_atoms)
    print("-" * 76)


def print_settings(settings,
                   phonon,
                   is_primitive_axes_auto,
                   unitcell_filename,
                   load_phonopy_yaml):
    primitive_matrix = phonon.primitive_matrix
    supercell_matrix = phonon.supercell_matrix
    interface_mode = phonon.calculator
    run_mode = settings.run_mode
    if interface_mode:
        print("Calculator interface: %s" % interface_mode)
    if (settings.cell_filename is not None and
        settings.cell_filename != unitcell_filename):
        print("\"%s\" was not able to be used."
              % settings.cell_filename)
    print("Crystal structure was read from \"%s\"." % unitcell_filename)
    physical_units = get_default_physical_units(interface_mode)
    print("Unit of length: %s" % physical_units['length_unit'])
    if is_band_auto(settings) and not is_primitive_axes_auto:
        print("Automatic band structure mode forced automatic choice "
              "of primitive axes.")
    if run_mode == 'band':
        if is_band_auto(settings):
            print("Band structure mode (Auto)")
        else:
            print("Band structure mode")
    if run_mode == 'mesh':
        print("Mesh sampling mode")
    if run_mode == 'band_mesh':
        print("Band structure and mesh sampling mode")
    if run_mode == 'anime':
        print("Animation mode")
    if run_mode == 'modulation':
        print("Modulation mode")
    if run_mode == 'irreps':
        print("Ir-representation mode")
    if run_mode == 'qpoints':
        if settings.write_dynamical_matrices:
            print("QPOINTS mode (dynamical matrices written out)")
        else:
            print("QPOINTS mode")
    if ((run_mode == 'band' or run_mode == 'mesh' or run_mode == 'qpoints') and
        settings.is_group_velocity):
        gv_delta_q = settings.group_velocity_delta_q
        if gv_delta_q is not None:
            print("  With group velocity calculation (dq=%3.1e)" % gv_delta_q)
        else:
            print('')
    if settings.create_displacements:
        print("Displacements creation mode")
        if not settings.is_plusminus_displacement == 'auto':
            if settings.is_plusminus_displacement:
                print("  Plus Minus displacement: full plus minus directions")
            else:
                print("  Plus Minus displacement: only one direction")
        if not settings.is_diagonal_displacement:
            print("  Diagonal displacement: off")
        if settings.random_displacements:
            print("  Number of supercells with random displacements: %d"
                  % settings.random_displacements)
            if settings.random_seed is not None:
                print("  Random seed: %d" % settings.random_seed)

    print("Settings:")
    if load_phonopy_yaml:
        if not settings.is_nac:
            print("  Non-analytical term correction (NAC): off")
    else:
        if settings.is_nac:
            print("  Non-analytical term correction (NAC): on")
            if settings.nac_q_direction is not None:
                print("  NAC q-direction: %s" % settings.nac_q_direction)
    if settings.fc_spg_symmetry:
        print("  Enforce space group symmetry to force constants: on")
    if load_phonopy_yaml:
        if not settings.fc_symmetry:
            print("  Force constants symmetrization: off")
    else:
        if settings.fc_symmetry:
            print("  Force constants symmetrization: on")
    if settings.lapack_solver:
        print("  Use Lapack solver via Lapacke: on")
    if settings.symmetry_tolerance is not None:
        print("  Symmetry tolerance: %5.2e"
              % settings.symmetry_tolerance)
    if run_mode == 'mesh' or run_mode == 'band_mesh':
        mesh = settings.mesh_numbers
        if type(mesh) is float:
            print("  Length for sampling mesh: %.1f" % mesh)
        elif mesh is not None:
            print("  Sampling mesh: %s" % np.array(mesh))
        if settings.is_thermal_properties:
            cutoff_freq = settings.cutoff_frequency
            if cutoff_freq is None:
                pass
            elif cutoff_freq < 0:
                print("  - Thermal properties are calculatd with "
                      "absolute phonon frequnecy.")
            else:
                print("  - Phonon frequencies > %5.3f are used to calculate "
                      "thermal properties." % cutoff_freq)
        elif (settings.is_thermal_displacements or
              settings.is_thermal_displacement_matrices):
            fmin = settings.min_frequency
            fmax = settings.max_frequency
            text = None
            if (fmin is not None) and (fmax is not None):
                text = "  - Phonon frequency from %5.3f to %5.3f" % (fmin,
                                                                     fmax)
                text += " are used to calculate\n"
                text += "    thermal displacements."
            elif (fmin is None) and (fmax is not None):
                text = "Phonon frequency < %5.3f" % fmax
                text = ("  - %s are used to calculate thermal displacements." %
                        text)
            elif (fmin is not None) and (fmax is None):
                text = "Phonon frequency > %5.3f" % fmin
                text = ("  - %s are used to calculate thermal displacements." %
                        text)
            if text:
                print(text)
    if (np.diag(np.diag(supercell_matrix)) - supercell_matrix).any():
        print("  Supercell matrix:")
        for v in supercell_matrix:
            print("    %s" % v)
    else:
        print("  Supercell: %s" % np.diag(supercell_matrix))
    if is_primitive_axes_auto or is_band_auto(settings):
        print("  Primitive matrix (Auto):")
        for v in primitive_matrix:
            print("    %s" % v)
    elif primitive_matrix is not None:
        print("  Primitive matrix:")
        for v in primitive_matrix:
            print("    %s" % v)


def write_displacements_files_then_exit(phonon,
                                        settings,
                                        confs,
                                        optional_structure_info,
                                        log_level):
    """Write supercells with displacements and displacement dataset

    Note
    ----
    From phonopy v1.15.0, displacement dataset is written into
    phonopy_disp.yaml.

    """

    cells_with_disps = phonon.supercells_with_displacements
    additional_info = {'supercell_matrix': phonon.supercell_matrix}
    write_supercells_with_displacements(phonon.calculator,
                                        phonon.supercell,
                                        cells_with_disps,
                                        optional_structure_info,
                                        additional_info=additional_info)

    if log_level > 0:
        print("\"phonopy_disp.yaml\" and supercells have been created.")

    settings.set_include_displacements(True)

    finalize_phonopy(log_level,
                     settings,
                     confs,
                     phonon,
                     filename="phonopy_disp.yaml")


def create_FORCE_SETS_from_settings(settings, symprec, log_level):
    if settings.create_force_sets:
        filenames = settings.create_force_sets
        force_sets_zero_mode = False
    elif settings.create_force_sets_zero:
        filenames = settings.create_force_sets_zero
        force_sets_zero_mode = True
    else:
        print_error_message("Something wrong for parsing arguments.")
        sys.exit(0)

    disp_filenames = files_exist(
        ['phonopy_disp.yaml', 'disp.yaml'], log_level, is_any=True)
    interface_mode = settings.calculator

    if disp_filenames[0] == 'phonopy_disp.yaml':
        try:
            phpy_yaml = PhonopyYaml()
            phpy_yaml.read('phonopy_disp.yaml')
            if phpy_yaml.calculator is not None:
                interface_mode = phpy_yaml.calculator  # overwrite
            disp_filename = 'phonopy_disp.yaml'
        except KeyError:
            file_exists('disp.yaml', log_level)
            if log_level > 0:
                print("\"phonopy_disp.yaml\" was found but wasn't used "
                      "because of the old-style format.")
            disp_filename = 'disp.yaml'
    else:
        disp_filename = disp_filenames[0]

    files_exist(filenames, log_level)

    create_FORCE_SETS(
        interface_mode,
        filenames,
        symmetry_tolerance=symprec,
        force_sets_zero_mode=force_sets_zero_mode,
        disp_filename=disp_filename,
        log_level=log_level)


def produce_force_constants(phonon,
                            settings,
                            phpy_yaml,
                            unitcell_filename,
                            log_level):
    num_satom = len(phonon.supercell)
    p2s_map = phonon.primitive.p2s_map
    is_full_fc = settings.fc_spg_symmetry or settings.is_full_fc

    if settings.read_force_constants:
        if settings.is_hdf5 or settings.readfc_format == 'hdf5':
            try:
                import h5py
            except ImportError:
                print_error_message("You need to install python-h5py.")
                if log_level:
                    print_error()
                sys.exit(1)

            file_exists("force_constants.hdf5", log_level)
            fc = read_force_constants_from_hdf5(
                filename="force_constants.hdf5",
                p2s_map=p2s_map,
                calculator=phonon.calculator)
            fc_filename = "force_constants.hdf5"
        else:
            file_exists("FORCE_CONSTANTS", log_level)
            fc = parse_FORCE_CONSTANTS(filename="FORCE_CONSTANTS",
                                       p2s_map=p2s_map)
            fc_filename = "FORCE_CONSTANTS"

        if log_level:
            print("Force constants are read from \"%s\"." % fc_filename)

        if fc.shape[1] != num_satom:
            error_text = ("Number of atoms in supercell is not consistent "
                          "with the matrix shape of\nforce constants read "
                          "from ")
            if settings.is_hdf5 or settings.readfc_format == 'hdf5':
                error_text += "force_constants.hdf5.\n"
            else:
                error_text += "FORCE_CONSTANTS.\n"
            error_text += ("Please carefully check DIM, FORCE_CONSTANTS, "
                           "and %s.") % unitcell_filename
            print_error_message(error_text)
            if log_level:
                print_error()
            sys.exit(1)

        # Compact fc is expanded to full fc when full fc is required.
        if is_full_fc and fc.shape[0] != fc.shape[1]:
            fc = compact_fc_to_full_fc(phonon, fc, log_level=log_level)

        phonon.force_constants = fc
    else:
        def read_force_sets_from_phonopy_yaml(phpy_yaml):
            if (phpy_yaml.dataset is not None and
                ('forces' in phpy_yaml.dataset or
                 ('first_atoms' in phpy_yaml.dataset and
                  'forces' in phpy_yaml.dataset['first_atoms'][0]))):
                return phpy_yaml.dataset
            else:
                return None

        force_sets = None

        if phpy_yaml is not None:
            force_sets = read_force_sets_from_phonopy_yaml(phpy_yaml)
            if log_level:
                if force_sets is None:
                    print("Force sets were not found in \"%s\"."
                          % unitcell_filename)
                else:
                    print("Forces and displacements were read from \"%s\"."
                          % unitcell_filename)

        if force_sets is None:
            file_exists("FORCE_SETS", log_level)
            force_sets = parse_FORCE_SETS(natom=num_satom)
            if log_level:
                print("Forces and displacements were read from \"%s\"."
                      % "FORCE_SETS")

        if (log_level and
            force_sets is not None and
            'displacements' in force_sets):
            print("%d snapshots were found."
                  % len(force_sets['displacements']))

        if 'natom' in force_sets:
            natom = force_sets['natom']
        else:
            natom = force_sets['forces'].shape[1]
        if natom != num_satom:
            error_text = "Number of atoms in supercell is not consistent with "
            error_text += "the data in FORCE_SETS.\n"
            error_text += ("Please carefully check DIM, FORCE_SETS,"
                           " and %s") % unitcell_filename
            print_error_message(error_text)
            if log_level:
                print_error()
            sys.exit(1)

        (fc_calculator,
         fc_calculator_options) = get_fc_calculator_params(settings)

        phonon.dataset = force_sets
        if log_level:
            if fc_calculator is not None:
                print("Force constants calculation by %s starts."
                      % fc_calculator_names[fc_calculator])
            else:
                print("Computing force constants...")

        if is_full_fc:
            # Need to calculate full force constant tensors
            phonon.produce_force_constants(
                fc_calculator=fc_calculator,
                fc_calculator_options=fc_calculator_options)
        else:
            # Only force constants between atoms in primitive cell and
            # supercell
            phonon.produce_force_constants(
                calculate_full_force_constants=False,
                fc_calculator=fc_calculator,
                fc_calculator_options=fc_calculator_options)


def compact_fc_to_full_fc(phonon, compact_fc, log_level=0):
    fc = np.zeros((compact_fc.shape[1], compact_fc.shape[1], 3, 3),
                  dtype='double', order='C')
    fc[phonon.primitive.p2s_map] = compact_fc
    distribute_force_constants_by_translations(
        fc, phonon.primitive, phonon.supercell)
    if log_level:
        print("Force constants were expanded to full format.")

    return fc


def store_force_constants(phonon,
                          settings,
                          phpy_yaml,
                          unitcell_filename,
                          load_phonopy_yaml,
                          log_level):
    physical_units = get_default_physical_units(phonon.calculator)
    p2s_map = phonon.primitive.p2s_map

    if load_phonopy_yaml:
        is_full_fc = settings.fc_spg_symmetry or settings.is_full_fc
        if phpy_yaml.force_constants is not None:
            if log_level:
                print("Force constants were read from \"%s\"."
                      % unitcell_filename)

            fc = phpy_yaml.force_constants

            if fc.shape[1] != len(phonon.supercell):
                error_text = ("Number of atoms in supercell is not consistent "
                              "with the matrix shape of\nforce constants read "
                              "from %s." % unitcell_filename)
                print_error_message(error_text)
                if log_level:
                    print_error()
                sys.exit(1)

            # Compact fc is expanded to full fc when full fc is required.
            if is_full_fc and fc.shape[0] != fc.shape[1]:
                fc = compact_fc_to_full_fc(phonon, fc, log_level=log_level)

            phonon.set_force_constants(fc, show_drift=(log_level > 0))

        if phpy_yaml.dataset is not None:
            if log_level:
                text = "Force sets were read from \"%s\"" % unitcell_filename
                if phpy_yaml.force_constants is not None:
                    text += " but not to be used."
                else:
                    text += "."
                print(text)

        if phpy_yaml.force_constants is None:
            (fc_calculator,
             fc_calculator_options) = get_fc_calculator_params(settings)
            set_dataset_and_force_constants(
                phonon,
                phpy_yaml.dataset,
                None,
                fc_calculator=fc_calculator,
                fc_calculator_options=fc_calculator_options,
{'fraction_non_alphanumeric': 0.08463411141343262, 'fraction_numerical': 0.004746220412839345, 'mean_word_length': 2.001657967454713, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 30, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5911759', 'n_tokens_mistral': 14748, 'n_tokens_neox': 13890, 'n_words': 3378}