#!/home/travis/miniconda/bin/python

# Copyright (C) 2011 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 numpy as np
import sys
from phonopy.units import *
from phonopy.interface.vasp import read_vasp, write_vasp
from phonopy import PhonopyQHA
from phonopy.file_IO import (read_thermal_properties_yaml, read_v_e,
                             read_cp, read_ve)


def create_unitcells(unitcell_filename,
                     plus_percent=10.0,
                     minus_percent=-5.0,
                     number_of_unitcells=10):
    cell = read_vasp(unitcell_filename)
    lattice = cell.get_cell()
    print("# Volumes of unit cells")
    for i, percent in enumerate(np.linspace(minus_percent,
                                            plus_percent,
                                            number_of_unitcells)):
        lattice_qha = lattice * (1 + percent / 100) ** (1.0 / 3)
        cell.set_cell(lattice_qha)
        print("%f" % np.linalg.det(lattice_qha))
        write_vasp("QHAPOSCAR-%02d" % (i + 1), cell, direct=True)

def get_options():
    # Parse options
    import argparse
    parser = argparse.ArgumentParser(
        description="Phonopy-QHA command-line-tool")
    parser.set_defaults(create_unitcells_range=None,
                        factor=1.0,
                        energy_plot_factor=None,
                        volume_factor=1.0,
                        pressure=0.0,
                        is_create_unitcells=False,
                        is_graph_plot=False,
                        is_graph_save=False,
                        is_bulk_modulus_only=False,
                        eos="vinet",
                        energy_shift=None,
                        symbol='o',
                        cp_filename=None,
                        ve_filename=None,
                        thin_number=10,
                        t_max=1000.0)
    parser.add_argument(
        "-b", dest="is_bulk_modulus_only", action="store_true",
        help="Just show Bulk modulus from v-e data")
    parser.add_argument(
        "--cp", "--heat_capacity", dest="cp_filename", metavar="FILE",
        help="Experimental data for heat capacity Cv")
    parser.add_argument(
        "--cu", "--create_unitcells", dest="is_create_unitcells",
        action="store_true",
        help="Create unitcells")
    parser.add_argument(
        "--cu_range", "--create_unitcells_range", dest="create_unitcells_range",
        help=("Range of created unitcells, plus, minus in percent, "
              "and number of unit cells"))
    parser.add_argument(
        "--eos", dest="eos",
        help="Choise of EOS among vinet, birch_murnaghan, and murnaghan")
    parser.add_argument(
        "--es", "--energy_shift", dest="energy_shift",
        help="Energy shift that is subtracted from Helmholtz free energy")
    parser.add_argument(
        "--exclude_imaginary", dest="exclude_imaginary", action="store_true",
        help="Exclude volumes that show imaginary modes")
    parser.add_argument(
        "--factor", dest="factor", type=float,
        help=("Conversion factor that is multiplied with all "
              "extensive variables"))
    parser.add_argument(
        "-p", "--plot", dest="is_graph_plot", action="store_true",
        help="Plot data")
    parser.add_argument(
        "--pressure", dest="pressure", type=float,
        help="Pressure in GPa")
    parser.add_argument(
        "-s", "--save", dest="is_graph_save", action="store_true",
        help="Save plot data in pdf")
    parser.add_argument(
        "--symbol", dest="symbol",
        help="Symbol used to plot experiment points")
    parser.add_argument(
        "--sparse", dest="thin_number", type=int,
        help=("Thin out the F-V plots of temperature. The value is "
              "used as deviser of number of temperature points."))
    parser.add_argument(
        "--tmax", dest="t_max", type=float,
        help="Maximum calculated temperature")
    parser.add_argument(
        "--ve", "--volume_expansion", dest="ve_filename", metavar="FILE",
        help="Experimental data for volume expansion")
    parser.add_argument(
        "--vf", "--volume_factor", dest="volume_factor", type=float,
        help="Conversion factor of volume to A^3")
    parser.add_argument(
        "--epf", "--energy_plot_factor", dest="energy_plot_factor", type=float,
        help=("Conversion factor of energy for plot after fitting, "
              "i.e., the input data have to have the proper physical"
              "units, but plotting is done with multiplied with this "
              "factor."))
    parser.add_argument(
        "filenames", nargs='*',
        help="Filenames of e-v.dat and thermal_properties.yaml's")
    args = parser.parse_args()
    return args

