# Import ASE to setup the automatic calculator
import ase
from ase.calculators.espresso import Espresso

# Lets import cellconstructor
import cellconstructor as CC
import cellconstructor.Phonons

# Import the sscha
import sscha, sscha.Ensemble, sscha.Relax, sscha.SchaMinimizer

from numpy import *
import numpy as np
import matplotlib.pyplot as plt

plt.ion() #Setup interactive plot mode

pseudo = {"H":"H.pbe-rrkjus_psl.1.0.0.UPF",
         "S" : "S.pbe-nl-rrkjus_psl.1.0.0.UPF"}

input_data = {"ecutwfc" : 35,
              "ecutrho" : 350,
              "occupations" : "smearing",
              "input_dft" : "blyp",
              "mixing_beta" : 0.2,
              "conv_thr" : 1e-9,
              "degauss" : 0.02,
              "smearing" : "mp",
              "pseudo_dir" : ".",
             "tprnfor" : True}


espresso_calc = Espresso(pseudopotentials =pseudo,
                        input_data = input_data,
                        kspacing = 0.06,
                        koffset = (1,1,1))

# Let us load the starting dynamical matrix
dyn = CC.Phonons.Phonons("matdyn", nqirr = 3)

# Apply the sum rule and delete the imaginary modes
dyn.Symmetrize()
dyn.ForcePositiveDefinite()

# Generate the ensemble and the minimizer objects
ensemble = sscha.Ensemble.Ensemble(dyn, 1000, supercell = dyn.GetSupercell())
minimizer = sscha.SchaMinimizer.SSCHA_Minimizer(ensemble)

# We setup all the minimization parameters
minimizer.min_step_dyn = 0.001
minimizer.kong_liu_ratio = 0.5

# Here we prepare a cluster
# Here we configure the cluster object MARCONI
import sscha.Cluster
my_hpc = sscha.Cluster.Cluster(pwd = None)

# We setup the connection info
my_hpc.hostname = "ekhi" # The command to connect via ssh to the cluster
#my_hpc.account_name = "IscrB_COMRED" # The name of the project for the computation
my_hpc.workdir = "/scratch/lorenzo/my_h3s_dir" # the directory in which the calculations are performed

# Now we need to setup the espresso
# First we must tell the cluster where to find him:
my_hpc.binary = "pw.x -npool NPOOL -i  PREFIX.pwi > PREFIX.pwo"
# Then we need to specify if some modules must be loaded in the submission script
my_hpc.load_modules = """
# Here this is a bash script at the beginning of the submission
# We can load modules

module load QuantumESPRESSO
export OMP_NUM_THREADS=1
"""

# All these information are independent from the calculation
# Now we need some more specific info, like the number of processors, pools and other stuff
my_hpc.n_cpu = 40 # We will use 32 processors
my_hpc.n_nodes = 1 #In 1 node
my_hpc.n_pool = 20 # This is an espresso specific tool, the parallel CPU are divided in 4 pools

# We can also choose in how many batch of jobs we want to submit simultaneously, and how many configurations for each job
my_hpc.batch_size = 10
my_hpc.job_number = 20
# In this way we submit 10 jobs, each one with 10 configurations (overall 100 configuration at time)

# We give 25 seconds of timeout
my_hpc.set_timeout(30)

# We can specify the time limit for each job,
my_hpc.time = "00:20:00" # 5 minutes

my_hpc.setup_workdir()

# We prepare the Automatic relaxation
relax = sscha.Relax.SSCHA(minimizer, ase_calculator = espresso_calc,
                         N_configs = 200,
                         max_pop = 6,
                         save_ensemble = True,
                         cluster = my_hpc)

relax.relax(ensemble_loc = "directory_of_the_ensemble")

relax.relax()

np.savetxt("all_frequencies.dat", all_frequencies)

# We reset the frequency array
all_frequencies = []

# We define a function that will be called by the code after each minimization step, 
# passing to us the minimizer at that point
# We will store the frequencies, to plot after the minimization their evolution.
def add_current_frequencies(minimizer):
    # Get the frequencies
    w, p = minimizer.dyn.DiagonalizeSupercell()
    
    all_frequencies.append(w)
    
    # In this way the file will be updated at each step
    np.savetxt("all_frequencies.dat", all_frequencies)

# We add this function to the relax ojbect
relax.setup_custom_functions(custom_function_post = add_current_frequencies)

# Now we are ready to
# ***************** RUN *****************
relax.relax(ensemble_loc = "data_ensemble_autorelax")
# ***************************************

# We can save the frequency list for furture analisys
# NOTE: This will overwrite the one given in the examples!
np.savetxt("all_frequencies.dat", all_frequencies)

# Lets plot the minimization
relax.minim.plot_results()

# Lets plot all the frequencies in their evolution
all_frequencies = np.loadtxt("all_frequencies.dat")
freqs = np.array(all_frequencies).T 
n_freqs = shape(freqs)[0]

plt.figure(dpi = 100)
plt.title("Frequencies")
plt.ylabel("Energy [cm-1]")
plt.xlabel("Steps")
for i in range(n_freqs):
    plt.plot(freqs[i, :] * CC.Units.RY_TO_CM)

# Compute the free energy hessian
free_energy_hessian = relax.minim.ensemble.get_free_energy_hessian(include_v4 = False)

# Now we can save the free energy hessian like it was a dynamical matrix
free_energy_hessian.save_qe("hessian")

# We can print its eigenvalues to see if there are imaginary phonons
w, p = free_energy_hessian.DiagonalizeSupercell()
print(w * CC.Units.RY_TO_CM)