import numpy as np
import metalwalls
import sys
import os
import math

me = 0

from mpi4py import MPI
comm = MPI.COMM_WORLD
fcomm = comm.py2f()
me = comm.Get_rank()
nprocs = comm.Get_size()

hartree_si = 4.359744650e-18   # hartree energy in J
avogadro_si = 6.022140857e+23  # (1/mol)
hartree2kjpermol = hartree_si * avogadro_si * 1.0e-3

# ION-O parameters
epsNa = 0.521578      
sigmaNa = 2.876500      
epsK = 0.5216   
sigmaK = 3.25
   
#Initialize mpi
my_parallel = metalwalls.mw_parallel.MW_parallel_t()
metalwalls.mw_tools.start_parallel_python(my_parallel, fcomm)

def run_lambda(lambd, num_steps_per_cycle, num_cycles):
   my_system = metalwalls.mw_system.MW_system_t()
   my_algorithms = metalwalls.mw_algorithms.MW_algorithms_t()
   
   #Initialize simulation, read input files
   do_output = metalwalls.mw_tools.initialize(my_system, my_parallel, my_algorithms)
   
   #Create two lj derived types for Na and K
   ljNa = metalwalls.mw_lennard_jones.copy_type(my_system.lj) 
   ljK = metalwalls.mw_lennard_jones.copy_type(my_system.lj) 
   
   #ION-O pairs: O=1, ION=4, only one pair because one ION 
   metalwalls.mw_lennard_jones.set_parameters(ljNa, 1, 4, epsNa, sigmaNa)
   metalwalls.mw_lennard_jones.set_parameters(ljK, 1, 4, epsK, sigmaK)
   
   #Setup simulation
   step_output_frequency = metalwalls.mw_tools.setup(my_system, my_parallel, my_algorithms, do_output)
  
   #Initialize run, timers 
   metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, 0, num_steps_per_cycle, do_output, step_output_frequency, True, False)
   
   energy_diff = np.zeros(num_cycles)
   for cycle in range(num_cycles):
      #Compute energy derivative dU/dlambd = U2 - U1
      metalwalls.mw_lookup_table.lut_set_all(my_system.energy, 0.0)
      my_system.lj = ljNa
      metalwalls.mw_system.energy_lj(my_system)
      hNa = metalwalls.mw_lookup_table.lut_get(my_system.energy, "Lennard-Jones")
      hNa += metalwalls.mw_lookup_table.lut_get(my_system.energy, "Lennard-Jones_cutoff")
      metalwalls.mw_lookup_table.lut_set_all(my_system.energy, 0.0)
      my_system.lj = ljK
      metalwalls.mw_system.energy_lj(my_system)
      hK = metalwalls.mw_lookup_table.lut_get(my_system.energy, "Lennard-Jones")
      hK += metalwalls.mw_lookup_table.lut_get(my_system.energy, "Lennard-Jones_cutoff")
      energy_diff[cycle] = hK - hNa
      
      #Run steps
      for step in range(num_steps_per_cycle):
         metalwalls.mw_tools.step_setup(my_system, my_algorithms, my_parallel, (cycle + 1) * num_steps_per_cycle + step)
         #Set forces to zero
         my_system.coulomb_forces_ions[:,:,:] = 0.0
         my_system.intermolecular_forces_ions[:,:,:] = 0.0
         my_system.intramolecular_forces_ions[:,:,:] = 0.0
         #Compute coulomb forces
         metalwalls.mw_system.forces_compute_coulomb(my_system)
         #Compute intermolecular forces
         uNa = np.zeros(np.shape(my_system.intermolecular_forces_ions))
         uK = np.zeros(np.shape(my_system.intermolecular_forces_ions))
         # Compute with model Na
         my_system.lj = ljNa
         metalwalls.mw_system.forces_compute_lj(my_system)
         uNa[:, :, :] = my_system.intermolecular_forces_ions[:, :, :]
         my_system.intermolecular_forces_ions[:,:,:] = 0.0
         my_system.lj = ljK
         metalwalls.mw_system.forces_compute_lj(my_system)
         uK[:, :, :] = my_system.intermolecular_forces_ions[:, :, :]
         my_system.intermolecular_forces_ions[:, :, :] = (1 - lambd) * uNa[:, :, :] + lambd * uK[:, :, :]
         #Gather all forces
         metalwalls.mw_system.forces_gather(my_system)
         metalwalls.mw_tools.step_update_output(my_system, my_algorithms, (cycle + 1) * num_steps_per_cycle + step, do_output, step_output_frequency)
   
   #Finalize run, timers
   metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, num_steps_per_cycle * num_cycles, 0, do_output, step_output_frequency, False, True)

   #Finalize   
   metalwalls.mw_tools.finalize(my_system, my_parallel, my_algorithms)

   student = 2.132 / math.sqrt(5)   # see https://en.wikipedia.org/wiki/Student%27s_t-distribution
   mean_student = np.zeros(5)
   np.random.shuffle(energy_diff)
   interval = int(num_cycles / 5)
   for block in range(5):
      mean_student[block] = np.mean(energy_diff[interval * block:interval * (block + 1) - 1])
   mean = np.mean(mean_student)
   std = np.std(mean_student) * student

   return mean, std

##################################################################

energy_deriv = np.zeros((11, 3))
for i, lambd in enumerate(np.arange(0, 1.1, 0.1)):
   energy_deriv[i, 0] = lambd
   energy_deriv[i, 1], energy_deriv[i, 2] = run_lambda(lambd, 100, 50)   
   if me == 0: 
      print('Lambda {:4.2f}: energy_difference {} hartree; {} kJ/mol'.format(lambd, energy_deriv[i, 1], energy_deriv[i, 1] * hartree2kjpermol)); sys.stdout.flush()
      os.system("mv trajectories.xyz trajectories_{:4.2f}.xyz".format(lambd))
      os.system("mv run.out run_{:4.2f}.out".format(lambd))
   comm.Barrier()

metalwalls.mw_tools.finalize_parallel()

import matplotlib.pyplot as plt

if me==0:
   energy_deriv[:, 1:] *= hartree2kjpermol
   np.savetxt('energy_derivative.out', energy_deriv)
   ti = np.trapz(energy_deriv[:, 1], x=energy_deriv[:, 0])

   plt.xlabel(r'$\lambda$')
   plt.xlim(-0.1, 1.1)
   plt.ylabel(r'$\frac{dU}{d\lambda}$ (kJ/mol)')
   plt.title(r'$\Delta \Delta _h F = ${:5.3f} kJ/mol'.format(ti))
   plt.errorbar(energy_deriv[:, 0], energy_deriv[:, 1], yerr=energy_deriv[:, 2])
   plt.savefig('ti_lj_lj.png', bbox_inches='tight')
   plt.show()

