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

# Run a simulation using a hybrid hamiltonian at a given lambd
# Model1: ion_type with charge q1
# Model2: ion_type with charge q2
def ti_charge_charge_lambda(q1, q2, ion_type, lambd, num_steps_per_cycle, num_cycles, my_parallel):
   # Setup the simulation
   my_system, my_parallel, my_algorithms, do_output, step_output_frequency = simulation.setup_simulation(False, my_parallel)
   my_system.num_steps = num_steps_per_cycle * num_cycles
   
   # Setup the timers 
   metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, 0, 0, do_output, step_output_frequency, True, False)
   
   # Cycle for num_cycles
   energy_derivative = np.zeros(num_cycles)
   for cycle in range(num_cycles):
      # Compute energy derivative dU/dlambd
      # Compute energy with charge1
      metalwalls.mw_lookup_table.lut_set_all(my_system.energy, 0.0)
      metalwalls.mw_tools.modify_charge(my_system, ion_type, q1)
      metalwalls.mw_system.energy_coulomb(my_system)
      h1 = metalwalls.mw_lookup_table.lut_get(my_system.energy, "Coulomb")
      # Compute energy with charge2
      metalwalls.mw_lookup_table.lut_set_all(my_system.energy, 0.0)
      metalwalls.mw_tools.modify_charge(my_system, ion_type, q2)
      metalwalls.mw_system.energy_coulomb(my_system)
      h2 = metalwalls.mw_lookup_table.lut_get(my_system.energy, "Coulomb")
      # Compute energy derivative dU/dlambd = h2-h1
      energy_derivative[cycle] = h2 - h1
      
      # Run MD steps
      for step in range(num_steps_per_cycle):
         # First step of MD step: propagate positions, compute charges...
         metalwalls.mw_tools.step_setup(my_system, my_algorithms, my_parallel, cycle * num_steps_per_cycle + step + 1)

         # Second step of MD step: compute forces
         # Compute intermolecular forces
         my_system.intermolecular_forces_ions[:,:,:] = 0.0
         metalwalls.mw_system.forces_compute_lj(my_system)
         metalwalls.mw_system.forces_compute_ft(my_system)
         metalwalls.mw_system.forces_compute_xft(my_system)
         metalwalls.mw_system.forces_compute_daim(my_system)
         # Compute intramolecular forces
         my_system.intramolecular_forces_ions[:,:,:] = 0.0
         metalwalls.mw_system.forces_compute_intramolecular(my_system)
         # Compute coulomb forces
         # with charge1
         my_system.coulomb_forces_ions[:,:,:] = 0.0
         f1 = np.zeros(np.shape(my_system.coulomb_forces_ions))
         metalwalls.mw_tools.modify_charge(my_system, ion_type, q1)
         metalwalls.mw_system.forces_compute_coulomb(my_system)
         f1[:, :, :] = my_system.coulomb_forces_ions[:,:,:]
         # with charge2
         my_system.coulomb_forces_ions[:,:,:] = 0.0
         f2 = np.zeros(np.shape(my_system.coulomb_forces_ions))
         metalwalls.mw_tools.modify_charge(my_system, ion_type, q2)
         metalwalls.mw_system.forces_compute_coulomb(my_system)
         f2[:, :, :] = my_system.coulomb_forces_ions[:,:,:]
         # Compute linear combination of forces
         my_system.coulomb_forces_ions[:, :, :] = (1 - lambd) * f1[:, :, :] + lambd * f2[:, :, :]
         # Gather all forces
         metalwalls.mw_system.forces_gather(my_system)

         # Third step of MD step: update velocities, output info
         metalwalls.mw_tools.step_update_output(my_system, my_algorithms, cycle * num_steps_per_cycle + step + 1, do_output, step_output_frequency)
   
   # Stop timers, output diagnostics
   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)

   mean, std = statistical_analysis(energy_derivative)

   return mean, std


