#!/usr/bin/env python

# Copyright (C) 2013 Atsushi Togo
# 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 os
import sys
import numpy as np
from optparse import OptionParser

from phonopy.interface.vasp import read_vasp, write_vasp
from phonopy.structure.cells import get_supercell, get_primitive, print_cell
from phonopy.structure.spglib import get_grid_point_from_address
from phonopy.units import VaspToTHz
from phonopy.harmonic.force_constants import show_drift_force_constants
from anharmonic.phonon3.fc3 import show_drift_fc3
from anharmonic.file_IO import parse_disp_fc4_yaml,\
    write_FORCES_FC4_vasp, parse_FORCES_FC4, \
    read_fc4_from_hdf5, read_fc3_from_hdf5, read_fc2_from_hdf5, \
    write_fc4_to_hdf5, write_fc3_to_hdf5, write_fc2_to_hdf5, \
    write_freq_shifts_to_hdf5, write_disp_fc4_yaml
from anharmonic.settings import Phono3pyConfParser
from anharmonic.phonon4.fc4 import show_drift_fc4
from anharmonic.phonon4 import Phono4py

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

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

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

def print_error(message):
    print message

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

# Parse options
parser = OptionParser()
parser.set_defaults(band_indices=None,
                    cell_poscar=None,
                    displacement_distance=None,
                    factor=None,
                    forces_fc4_mode=False,
                    grid_points=None,
                    input_filename=None,
                    input_output_filename=None,
                    is_nodiag=False,
                    is_displacement=False,
                    is_nosym=False,
                    is_plusminus_displacements=False,
                    is_translational_symmetry=False,
                    is_symmetrize_fc2=False,
                    is_symmetrize_fc3_r=False,
                    is_symmetrize_fc4_r=False,
                    log_level=None,
                    mesh_numbers=None,
                    primitive_axis=None,
                    quiet=False,
                    read_fc2=False,
                    read_fc3=False,
                    read_fc4=False,
                    output_filename=None,
                    supercell_dimension=None,
                    symprec=1e-5,
                    temperatures=None,
                    tmax=None,
                    tmin=None,
                    tstep=None,
                    verbose=False)
parser.add_option("--amplitude", dest="displacement_distance", type="float",
                  help="Distance of displacements")
parser.add_option("--bi", "--band_indices", dest="band_indices",
                  type="string",
                  help="Band indices where life time is calculated")
parser.add_option("-c", "--cell", dest="cell_poscar",
                  action="store", type="string",
                  help="Read unit cell", metavar="FILE")
parser.add_option("--cf4", "--create_f4",
                  dest="forces_fc4_mode",
                  action="store_true",
                  help="Create FORCES_FC4")
parser.add_option("-d", "--disp", dest="is_displacement",
                  action="store_true",
                  help="As first stage, get least displacements")
parser.add_option("--dim",
                  dest="supercell_dimension",
                  type="string",
                  help="Supercell dimension")
parser.add_option("--factor", dest="factor", type="float",
                  help="Conversion factor to favorite frequency unit")
parser.add_option("--fc2",
                  dest="read_fc2",
                  action="store_true",
                  help="Read second order force constants")
parser.add_option("--fc3",
                  dest="read_fc3",
                  action="store_true",
                  help="Read third order force constants")
parser.add_option("--fc4",
                  dest="read_fc4",
                  action="store_true",
                  help="Read fourth order force constants")
parser.add_option("--ga", "--grid_addresses",
                  dest="grid_addresses", type="string",
                  help=("Fixed grid addresses where anharmonic properties "
                        "are calculated"))
parser.add_option("--gp", "--grid_points",
                  dest="grid_points",
                  type="string",
                  help=("Fixed grid points where damping functions are "
                        "calculated "))
parser.add_option("-i", dest="input_filename", type="string",
                  help="Input filename extension")
parser.add_option("--io", dest="input_output_filename", type="string",
                  help="Input and output filename extension")
parser.add_option("--mesh",
                  dest="mesh_numbers",
                  type="string",
                  help="Mesh numbers")
parser.add_option("--nodiag", dest="is_nodiag",
                  action="store_true",
                  help="Set displacements parallel to axes")
parser.add_option("--nosym", dest="is_nosym",
                  action="store_true",
                  help="No symmetrization of triplets")
parser.add_option("-o", dest="output_filename", type="string",
                  help="Output filename extension")
parser.add_option("--pa", "--primitive_axis", dest="primitive_axis",
                  action="store", type="string",
                  help="Same as PRIMITIVE_AXIS tags")
parser.add_option("--pm", dest="is_plusminus_displacements",
                  action="store_true",
                  help="Set plus minus displacements")