def main(args):
    if args.is_graph_save:
        import matplotlib
        matplotlib.use('Agg')
        matplotlib.rc('pdf', fonttype=42)

    ####################
    # Create unitcells #
    ####################
    if args.is_create_unitcells:
        if args.create_unitcells_range:
            plus, minus, num_div = args.create_unitcells_range.split()
            create_unitcells(args.filenames[0],
                             float(plus), float(minus), int(num_div))
        else:
            create_unitcells(args.filenames[0])
        sys.exit(0)

    ########################
    # Read data from files #
    ########################
    volumes, electronic_energies = read_v_e(args.filenames[0],
                                            factor=args.factor,
                                            volume_factor=args.volume_factor,
                                            pressure=args.pressure)
    if args.energy_shift is not None:
        electronic_energies -= args.energy_shift

    # Choose EOS
    if args.eos == "birch_murnaghan":
        print("# Third-order Birch-Murnaghan EOS")
    elif args.eos == "murnaghan":
        print("# Murnaghan EOS")
    else:
        print("# Vinet EOS")

    # Show bulk modulus of v-e data
    if args.is_bulk_modulus_only:
        bulk_modulus = PhonopyQHA(volumes,
                                  electronic_energies,
                                  eos=args.eos)
        parameters = bulk_modulus.get_bulk_modulus_parameters()
        print("Volume: %f" % parameters[3])
        print("Energy: %f" % parameters[0])
        print("Bulk modulus: %f" % (parameters[1] * EVAngstromToGPa))
        print("Parameters: %f %f %f %f" % tuple(parameters))
        if args.is_graph_plot:
            bulk_modulus.plot_bulk_modulus().show()
        sys.exit(0)

    # Check number of files
    if len(volumes) != len(args.filenames[1:]):
        print("The number of thermal_properites.yaml files (%d) "
              "is inconsisten with" % len(args.filenames[1:]))
        print("the number of e-v data (%d)." % len(volumes))
        sys.exit(1)

    # read data from thermal_properties
    t_props = read_thermal_properties_yaml(args.filenames[1:(len(volumes)+1)],
                                           factor=args.factor)
    T, cv, entropy, fe_phonon, num_modes, num_integrated_modes = t_props

    # Treatment of thermal properties with imaginary modes
    if args.exclude_imaginary and num_modes:
        indices = []
        num_imag_modes = np.array(num_modes) - np.array(num_integrated_modes)
        for i, nim in enumerate(num_imag_modes):
            if nim < 4:
                indices.append(i)
    else:
        indices = range(len(volumes))

    ##########################
    # Analyzing and plotting #
    ##########################
    phonopy_qha = PhonopyQHA(volumes[indices],
                             electronic_energies[indices],
                             eos=args.eos,
                             temperatures=T,
                             free_energy=fe_phonon[:, indices],
                             cv=cv[:, indices],
                             entropy=entropy[:, indices],
                             t_max=args.t_max,
                             energy_plot_factor=args.energy_plot_factor,
                             verbose=True)

    if num_modes:
        num_imag_modes = np.array(num_modes) - np.array(num_integrated_modes)
        for filename, nim in zip(args.filenames[1:(len(volumes)+1)],
                                 num_imag_modes):
            if nim > 3:
                if args.exclude_imaginary:
                    print("# %s has been excluded." % filename)
                else:
                    print("# Warning: %s has imaginary modes." % filename)

    if args.is_graph_plot and not args.is_graph_save:
        # Plot on display
        # - Volume vs Helmholtz free energy
        # - Volume vs Temperature
        # - Thermal expansion coefficient
        if args.ve_filename is None:
            exp_data = None
        else:
            exp_data = parse_ve(args.ve_filename)

        phonopy_qha.plot_qha(thin_number=args.thin_number,
                             volume_temp_exp=exp_data).show()

    if args.is_graph_save:
        # Volume vs Helmholts free energy
        phonopy_qha.plot_pdf_helmholtz_volume(thin_number=args.thin_number)

        # Volume vs Temperature
        if args.ve_filename is None:
            exp_data = None
        else:
            exp_data = parse_ve(args.ve_filename)
        phonopy_qha.plot_pdf_volume_temperature(exp_data=exp_data)

        # Thermal expansion coefficient
        phonopy_qha.plot_pdf_thermal_expansion()

        # Volume expansion
        if args.ve_filename is None:
            exp_data = None
        else:
            exp_data = parse_ve(args.ve_filename)
        phonopy_qha.plot_pdf_volume_expansion(exp_data=exp_data,
                                      symbol=args.symbol)

        # G vs Temperature
        phonopy_qha.plot_pdf_gibbs_temperature()

        # Bulk modulus vs Temperature
        phonopy_qha.plot_pdf_bulk_modulus_temperature()

        # C_P vs Temperature
        if args.cp_filename is None:
            exp_data = None
        else:
            exp_data = read_cp(args.cp_filename)
        phonopy_qha.plot_pdf_heat_capacity_P_numerical(exp_data=exp_data)

        # C_P vs Temperature (poly fit)
        if args.cp_filename is None:
            exp_data = None
        else:
            exp_data = read_cp(args.cp_filename)
        phonopy_qha.plot_pdf_heat_capacity_P_polyfit(exp_data=exp_data)

        # Gruneisen parameter vs Temperature
        phonopy_qha.plot_pdf_gruneisen_temperature()

    phonopy_qha.write_helmholtz_volume()
    phonopy_qha.write_volume_temperature()
    phonopy_qha.write_thermal_expansion()
    phonopy_qha.write_volume_expansion()
    phonopy_qha.write_gibbs_temperature()
    phonopy_qha.write_bulk_modulus_temperature()
    phonopy_qha.write_heat_capacity_P_numerical()
    phonopy_qha.write_heat_capacity_P_polyfit()
    phonopy_qha.write_gruneisen_temperature()

if __name__ == "__main__":
    main(get_options())
