#!/usr/bin/env python
import dynaphopy.interface.iofile as iofile
import dynaphopy.interface.phonopy_link as phonopy_link

import numpy as np
import yaml
import argparse
import glob

from phonopy import PhonopyQHA
from scipy.interpolate import interp1d

parser = argparse.ArgumentParser(description='thermal_extractor options')
parser.add_argument('input_file', metavar='data_file', type=str,
                    help='input file containing structure related data')

parser.add_argument('-fc', metavar='data_file', type=str, nargs='*', required=True,
                    help='phonopy force constant files')

parser.add_argument('-tp', metavar='data_file', type=str, nargs='*', required=True,
                    help='phonopy thermal properties files')

parser.add_argument('-ev', metavar='data', type=str, required=True,
                    help='Energy volume file')

parser.add_argument('-t', metavar='F', type=float, required=True,
                    help='temperature at which the renormalized force constants will be calculated')

parser.add_argument('-p', action='store_true',
                    help='plot QHA data')

parser.add_argument('-sfc', metavar='filename', type=str, default='FORCE_CONSTANTS_QHA',
                    help='select custom filename to store the renormalized force constants')


args = parser.parse_args()

# Read energy volume data
target_temperature = args.t
ev_file = np.loadtxt(args.ev)
volumes = ev_file.T[0]
electronic_energies = ev_file.T[1]

# Read force constants
fc_filenames = []
for i in args.fc:
    fc_filenames += list(sorted(glob.iglob(i)))

# Read and setup thermal_properties
tp_filenames = []
for i in args.tp:
    tp_filenames += list(sorted(glob.iglob(i)))

temperatures = []
fe_phonon = []
entropy = []
cv = []

for filename in tp_filenames:
    temperatures = []
    entropy_i = []
    fe_i = []
    cv_i = []

    with open(filename, 'r') as stream:
        thermal_properties = dict(yaml.load(stream))
        for tp in thermal_properties['thermal_properties']:
            temperatures.append(tp['temperature'])
            entropy_i.append(tp['entropy'])
            fe_i.append(tp['free_energy'])
            cv_i.append(tp['heat_capacity'])

    fe_phonon.append(fe_i)
    entropy.append(entropy_i)
    cv.append(cv_i)

sort_index = np.argsort(volumes)

volumes = np.array(volumes)[sort_index]
electronic_energies = np.array(electronic_energies)[sort_index]
temperatures = np.array(temperatures)
fe_phonon = np.array(fe_phonon).T[:, sort_index]
entropy = np.array(entropy).T[:, sort_index]
cv = np.array(cv).T[:, sort_index]

# Apply QHA using phonopy
phonopy_qha = PhonopyQHA(np.array(volumes),
                         np.array(electronic_energies),
                         eos="vinet",
                         temperatures=np.array(temperatures),
                         free_energy=np.array(fe_phonon),
                         cv=np.array(cv),
                         entropy=np.array(entropy),
                         #t_max=target_temperature,
                         verbose=False)

if args.p:
    phonopy_qha.plot_qha().show()

volume_temperature = phonopy_qha.get_volume_temperature()
qha_temperatures = phonopy_qha._qha._all_temperatures[:len(volume_temperature)]

# Fit force constants as a function of the temperature
fit_vt = interp1d(qha_temperatures, volume_temperature, kind='quadratic')
try:
    target_volume = fit_vt(target_temperature)
except ValueError:
    print('Error interpolating this temperature!\n'
          'Make sure the requested temperature is within the range of thermal_properties.yaml files')
    exit()

input_parameters = iofile.read_parameters_from_input_file(args.input_file)

if 'structure_file_name_outcar' in input_parameters:
    structure = iofile.read_from_file_structure_outcar(input_parameters['structure_file_name_outcar'])
else:
    structure = iofile.read_from_file_structure_poscar(input_parameters['structure_file_name_poscar'])
structure.get_data_from_dict(input_parameters)

fc_supercell = input_parameters['supercell_phonon']

force_constants_mat = []
for filename in fc_filenames:
    force_constants = phonopy_link.get_force_constants_from_file(filename, fc_supercell=fc_supercell)
    force_constants_mat.append(force_constants.get_array())

force_constants_mat = np.array(force_constants_mat).T
f_temperature = interp1d(volumes, force_constants_mat, kind='quadratic')

# Get force constants at the requested temperature
try:
    target_fc = f_temperature(target_volume).T
except ValueError:
    print('Error interpolating minimum volume!\n'
          'Make sure the free energy-volume curve has a minimum ')
    exit()

phonopy_link.write_FORCE_CONSTANTS(target_fc,filename=args.sfc)

print ('QHA Renormalized force constants written in file: {}'.format(args.sfc))