parser.add_option("-q", "--quiet", dest="quiet", action="store_true",
                  help="Print out smallest information")
parser.add_option("--sym_fc2", dest="is_symmetrize_fc2",
                  action="store_true",
                  help="Symmetrize fc2 by index exchange")
parser.add_option("--sym_fc3r", dest="is_symmetrize_fc3_r",
                  action="store_true",
                  help="Symmetrize fc3 in real space by index exchange")
parser.add_option("--sym_fc4r", dest="is_symmetrize_fc4_r",
                  action="store_true",
                  help="Symmetrize fc4 in real space by index exchange")
parser.add_option("--tmax", dest="tmax", type="string",
                  help="Maximum calculated temperature")
parser.add_option("--tmin", dest="tmin", type="string",
                  help="Minimum calculated temperature")
parser.add_option("--ts", dest="temperatures", 
                  type="string", help="Temperatures for damping functions")
parser.add_option("--tstep", dest="tstep", type="string",
                  help="Calculated temperature step")
parser.add_option("--tsym", dest="is_translational_symmetry",
                  action="store_true",
                  help="Impose translational invariance condition")
parser.add_option("--tolerance", dest="symprec", type="float",
                  help="Symmetry tolerance to search")
parser.add_option("-v", "--verbose", dest="verbose",
                  action="store_true",
                  help="Detailed run-time information is displayed")
parser.add_option("--loglevel", dest="log_level", type="int",
                  help="Log level")
(options, args) = parser.parse_args()
option_list = parser.option_list

# Log level
log_level = 1
if options.verbose:
    log_level = 2
if options.quiet:
    log_level = 0
if options.log_level is not None:
    log_level=options.log_level

# Input and output filename extension
input_filename = options.input_filename
output_filename = options.output_filename
if options.input_output_filename is not None:
    input_filename = options.input_output_filename
    output_filename = options.input_output_filename
    
# Title
if log_level:
    print_phono4py()

#####################
# Create FORCES_FC4 #
#####################
if options.forces_fc4_mode:
    if input_filename is None:
        filename = 'disp_fc4.yaml'
    else:
        filename = 'disp_fc4.' + input_filename + '.yaml'
    file_exists(filename, log_level)
    if log_level:
        print "Displacement dataset is read from %s." % filename
    disp_dataset = parse_disp_fc4_yaml()

    write_FORCES_FC4_vasp(args, disp_dataset)
    
    if log_level:
        print "FORCES_FC4 has been created."
        print_end()
    exit(0)

##################
# Parse settings #
##################
if len(args) > 0:
    phono3py_conf = Phono3pyConfParser(filename=args[0],
                                       options=options,
                                       option_list=option_list)
    settings = phono3py_conf.get_settings()

else:
    phono3py_conf = Phono3pyConfParser(options=options,
                                       option_list=option_list)
    settings = phono3py_conf.get_settings()

###################################
# Read crystal structure (POSCAR) #
###################################
if options.cell_poscar is None:
    file_exists('POSCAR', log_level)
    unitcell_filename = 'POSCAR'
else:
    file_exists(options.cell_poscar, log_level)
    unitcell_filename = options.cell_poscar
unitcell = read_vasp(unitcell_filename, settings.get_chemical_symbols())

#################################################
# Create supercells with displacements and exit #
#################################################
if options.is_displacement:
    if settings.get_displacement_distance() is None:
        displacement_distance = 0.03
    else:
        displacement_distance = settings.get_displacement_distance()
    phono4py = Phono4py(unitcell,
                        settings.get_supercell_matrix(),
                        symprec=options.symprec)
    phono4py.generate_displacements(
        distance=displacement_distance,
        is_plusminus=settings.get_is_plusminus_displacement(),
        is_diagonal=settings.get_is_diagonal_displacement())
    dds = phono4py.get_displacement_dataset()

    if log_level:
        print
        print "Displacement distance:", displacement_distance

    if output_filename is None:
        filename = 'disp_fc4.yaml'
    else:
        filename = 'disp_fc4.' + output_filename + '.yaml'
        
    supercell = phono4py.get_supercell()
    num_disps = write_disp_fc4_yaml(dds, supercell, filename=filename)
    for i, dcell in enumerate(phono4py.get_supercells_with_displacements()):
        write_vasp('POSCAR-%05d' % (i + 1), dcell, direct=True)

    if log_level:
        print "Total number of displacements:", num_disps
        print_end()
    sys.exit(0)
    
##############
# Initialize #
##############
if settings.get_temperatures() is None:
    t_max=settings.get_max_temperature()
    t_min=settings.get_min_temperature()
    t_step=settings.get_temperature_step()
    temperatures = np.arange(t_min, t_max + float(t_step) / 10, t_step)