# Run a series of TI simulations with different lambda values given by lambd_list
# Model1: ion_type with charge q1
# Model2: ion_type with charge q2
def ti_charge_charge(q1, q2, ion_type, lambd_list, num_steps_per_cycle, num_cycles, output_file='energy_derivative.out'):
   # Define constants
   hartree_si = 4.359744650e-18   # hartree energy in J
   avogadro_si = 6.022140857e+23  # (1/mol)
   hartree2kjpermol = hartree_si * avogadro_si * 1.0e-3

   # Initialize mpi
   my_parallel, comm = simulation.initialize_parallel()
   fcomm = comm.py2f()
   metalwalls.mw_tools.start_parallel_python(my_parallel, fcomm)
   
   energy_derivative = np.zeros((np.size(np.array(lambd_list)), 3))
   for i, lambd in enumerate(lambd_list):
      energy_derivative[i, 0] = lambd
      energy_derivative[i, 1], energy_derivative[i, 2] = ti_charge_charge_lambda(q1, q2, ion_type, lambd, num_steps_per_cycle, num_cycles, my_parallel)

      if my_parallel.comm_rank == 0: 
         print('Lambda {:4.2f}: energy_derivative {} hartree; {} kJ/mol'.format(lambd, energy_derivative[i, 1], energy_derivative[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()

   if my_parallel.comm_rank == 0:
      energy_derivative[:, 1:] *= hartree2kjpermol
      np.savetxt(output_file, energy_derivative, header='# Lambda   dU/dlambda (kJ/mol)   std(dU/dlambda) (kJ/mol)\n')
   
   metalwalls.mw_tools.finalize_parallel()



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

# Run a simulation using a hybrid hamiltonian at a given lambd
# Model1 with eps1, sigma1
# Model2 with eps2, sigma2
def ti_lj_lj_lambda(eps1, sigma1, eps2, sigma2, lambd, num_steps_per_cycle, num_cycles, my_parallel):
   # Setup the simulation
   my_system, my_parallel, my_algorithms, do_output, step_output_frequency = simulation.setup_simulation(False, my_parallel)
   my_system.num_steps = num_steps_per_cycle * num_cycles

   # Setup the timers 
   metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, 0, 0, do_output, step_output_frequency, True, False)

   # Create two lj derived types for 1 and 2 as copies of the runtime (reads cutoff too)
   lj1 = metalwalls.mw_lennard_jones.copy_type(my_system.lj) 
   lj2 = metalwalls.mw_lennard_jones.copy_type(my_system.lj) 
   lj_diff = metalwalls.mw_lennard_jones.copy_type(my_system.lj) 
   lj_lc = metalwalls.mw_lennard_jones.copy_type(my_system.lj) 
   
   # Fill lj1 and lj2
   for typeA in range(my_system.lj.num_species): # Fortran indexes
      for typeB in range(typeA, my_system.lj.num_species):
         metalwalls.mw_lennard_jones.set_parameters(lj1, typeA+1, typeB+1, eps1[typeA, typeB], sigma1[typeA, typeB])
         metalwalls.mw_lennard_jones.set_parameters(lj2, typeA+1, typeB+1, eps2[typeA, typeB], sigma2[typeA, typeB])
   
   # Set the system's lj as the linear combination of A=4*eps*sigma**12 and B*eps*sigma**6
   for typeA in range(my_system.lj.num_species): # Python indexes
      for typeB in range(my_system.lj.num_species):
         lj_lc.a[typeA, typeB] = (1-lambd) * lj1.a[typeA, typeB] + lambd * lj2.a[typeA, typeB]
         lj_lc.b[typeA, typeB] = (1-lambd) * lj1.b[typeA, typeB] + lambd * lj2.b[typeA, typeB]
         lj_diff.a[typeA, typeB] = lj2.a[typeA, typeB] - lj1.a[typeA, typeB]
         lj_diff.b[typeA, typeB] = lj2.b[typeA, typeB] - lj1.b[typeA, typeB]
   
   my_system.lj = lj_lc

   energy_derivative = np.zeros(num_cycles)
   for cycle in range(num_cycles):
      # Compute energy derivative dU/dlambd
      # Compute energy difference
      metalwalls.mw_lookup_table.lut_set_all(my_system.energy, 0.0)
      my_system.lj = lj_diff
      metalwalls.mw_system.energy_lj(my_system)
      energy_derivative[cycle] = metalwalls.mw_lookup_table.lut_get(my_system.energy, "Lennard-Jones")
      energy_derivative[cycle] += metalwalls.mw_lookup_table.lut_get(my_system.energy, "Lennard-Jones_cutoff")
      
      # Run MD steps
      my_system.lj = lj_lc
      metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, num_steps_per_cycle * (cycle + 1), num_steps_per_cycle, do_output, step_output_frequency, False, False)
   
   # Stop timers, output diagnostics
   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)

   mean, std = statistical_analysis(energy_derivative)

   return mean, std


