import numpy as np
import metalwalls
import simulation

# Increase (or decrease) the temperature using a 'stairs' function

#  Temperature
#       ^
#       |   
# temp2 |            ---
#       |         ---
#       |      --- 
#       |   ---
# temp1 |---
#       ----------------------------->
#       0             num_steps       Time

def run_temp_ramp(temp1, temp2, ct=1, num_steps=-1):
   # Setup the simulation
   my_system, my_parallel, my_algorithms, do_output, step_output_frequency = simulation.setup_simulation()

   # Read optional parameter or the number of steps provided in the runtime
   if num_steps == -1:
      num_steps = my_system.num_steps
   else:
      my_system.num_steps = num_steps

   # Setup the timers 
   metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, 0, 0, do_output, step_output_frequency, True, False)
   
   dT = (temp2 - temp1) / my_system.num_steps * ct
   nct = my_system.num_steps // ct

   print('Changing temperature by {} every {} step'.format(dT, ct))

   # Run steps
   for istep in range(nct):
      # Change target temperature
      try:
         my_system.barostat.temperature = temp1 + istep * dT
      except:
         my_system.thermostat.temperature = temp1 + istep * dT
         
      step = 1 + istep * ct 
      metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, step, ct, do_output, step_output_frequency, False, False)
   
   # Stop timers, output diagnostics
   metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, num_steps, 0, do_output, step_output_frequency, False, True)
   
   simulation.finalize_simulation(my_system, my_parallel, my_algorithms)

