from multiprocessing import Pool
import EnvEq as ee
import numpy as np
import itertools as it
import os

#parsing input into numpy arrays
from input import *
y0=np.array([y0_Tpos,y0_Tpro,y0_Tneg,y0_o2,y0_test])
p=np.array([p_o2,p_test])
mu=np.array([[mu_o2Tpos,mu_o2Tpro,mu_o2Tneg],[mu_testTpos,mu_testTpro,0]])
lam=np.array([lam_o2,lam_test])
t_D=np.array([t_DTpos,t_DTpro,t_DTneg])
r=np.array([r_Tpos,r_Tpro,r_Tneg])
delta=np.array([delta_Tpos,delta_Tpro,delta_Tneg])
rho=np.array([rho_Tpos,rho_Tpro,rho_Tneg])
lim=np.array([[[l_lim_o2Tpos,u_lim_o2Tpos],[l_lim_o2Tpro,u_lim_o2Tpro],[l_lim_o2Tneg,u_lim_o2Tneg]],[[l_lim_testTpos,u_lim_testTpos],[l_lim_testTpro,u_lim_testTpro],[0,0]]],dtype=np.float64)

#make directories for saving raw_outputs
try:
    os.makedirs("../../raw_output/EnvEq/"+f_name)
except:
    pass

#iterator over these 
o2_lim_arr=np.empty([0,2])
for ulim_Tpro in np.arange(0.1,1,0.2):
    for ulim_Tneg in np.arange(0.1,1,0.2):
        o2_lim_arr=np.append(o2_lim_arr,[[ulim_Tpro,ulim_Tneg]],axis=0)


def solve_parm(u_lim_o2): #calls the solve_eq function with all default inputs other than o2_lim
    f_name_i=f_name+"{:.1f}".format(u_lim_o2[0])+"-"+"{:.1f}".format(u_lim_o2[1])
    lim[0,1,1]=u_lim_o2[0]
    lim[0,2,1]=u_lim_o2[1]
    ee.solve_eq(t_max,dt,y0,p,mu,lam,r,K,delta,rho,lim,f_name_i)

if __name__ == '__main__':
    pool = Pool(4)
    pool.map(solve_parm,o2_lim_arr) #iterate over the o2_lims
    pool.close()
    pool.join()

[(9, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (10, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (11, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (12, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (13, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (14, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (15, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (16, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (17, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (26, 'numpy.empty', 'np.empty', True, 'import numpy as np\n'), (27, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (21, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (28, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (36, 'EnvEq.solve_eq', 'ee.solve_eq', True, 'import EnvEq as ee\n'), (39, 'multiprocessing.Pool', 'Pool', False, 'from multiprocessing import Pool\n'), (29, 'numpy.append', 'np.append', True, 'import numpy as np\n')]