import py21cmfast as p21c
import numpy as np
import time

FileName = 'MCG_Template.h5'
LC_Quantities = ('brightness_temp','Ts_box','xH_box','Tk_box','Trad_box')
GLB_Quantities = ('brightness_temp','Ts_box','xH_box','dNrec_box','z_re_box','Gamma12_box','J_21_LW_box','density','Trad_box','Tk_box','Fcoll')

def Get_EDGES_Data():
    data = np.genfromtxt("/home/dm/gaolq/cjs/Radio_Excess_EDGES/data/Data_EDGES.txt")
    freq = data[:, 0]
    wght = data[:, 1]
    tsky = data[:, 2]
    freq = freq[wght>0]
    tsky = tsky[wght>0]
    redshifts = 1420/freq - 1
    length = len(freq)
    EDGES_Data = {'freq':freq, 'tsky':tsky, 'redshifts':redshifts, 'length':length}
    return EDGES_Data

EDGES_Data = Get_EDGES_Data()
freq = EDGES_Data['freq']
observed_redshifts = 1420/freq - 1,
redshift = np.min(observed_redshifts) - 0.1
max_redshift=np.max(observed_redshifts) + 0.1,

user_params = p21c.UserParams(
  HII_DIM = 128,
  N_THREADS = 15,
  USE_RELATIVE_VELOCITIES = True,
  USE_INTERPOLATION_TABLES = True,
  FAST_FCOLL_TABLES = True,
  BOX_LEN = 250,
  )

astro_params = p21c.AstroParams(
  # -- Fix these params --
  F_STAR10 = -1.5,
  ALPHA_STAR = 0.5,
  ALPHA_STAR_MINI = 0.0,
  F_ESC10 = -1.2,
  ALPHA_ESC = -0.3,
  L_X = 40.5,
  NU_X_THRESH = 200.0,
  aR_mini = 0.62,
  )

flag_options = p21c.FlagOptions(
  USE_RADIO_ACG = False,
  USE_RADIO_MCG = True,
  USE_MINI_HALOS = True,
  USE_MASS_DEPENDENT_ZETA = True,
  INHOMO_RECO = True,
  USE_TS_FLUCT = True,
  )


# ---- Initialise ----
start_time = time.time()
lc = p21c.run_lightcone(
  redshift=redshift, 
  max_redshift=max_redshift,
  astro_params=astro_params, 
  flag_options=flag_options,
  user_params = user_params,
  lightcone_quantities=LC_Quantities,
  global_quantities=GLB_Quantities
  )
end_time = time.time()
print('--------RunTime--------')
print(end_time - start_time)

try:
  lc.save(FileName)
except:
  # Remove pre-existing file
  os.remove(FileName)
  lc.save(FileName)