else:
    temperatures = settings.get_temperatures()

if options.factor is None:
    factor = VaspToTHz
else:
    factor = options.factor

mesh = settings.get_mesh_numbers()

phono4py = Phono4py(unitcell,
                    settings.get_supercell_matrix(),
                    primitive_matrix=settings.get_primitive_matrix(),
                    mesh=mesh,
                    band_indices=settings.get_band_indices(),
                    frequency_factor_to_THz=factor,
                    cutoff_frequency=1e-2,
                    is_nosym=options.is_nosym,
                    symprec=options.symprec,
                    log_level=log_level)

primitive = phono4py.get_primitive()
supercell = phono4py.get_supercell()

if log_level:
    print "------------------------ primitive cell for fc ---------------------------"
    print_cell(primitive)
    print "-------------------------- supercell for fc ------------------------------"
    print_cell(supercell, mapping=primitive.get_supercell_to_primitive_map())
    print "----------------- ratio (supercell for fc)/(primitive) -------------------"
    for vec in np.dot(supercell.get_cell(), np.linalg.inv(primitive.get_cell())):
        print "%5.2f"*3 % tuple(vec)


#############################
# Calculate force constants #
#############################
if not options.read_fc4:
    displacements = parse_disp_fc4_yaml()
    forces_fc4 = parse_FORCES_FC4(displacements)
    translational_symmetry_type = options.is_translational_symmetry * 1
    phono4py.produce_fc4(
        forces_fc4,
        displacements,
        translational_symmetry_type=translational_symmetry_type,
        is_permutation_symmetry=options.is_symmetrize_fc4_r,
        is_permutation_symmetry_fc3=options.is_symmetrize_fc3_r,
        is_permutation_symmetry_fc2=options.is_symmetrize_fc2)

if options.read_fc2:
    if input_filename is None:
        filename = "fc2.hdf5"
    else:
        filename = "fc2.%s.hdf5" % input_filename
    if log_level:
        print  "Reading fc2 from %s" % filename
        sys.stdout.flush()
    phono4py.set_fc2(read_fc2_from_hdf5(filename=filename))
else:
    if output_filename is None:
        filename = "fc2.hdf5"
    else:
        filename = "fc2.%s.hdf5" % output_filename
    if log_level:
        print "Writing fc2 to %s" % filename
    write_fc2_to_hdf5(phono4py.get_fc2(), filename=filename)
fc2 = phono4py.get_fc2()
show_drift_force_constants(fc2, name='fc2')

if options.read_fc3:
    if input_filename is None:
        filename = "fc3.hdf5"
    else:
        filename = "fc3.%s.hdf5" % input_filename
    if log_level:
        print  "Reading fc3 from %s" % filename
        sys.stdout.flush()
    phono4py.set_fc3(read_fc3_from_hdf5(filename=filename))
else:
    if output_filename is None:
        filename = "fc3.hdf5"
    else:
        filename = "fc3.%s.hdf5" % output_filename
    if log_level:
        print "Writing fc3 to %s" % filename
    write_fc3_to_hdf5(phono4py.get_fc3(), filename=filename)
fc3 = phono4py.get_fc3()
show_drift_fc3(fc3)

if options.read_fc4:
    if input_filename is None:
        filename = "fc4.hdf5"
    else:
        filename = "fc4.%s.hdf5" % input_filename
    if log_level:
        print  "Reading fc4 from %s" % filename
        sys.stdout.flush()
    phono4py.set_fc4(read_fc4_from_hdf5(filename=filename))
else:
    if output_filename is None:
        filename = "fc4.hdf5"
    else:
        filename = "fc4.%s.hdf5" % output_filename
    if log_level:
        print "Writing fc4 to %s" % filename
    write_fc4_to_hdf5(phono4py.get_fc4(), filename=filename)
fc4 = phono4py.get_fc4()
show_drift_fc4(fc4)


###################
# Frequency shift #
###################
if mesh is not None:
    grid_points = settings.get_grid_points()
    grid_addresses = settings.get_grid_addresses()
    if grid_addresses is not None:
        grid_points = [get_grid_point_from_address(ga, mesh)
                       for ga in grid_addresses]

    phono4py.set_frequency_shift(temperatures=temperatures)
    phono4py.set_dynamical_matrix(fc2,
                                  supercell,
                                  primitive)
    phono4py.run_frequency_shift(grid_points)
    freq_shifts = phono4py.get_frequency_shift()
    write_freq_shifts_to_hdf5(freq_shifts)
                            
if log_level:
    print_end()


    
