#!/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
from phonopy.structure.cells import get_supercell
from phonopy.structure.symmetry import Symmetry
from phonopy.harmonic.force_constants import show_drift_force_constants
from phonopy.file_IO import parse_FORCE_SETS
from force_fit.fc2 import FC2Fit
from force_fit.fc3 import FC3Fit
from force_fit.fc4 import FC4Fit
from anharmonic.file_IO import (parse_disp_fc4_yaml, parse_FORCES_FC4,
                                parse_disp_fc3_yaml, parse_FORCES_FC3,
                                parse_disp_fc2_yaml, parse_FORCES_FC2)
from anharmonic.file_IO import (write_fc3_dat, write_fc4_dat, write_fc4_to_hdf5,
                                write_fc3_to_hdf5, write_fc2_to_hdf5)
from anharmonic.phonon3.fc3 import show_drift_fc3
from anharmonic.phonon4.fc4 import show_drift_fc4

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

def print_error(message):
    print message

parser = OptionParser()
parser.set_defaults(cell_poscar=None,
                    coef_invariants=None,
                    pinv_cutoff=None,
                    read_phonopy_files=False,
                    fc2=False,
                    fc3=False,
                    fc4=False,
                    rot_inv=False,
                    supercell_dimension=None,
                    symprec=1e-5,
                    trans_inv=False,
                    verbose=False)
parser.add_option("-c", "--cell", dest="cell_poscar",
                  action="store", type="string",
                  help="Read unit cell", metavar="FILE")
parser.add_option("--ci", dest="coef_invariants", type="float",
                  help="Coefficient to be multiplied with invariat matrix")
parser.add_option("--dim", dest="supercell_dimension",
                  type="string", help="Supercell dimension")
parser.add_option("--fc2", dest="fc2",
                  action="store_true",
                  help="Calculate fc2")
parser.add_option("--fc3", dest="fc3",
                  action="store_true",
                  help="Calculate fc3")
parser.add_option("--fc4", dest="fc4",
                  action="store_true",
                  help="Calculate fc4")
parser.add_option("--phonopy", dest="read_phonopy_files",
                  action="store_true",
                  help="Read disp.yaml and FORCE_SETS")
parser.add_option("--pinv_cutoff", dest="pinv_cutoff", type="float",
                  help="Cutoff value for pseudo-inversion")
parser.add_option("--ri", dest="rot_inv",
                  action="store_true",
                  help="Enforce rotational invariance")
parser.add_option("--tolerance", dest="symprec", type="float",
                  help="Symmetry tolerance to search")
parser.add_option("--ti", dest="trans_inv",
                  action="store_true",
                  help="Enforce translational invariance")
parser.add_option("-v", "--verbose", dest="verbose", action="store_true",
                  help="Detailed run-time information is displayed")
(options, args) = parser.parse_args()

if options.supercell_dimension is None:
    print_error("\'--dim\' has to be set")
    sys.exit(1)
    
dimension = np.array([int(x) for x in options.supercell_dimension.split()])
if len(dimension) == 3:
    dimension = np.diag(dimension)

# Read POSCAR
if options.cell_poscar is None:
    unitcell_filename = 'POSCAR'
else:
    unitcell_filename = options.cell_poscar
if os.path.exists(unitcell_filename):
    unitcell = read_vasp(unitcell_filename)
else:
    print_error("%s could not be found." % unitcell_filename)
    sys.exit(1)

# Supercell and Symmetry 
supercell = get_supercell(unitcell, dimension)
symmetry = Symmetry(supercell, options.symprec)

print "Spacegroup: ", symmetry.get_international_table()

if options.fc2:
    if options.read_phonopy_files:
        file_exists("FORCE_SETS")
        disp_dataset = parse_FORCE_SETS("FORCE_SETS")
    else:
        file_exists("disp_fc3.yaml")
        file_exists("FORCES_FC2")
        disp_dataset = parse_disp_fc2_yaml()
        forces_fc2 = parse_FORCES_FC2(disp_dataset)
        for forces, disp1 in zip(forces_fc2, disp_dataset['first_atoms']):
            disp1['forces'] = forces
    if options.coef_invariants is not None:
        print "Adjustment parameter: %e" % options.coef_invariants
    if options.pinv_cutoff is not None:
        print "Cutoff value for pseudo inversion: %e" % options.pinv_cutoff
    fc2fit = FC2Fit(supercell,
                    disp_dataset,
                    symmetry,
                    translational_invariance=options.trans_inv,
                    rotational_invariance=options.rot_inv,
                    coef_invariants=options.coef_invariants,
                    pinv_cutoff=options.pinv_cutoff)
    fc2fit.run()
    fc2 = fc2fit.get_fc2()
    print "Writing fc2..."
    write_fc2_to_hdf5(fc2, 'fc2.fit.hdf5')

if options.fc3:
    file_exists("disp_fc3.yaml")
    file_exists("FORCES_FC3")
    disp_dataset = parse_disp_fc3_yaml()
    forces_fc3 = parse_FORCES_FC3(disp_dataset)
    for forces, disp1 in zip(forces_fc3, disp_dataset['first_atoms']):
        disp1['forces'] = forces
    count = len(disp_dataset['first_atoms'])
    for disp1 in disp_dataset['first_atoms']:
        for disp2 in disp1['second_atoms']:
            disp2['forces'] = forces_fc3[count]
            count += 1
    
    fc3fit = FC3Fit(supercell, disp_dataset, symmetry, verbose=options.verbose)
    fc3fit.run()
    fc3 = fc3fit.get_fc3()
    print "Calculating drift fc3..."
    show_drift_fc3(fc3)
    print "Writing fc3..."
    write_fc3_to_hdf5(fc3, 'fc3.fit.hdf5')

if options.fc4:
    file_exists("disp_fc4.yaml")
    file_exists("FORCES_FC4")
    disp_dataset = parse_disp_fc4_yaml()
    forces_fc4 = parse_FORCES_FC4(disp_dataset)
    count = 0
    for disp1 in disp_dataset['first_atoms']:
        disp1['forces'] = forces_fc4[count]
        count += 1
    for disp1 in disp_dataset['first_atoms']:
        for disp2 in disp1['second_atoms']:
            disp2['forces'] = forces_fc4[count]
            count += 1
    for disp1 in disp_dataset['first_atoms']:
        for disp2 in disp1['second_atoms']:
            for disp3 in disp2['third_atoms']:
                disp3['forces'] = forces_fc4[count]
                count += 1
    fc4fit = FC4Fit(supercell, disp_dataset, symmetry)
    fc4fit.run()
    fc4 = fc4fit.get_fc4()
    print "Calculating drift fc4..."
    show_drift_fc4(fc4)
    print "Writing fc4..."
    write_fc4_to_hdf5(fc4, 'fc4.fit.hdf5')
    
    fc3 = fc4fit.get_fc3()
    print "Calculating drift of fc3..."
    show_drift_fc3(fc3)
    print "Writing fc3..."
    write_fc3_to_hdf5(fc3, 'fc3.fit.hdf5')

    fc2 = fc4fit.get_fc2()
    show_drift_force_constants(fc2, name="fc2")
    print "Writing fc2..."
    write_fc2_to_hdf5(fc2, 'fc2.fit.hdf5')