# Run a series of TI simulations with different lambda values given by lambd_list
# Model1 with eps1, sigma1
# Model2 with eps2, sigma2
def ti_lj_lj(eps1, sigma1, eps2, sigma2, lambd_list, num_steps_per_cycle, num_cycles, output_file='energy_derivative.out'):
   # Define constants
   hartree_si = 4.359744650e-18   # hartree energy in J
   avogadro_si = 6.022140857e+23  # (1/mol)
   hartree2kjpermol = hartree_si * avogadro_si * 1.0e-3

   # Initialize mpi
   my_parallel, comm = simulation.initialize_parallel()
   fcomm = comm.py2f()
   metalwalls.mw_tools.start_parallel_python(my_parallel, fcomm)

   energy_derivative = np.zeros((np.size(np.array(lambd_list)), 3))
   for i, lambd in enumerate(lambd_list):
      energy_derivative[i, 0] = lambd
      energy_derivative[i, 1], energy_derivative[i, 2] = ti_lj_lj_lambda(eps1, sigma1, eps2, sigma2, lambd, num_steps_per_cycle, num_cycles, my_parallel)   

      if my_parallel.comm_rank == 0:
         print('Lambda {:4.2f}: energy_derivative {} hartree; {} kJ/mol'.format(lambd, energy_derivative[i, 1], energy_derivative[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()
   
   if my_parallel.comm_rank == 0:
      energy_derivative[:, 1:] *= hartree2kjpermol
      np.savetxt(output_file, energy_derivative, header='# Lambda   dU/dlambda (kJ/mol)   std(dU/dlambda) (kJ/mol)\n')

   metalwalls.mw_tools.finalize_parallel()
   


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

# Run a simulation using a hybrid hamiltonian at a given lambd
# Model1 with eps1, sigma1
# Model2 with eps2, sigma2
# Atoms of ion_type are described with a soft-core lj potential
def ti_soft_core_lambda(eps1, sigma1, eps2, sigma2, ion_type, p, alpha, lambd, num_steps_per_cycle, num_cycles, my_parallel):
   # Setup the simulation
   my_system, my_parallel, my_algorithms, do_output, step_output_frequency = simulation.setup_simulation(False, my_parallel)
   my_system.num_steps = num_steps_per_cycle * num_cycles

   # Setup the timers 
   metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, 0, 0, do_output, step_output_frequency, True, False)

   # Create two lj derived types for 1 and 2 as copies of the runtime (reads cutoff too)
   lj1 = metalwalls.mw_lennard_jones.copy_type(my_system.lj)
   lj2 = metalwalls.mw_lennard_jones.copy_type(my_system.lj)
   
   # Fill lj1 and lj2
   for typeA in range(my_system.lj.num_species): # Fortran indexes
      for typeB in range(typeA, my_system.lj.num_species):
         metalwalls.mw_lennard_jones.set_parameters(lj1, typeA+1, typeB+1, eps1[typeA, typeB], sigma1[typeA, typeB])
         metalwalls.mw_lennard_jones.set_parameters(lj2, typeA+1, typeB+1, eps2[typeA, typeB], sigma2[typeA, typeB])

   # Put shift to 0 for all interactions without ion_type, only the ion_type-* interactions should be soft-core
   # lj.shift contains lj.sigma**6
   # modify to have lj.shift1 = alpha * lj.sigma1**6 * lambda**p
   # modify to have lj.shift2 = alpha * lj.sigma2**6 * (1-lambda)**p
   for typeA in range(my_system.lj.num_species):
      for typeB in range(my_system.lj.num_species):
         if (typeA == ion_type-1 or typeB == ion_type-1) and not(typeA == ion_type-1 and typeB == ion_type-1):
            lj1.shift[typeA, typeB] *= alpha * lambd**p
            lj2.shift[typeA, typeB] *= alpha * (1 - lambd)**p
            lj1.alpha[typeA, typeB] = alpha
            lj2.alpha[typeA, typeB] = alpha
         elif not(typeA == ion_type-1 or typeB == ion_type-1):
            lj1.shift[typeB, typeB] = 0.0
            lj2.shift[typeA, typeB] = 0.0
            lj1.alpha[typeA, typeB] = 0.0
            lj2.alpha[typeA, typeB] = 0.0
   lj1.shift[ion_type-1, ion_type-1] *= alpha * lambd**p
   lj2.shift[ion_type-1, ion_type-1] *= alpha * (1 - lambd)**p

   energy_derivative = np.zeros(num_cycles)
   for cycle in range(num_cycles):
      # Compute energy derivative dU/dlambd
      # Compute energy with model1
      h1 = np.zeros(4)
      my_system.lj = lj1
      metalwalls.mw_system.energy_lj_soft_core(my_system, h1)
      # Compute energy with model1
      h2 = np.zeros(4)
      my_system.lj = lj2
      metalwalls.mw_system.energy_lj_soft_core(my_system, h2)
      # Compute energy derivative dU/dlambd = h2-h1
      part2 = h2[0] - alpha * p / 6 * lambd * (1 - lambd)**(p-1) * h2[1] + h2[2] + h2[3]
      part1 = h1[0] - alpha * p / 6 * (1 - lambd) * lambd**(p-1) * h1[1] + h1[2] + h1[3]
      energy_derivative[cycle] = part2 - part1
      
      # Run MD steps
      for step in range(num_steps_per_cycle):
         # First step of MD step: propagate positions, compute charges...
         metalwalls.mw_tools.step_setup(my_system, my_algorithms, my_parallel, cycle * num_steps_per_cycle + step + 1)
         # Second step of MD step: compute forces
         #Compute coulomb forces
         my_system.coulomb_forces_ions[:,:,:] = 0.0
         metalwalls.mw_system.forces_compute_coulomb(my_system)
         # Compute intramolecular forces
         my_system.intramolecular_forces_ions[:,:,:] = 0.0
         metalwalls.mw_system.forces_compute_intramolecular(my_system)
         # Compute intermolecular forces
         # Compute with model V
         my_system.intermolecular_forces_ions[:,:,:] = 0.0
         f1 = np.zeros(np.shape(my_system.intermolecular_forces_ions))
         my_system.lj = lj1
         metalwalls.mw_system.forces_compute_lj_soft_core(my_system)
         f1[:, :, :] = my_system.intermolecular_forces_ions[:, :, :]
         # Compute with model Na
         my_system.intermolecular_forces_ions[:,:,:] = 0.0
         f2 = np.zeros(np.shape(my_system.intermolecular_forces_ions))
         my_system.lj = lj2
         metalwalls.mw_system.forces_compute_lj_soft_core(my_system)
         f2[:, :, :] = my_system.intermolecular_forces_ions[:, :, :]
         # Compute linear combination of forces
         my_system.intermolecular_forces_ions[:, :, :] = (1 - lambd) * f1[:, :, :] + lambd * f2[:, :, :]
         # Add eventual Fumi-Tosi or DAIM contributions
         metalwalls.mw_system.forces_compute_ft(my_system)
         metalwalls.mw_system.forces_compute_xft(my_system)
         metalwalls.mw_system.forces_compute_daim(my_system)
         # Gather all forces
         metalwalls.mw_system.forces_gather(my_system)

         # Third step of MD step: update velocities, output info
         metalwalls.mw_tools.step_update_output(my_system, my_algorithms, cycle * num_steps_per_cycle + step + 1, do_output, step_output_frequency)
   
   # Stop timers, output diagnostics
   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)

   mean, std = statistical_analysis(energy_derivative)

   return mean, std


# Run a series of TI simulations with different lambda values given by lambd_list
# Model1 with eps1, sigma1
# Model2 with eps2, sigma2
# Atoms of ion_type are described with a soft-core lj potential
def ti_soft_core(eps1, sigma1, eps2, sigma2, ion_type, p, alpha, lambd_list, num_steps_per_cycle, num_cycles, output_file='energy_derivative.out'):
   # Define constants
   hartree_si = 4.359744650e-18   # hartree energy in J
   avogadro_si = 6.022140857e+23  # (1/mol)
   hartree2kjpermol = hartree_si * avogadro_si * 1.0e-3

   # Initialize mpi
   my_parallel, comm = simulation.initialize_parallel()
   fcomm = comm.py2f()
   metalwalls.mw_tools.start_parallel_python(my_parallel, fcomm)

   energy_derivative = np.zeros((np.size(np.array(lambd_list)), 3))
   for i, lambd in enumerate(lambd_list):
      energy_derivative[i, 0] = lambd
      energy_derivative[i, 1], energy_derivative[i, 2] = ti_soft_core_lambda(eps1, sigma1, eps2, sigma2, ion_type, p, alpha, lambd, num_steps_per_cycle, num_cycles, my_parallel)

      if my_parallel.comm_rank == 0:
         print('Lambda {:4.2f}: energy_derivative {} hartree; {} kJ/mol'.format(lambd, energy_derivative[i, 1], energy_derivative[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()

   if my_parallel.comm_rank == 0:
      energy_derivative[:, 1:] *= hartree2kjpermol
      np.savetxt(output_file, energy_derivative, header='# Lambda   dU/dlambda (kJ/mol)   std(dU/dlambda) (kJ/mol)\n')

   metalwalls.mw_tools.finalize_parallel()



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

def ti_plot(inpt_file='energy_derivative.out', figure_file='ti.png'):
   import matplotlib.pyplot as plt
   energy_derivative = np.loadtxt(inpt_file)
   ti = np.trapz(energy_derivative[:, 1], x=energy_derivative[:, 0])
   plt.figure(figsize=(20,10))
   plt.xlabel(r'$\lambda$')
   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_derivative[:, 0], energy_derivative[:, 1], yerr=energy_derivative[:, 2])
   plt.savefig(figure_file, bbox_inches='tight')


def statistical_analysis(energy_derivative):
   # Statistical analysis
   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_derivative)
   num_cycles = np.shape(energy_derivative)[0]
   interval = int(num_cycles / 5)
   for block in range(5):
      mean_student[block] = np.mean(energy_derivative[interval * block:interval * (block + 1) - 1])
   mean = np.mean(mean_student)
   std = np.std(mean_student) * student

   return mean, std
