import os
import subprocess
import math
import numpy as np

# input values
directory_name = 'Calc_ScatLen_NaCs' # the name of the directory where the values will be stored
delta_over_omega = 1.0
omega_start = 3.5 # MHz
omega_stop = 3.5 # MHz
n_steps = 2
spacing = 'linear' # choose log or linear
middle_range = (0.5,2) # the middle range that should have a higher density of data points
n_middle = 1 # how many more data points should be in the middle compared to the rest; set to 1 for a uniform distribution

# constants and conversions
hartree_in_inv_cm = 2.1947463136314e5
atomic_mass_constant = 1.66053906892e-27
electron_mass = 9.1093837139e-31
elementary_charge_in_SI = 1.6021766340e-19
bohr_in_SI = 0.529177210544e-10
hz_in_inv_cm = 3.335640951e-11
debye_to_SI = hz_in_inv_cm/1e19
debye_to_au = (debye_to_SI/elementary_charge_in_SI)/bohr_in_SI
hartree_in_inv_cm = 2.1947463136314e5
hartree_in_SI = 4.3597447222060e-18
planck_const_in_SI = 6.62607015e-34
pi_value = 3.141592653589793
hbar_in_SI = planck_const_in_SI/2.0/pi_value


# create directory for outputs and move into it
if not os.path.exists(directory_name):
    os.mkdir(directory_name)
os.chdir(directory_name)

# create file
def create_file(
        basename:str,i:int,rmin:float|int,rmid:float|int,rmax:float|int,
        drs:float|int,drl:float|int,c6_0:float,c6_2:float,c6_4:float,c3:float,
        ured:float,omega:float
        ) -> None:
    """
    Creates an input file and a slurm file in a separate directory and runs them in this directory.
    All files can be found in the subdirectory labelled dir{i}.
    The value for r_inv_sqrt is included in the LABEL parameter in the &INPUT block.

    Args:
        basename (string)   : the basename used for the input and output files.
        i (int)             : the number of this job
        a (float)           : the coefficient for the dipole-dipole interaction in cm-1
        rmid (int or float) : the value for RMID
        rmax (int or float) : the value for RMAX
        drs (int or float)  : the value for DRS
        drl (int or float)  : the value for DRL
        emin (float)        : the value for EMIN, the minimum energy value where we are searching,
                                relative to the bound state (i.e. bound state = 0)
        emax (float)        : the value for EMAX, the maximum energy value where we are searching
        freq_harm (float)   : the harmonic frequency of the harmonic potential

    Returns:
        None : creates and runs files
    """
    input_text = f"""&INPUT
      LABEL  = 'omega = {omega}',
      IPRINT = 20,
      URED   = {'{0:.15f}'.format(ured)}D0,
      RMIN   = {rmin}D0,
      RMID   = {rmid}D0,
      RMAX   = {rmax}.0D0,
      BCYOMN = 1.D8,
      DRS    = {drs}D0,
      DRL    = {drl}D0,
      IRMSET = 0,
      IPROPS = 6,
      IPROPL = 9,
      TOLHI  = 0.0,
      POWRL = 2.0D0,
      ISIGPR = 1,
      ISCRU =  0,
      EUNITS = 2, ! units = K
      JTOTL  = 0,
      JTOTU  = 0,
      IBFIX = 1,
      NNRG = 1,
      ENERGY = 1.128115D-11, ! Ecoll=10E_3
      &END

    &BASIS
      ITYPE  = 21,
      IBOUND = 0,
      BCT    = .TRUE.,
      JMIN   = 0,
      JMAX   = 10,
      JSTEP  = 2,
      JZCSMX = -4,
      ROTI   = 1.D-32,
      BE     = 1.D-32,
      &END

    &POTL
      RM     = 0.529177210903D0,
      EPSIL  = 1.0,
      MXLAM  = 4,
      LAMBDA = 0,2,4,2,
      NTERM  = 1,1,1,1,
      A      = {'{0:.20f}'.format(c6_0)}D0,{'{0:.20f}'.format(c6_2)}D0,{'{0:.20f}'.format(c6_4)}D0,{'{0:.20f}'.format(c3)}D0,
      NPOWER = -6,-6,-6,-3,
      &END
    """

    subdir_name = f"dir{i}"
    if not os.path.exists(subdir_name):
        os.mkdir(subdir_name)
    #os.chdir(subdir_name)

    with open(f"{subdir_name}/{basename}.input","w") as f:
        f.write(input_text)


