#!/usr/bin/python3.5

import math
import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize as opt

unit_bohr2angstrom = 0.529177249
unit_hartree2cm    = 219474.631371

def func_LJ(x,a,c):
    # The Lennard-Jones model
    return a/x**12 - c/x**6

def func_BM(x,a,b,c):
    # The Born-Mayer model
    #return a*np.exp(-b*x) - c/x**6
    return np.exp(-b*(x-a)) - c/x**6

data = np.transpose(np.loadtxt('Ar2_MP2.txt'))
rvals = data[0]/unit_bohr2angstrom # convert from Angstrom to Bohr
evals = data[1]/unit_hartree2cm    # convert from cm-1 to Hartree


# Standard deviations all set to 1.0
sigma = sigma = np.full(rvals.size,1.0)

print("sigma array: \n")
print(sigma)

# LJ fit
# ======
# Starting values for parameters
x0=np.array([1.0,1.0])
# Do the fitting:
popt, pcov = opt.curve_fit(func_LJ,rvals,evals,p0=x0,sigma=sigma)
#
# popt : optimized parameters
# pcov : covariance matrix
print("\n\n")
print(" a = %12.4f  c = %12.4f  \n" % (popt[0],popt[1]))
print("Covariance matrix: \n")
ndim = pcov.shape[0]
for i in range(0,ndim):
    for j in range(0,ndim):
        pcov[i][j] = pcov[i][j]/np.sqrt(pcov[i][i]*pcov[j][j])
print(pcov)
print("\n\n")
#
# Plot
plt.plot(rvals,evals,'b-', label='data',marker='o')
plt.plot(rvals, func_LJ(rvals, *popt), 'g--', label='fit')
plt.show()

# BM fit
# ======
# Starting values for parameters
x0=np.array([1.0,1.0,1.0])
# Do the fitting:
popt, pcov = opt.curve_fit(func_BM,rvals,evals,p0=x0,sigma=sigma)
#
# popt : optimized parameters
# pcov : covariance matrix
print("\n\n")
print(" a = %12.4f  b = %12.4f  c = %12.4f  \n" % (popt[0],popt[1],popt[2]))
print("Covariance matrix: \n")
ndim = pcov.shape[0]
for i in range(0,ndim):
    for j in range(0,ndim):
        pcov[i][j] = pcov[i][j]/np.sqrt(pcov[i][i]*pcov[j][j])
print(pcov)
print("\n\n")
#
# Plot
plt.plot(rvals,evals,'b-', label='data',marker='o')
plt.plot(rvals, func_BM(rvals, *popt), 'g--', label='fit')
plt.show()


# Now try a weighting of exp(-(e-e0)/kT)
e0 =   80.0/unit_hartree2cm  # both in units of evals
kT =  400.0/unit_hartree2cm
#sigma = np.exp(-np.abs(evals-e0)/kT)
sigma = np.exp(-np.abs(evals/e0))

print("sigma array: \n")
print(sigma)

# LJ fit
# ======
# Starting values for parameters
x0=np.array([1.0,1.0])
# and re-fit
popt, pcov = opt.curve_fit(func_LJ,rvals,evals,p0=x0,sigma=sigma)
#
# popt : optimized parameters
# pcov : covariance matrix
print("\n\n")
print(" a = %12.4f  c = %12.4f  \n" % (popt[0],popt[1]))
print("Covariance matrix: \n")
ndim = pcov.shape[0]
for i in range(0,ndim):
    for j in range(0,ndim):
        pcov[i][j] = pcov[i][j]/np.sqrt(pcov[i][i]*pcov[j][j])
print(pcov)
print("\n\n")
#
# Plot
plt.plot(rvals,evals,'b-', label='data',marker='o')
plt.plot(rvals, func_LJ(rvals, *popt), 'g--', label='fit')
plt.show()