#!/usr/bin/env python3

import numpy as np
import sys

# We do it all in atomic units and convert only in the end
#hbar_J = 1.0545718e-34  # J*s
#kb_J = 1.38064852e-23  # J/K
#kb_ev = 8.6173303e-5  # eV/K
#ev2J = 1.60217662e-19
#J2ev = 1 / ev2J
#atomic2s = 2.418e-17
temptohartree=3.1668116e-06
evtohartree=0.036749322

def ipi_harmonic_free_energy(fpath, T, U_total):
    try:
        # We skip the 6 first eigenvalues corresponding to
        # free translations and rotations. Remember that skiprows counts commented lines
        omega_kw = np.loadtxt(fpath, skiprows=7)
    except:
        print("We cannot open the file '{}'.\n".format(fpath))
        sys.exit()
    omega_kw_abs = np.abs(omega_kw)
    omega = np.sqrt(omega_kw_abs)

    #freq = omega / (2 * np.pi)
    #freq = freq / atomic2s
    #omega = freq * 2 * np.pi

    Fc = []
    for T_ in T:
        beta = 1 / (T_)
        Fc.append(np.sum(np.log(beta * omega) / beta))

    return np.array(Fc) + U_total


# Check syntax
if len(sys.argv) != 3:
    print("")
    print("Error: wrong number of arguments.")
    print("This script takes exactly 1 input and 1 output file.\n")
    sys.exit(1)
else:
    input_file = sys.argv[1]
    output_file = sys.argv[2]

# Calculating harmonic free energy
U_total = 0.  # For simplicity, we ignore baseline pot. energy. If any other number is written here, remember to convert to Hartree.
T = np.array([10, 20, 40, 60, 80, 100, 120, 140, 160, 180, 200])
T=T*temptohartree
F = ipi_harmonic_free_energy(input_file, T, U_total)
F=F/evtohartree
#F += 3* kb_ev * T/2. # rotational contribution would be this term minus T* entropic term

data = np.stack((T/temptohartree, F), axis=1)
np.savetxt(output_file, data, fmt="%3g  %.6f", header="T (K)   F_harm (eV)")
np.savetxt(sys.stdout.buffer, data, fmt="%3g  %.6f", header="T (K)   F_harm (eV)")