#omega_in_au = omega * planck_const_in_SI*1e6/hartree_in_SI #MHz -> Hatree
m_amu = 77.94761061999999 # reduced mass in amu
m = m_amu * (atomic_mass_constant/electron_mass) # reduced mass in atomic units
mu_debye = 4.75 # dipole moment in debye
mu = mu_debye*debye_to_au # dipole moment in a.u.
mu_SI = mu_debye * debye_to_SI
#d = mu / math.sqrt(12*(1+(delta_over_omega)**2)) # effective dipole in atomic units
#c3_in_au = 2*d*d # coefficient of dipole moment in atomic units
c3_in_au = 2*mu*mu/(12*(1+delta_over_omega**2))
#c6_in_au = mu**4 / (8*omega_in_au*(1+(delta_over_omega)**2)**(3/2))
c3 = c3_in_au*hartree_in_inv_cm # c3 in cm-1 (times bohr radius^3)
#c6 = c6_in_au * hartree_in_inv_cm # c6 in cm-1 (times bohr radius^6)
#c6_0 = c6*4/5
#c6_2 = c6*(-4/7)
#c6_4 = c6*(-8/35)

# E (in J) = 1/2 * hbar * freq (in Hz)
# E (in MHz) = E_in_J * 10^-6 / h
# --> E (in MHz) = 1/2 * 1/2*pi * h(in MHz)
#emax = 1/(2*2*pi_value) * freq_harm # emax in MHz (equals the energy of separated particles)

drs = 0.125
drl = 0.25
rmin = 50
rmax = 1000000
rmid = 10000


#stepsize = (omega_stop-omega_start)/(n_steps-1)
if spacing == 'linear':
    step_val = np.linspace(omega_start,omega_stop,n_steps)
elif spacing == 'log':
    if n_middle == 1:
        step_val = np.logspace(np.log10(omega_start),np.log10(omega_stop),n_steps)
    else:
        m1,m2 = middle_range
        total_n = 2+n_middle
        sample1 = int(n_steps/total_n)
        sample3 = sample1
        sample2 = int(n_steps*n_middle/total_n)
        # make sure the number of samples still sum upt to the correct value, i.e. no samples lost to roundoffs
        sample2 += n_steps - sample1 - sample2 - sample3 # add any lost samples to the midpoint
        step_val1 = np.logspace(np.log10(omega_start),np.log10(m1),sample1,endpoint=False)
        step_val2 = np.logspace(np.log10(m1),np.log10(m2),sample2,endpoint=False)
        step_val3 = np.logspace(np.log10(m2),np.log10(omega_stop),sample3,endpoint=True)
        step_val = np.concatenate([step_val1,step_val2,step_val3])

else:
    raise ValueError("spacing did not match 'linear' or 'log'.")

for i in range(n_steps):
    #omega = omega_start + i*stepsize
    omega = step_val[i]
    omega_in_au = omega * planck_const_in_SI*1e6/hartree_in_SI #MHz -> Hatree
    c6_in_au = mu**4 / (8*omega_in_au*(1+(delta_over_omega)**2)**(3/2))
    c6 = c6_in_au * hartree_in_inv_cm # c6 in cm-1 (times bohr radius^6)
    c6_0 = c6*4/5
    c6_2 = c6*(-4/7)
    c6_4 = c6*(-8/35)

    create_file(
        basename='molscat_NaCs',
        i=i,
        rmin=rmin,
        rmid=rmid,
        rmax=rmax,
        drs=drs,
        drl=drl,
        c6_0=c6_0,
        c6_2=c6_2,
        c6_4=c6_4,
        c3=c3,
        ured=m_amu,
        omega=omega
        )

