sim_low.py

import configparser
import subprocess
import json
import logging
import os
import sys
import glob

import astropy
from astropy.io import fits
from astropy.io.fits import Header
from astropy.time import Time, TimeDelta
from astropy.table import Table
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
import astropy.units as u
from ARatmospy import ArScreens

#import healpy as hp

import colorednoise as cn
import h5py

#from healpy.newvisufunc import projview, newprojplot
from reproject import reproject_from_healpix
from scipy.interpolate import interp1d
from scipy.stats import describe
from scipy import constants

from pygdsm import GlobalSkyModel16 #A
import numpy as np
from struct import pack,unpack
from numpy.random import default_rng
import math
import numpy.matlib
#import oskar #A
import multiprocessing as mp

LOG = logging.getLogger()

handler = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
LOG.addHandler(handler)
LOG.setLevel(logging.INFO)

result_list = []
def log_result(result):
    result_list.append(result)
    return
    
def get_start_time(ra0_deg, length_sec):
    """Returns optimal start time for field RA and observation length."""
    t = Time('2021-09-21 00:00:00', scale='utc', location=('116.764d', '0d'))
    dt_hours = 24.0 - t.sidereal_time('apparent').hour + (ra0_deg / 15.0)
    start = t + TimeDelta(dt_hours * 3600.0 - length_sec / 2.0, format='sec')
    return start.value

def get_dt_days(ra0_deg):
    """Returns optimal centre time for field RA."""
    t = Time('2021-09-21 00:00:00', scale='utc', location=('116.764d', '0d'))
    dt_days = (24.0 - t.sidereal_time('apparent').hour + (ra0_deg / 15.0))/24.
    return dt_days

def bright_sources():
    """Returns a list of bright A-team sources."""
    # Sgr A: Guesstimate from Fred Dulwich, should check
    # For A: data from the Molonglo Southern 4 Jy sample (VizieR).
    # Others from GLEAM reference paper, Hurley-Walker et al. (2017), Table 2.
    return numpy.array((
        [266.41683, -29.00781,  2000,0,0,0, 178e6, 0.0,0,3600,3600,0], # Sgr A
        [ 50.67375, -37.20833,   528,0,0,0, 178e6, -0.51, 0, 0, 0, 0],  # For
        [201.36667, -43.01917,  1370,0,0,0, 200e6, -0.50, 0, 0, 0, 0],  # Cen
        [139.52500, -12.09556,   280,0,0,0, 200e6, -0.96, 0, 0, 0, 0],  # Hyd
        [ 79.95833, -45.77889,   390,0,0,0, 200e6, -0.99, 0, 0, 0, 0],  # Pic
        [252.78333,   4.99250,   377,0,0,0, 200e6, -1.07, 0, 0, 0, 0],  # Her
        [187.70417,  12.39111,   861,0,0,0, 200e6, -0.86, 0, 0, 0, 0],  # Vir
        [ 83.63333,  22.01444,  1340,0,0,0, 200e6, -0.22, 0, 0, 0, 0],  # Tau
        [299.86667,  40.73389,  7920,0,0,0, 200e6, -0.78, 0, 0, 0, 0],  # Cyg
        [350.86667,  58.81167, 11900,0,0,0, 200e6, -0.41, 0, 0, 0, 0]   # Cas
        ))

def make_sky_model(sky0, sky1, fgd_att, fsl_att, settings, radius_deg, flux_min_inner_jy, flux_min_outer_jy):
    """Filter sky models. 
    sky0 has full flux density, sky1 has attenuated flux density
    fsl_att is outer sky attenuation factor
    fgd_att is extra attenuation factor applied to all sky models
    Includes all sources within the given radius, and sources above the
    specified flux outside this radius.
    """
    # Get pointing centre.
    ra0_deg = float(settings['observation/phase_centre_ra_deg'])
    dec0_deg = float(settings['observation/phase_centre_dec_deg'])

    # Create "inner" and "outer" sky models.
    sky_inner = sky0.create_copy()
    sky_outer = sky1.create_copy()
    sky_inner.filter_by_radius(0.0, radius_deg, ra0_deg, dec0_deg)
    sky_inner.filter_by_flux((fgd_att*flux_min_inner_jy), 1e9)
    sky_outer.filter_by_radius(radius_deg, 180.0, ra0_deg, dec0_deg)
    sky_outer.filter_by_flux((fgd_att*fsl_att*flux_min_outer_jy), 1e9)
    LOG.info("Number of sources in sky0: %d", sky0.num_sources)
    LOG.info("Number of sources in inner compact sky model: %d", sky_inner.num_sources)
    LOG.info("Number of sources in outer compact sky model above %.3f Jy: %d",
             flux_min_outer_jy, sky_outer.num_sources)
    sky_outer.append(sky_inner)
    LOG.info("Number of sources in output compact sky model: %d", sky_outer.num_sources)
    return sky_outer

def crop_frac(img,fracx,fracy):
    z,y,x = img.shape
    cropx = int(x*fracx)
    cropy = int(y*fracy)
    startx = x//2-(cropx//2)
    starty = y//2-(cropy//2)    
    return img[:,starty:starty+cropy,startx:startx+cropx]
                       
def casa_proc(itext,sifrq):
    subprocess.run(itext, shell='True')
    otext = 'Finished uvfits for channel: '+sifrq
    return otext

def subp_proc(itext,sifrq):
    subprocess.run(itext, shell='True')
    otext = 'Finished subprocess run for channel: '+sifrq
    return otext

def oskar_proc(settings_dict,ifrq,out_root,do_TRECS,gxypix_asec,txypix_asec,num_gpu,rseed,sky_c_name):
    sifrq = str(ifrq).zfill(len(str(num_freq))+1)
    gsm = out_root+'_GSM'
    tsm = out_root+'_TSM'
    gsmxxx = out_root+'_GSMXXX'
    # Load base oskar settings
    settings = oskar.SettingsTree('oskar_sim_interferometer')
    settings.from_dict(settings_dict)
    # Start with fresh compact source sky model
    sky = oskar.Sky()
    sky_c = oskar.Sky.load(sky_c_name)
    sky.append(sky_c)
    # define some names
    gsm_frq = gsm+'_IFRQ_'+sifrq
    gsm_frqf = gsm+'_IFRQ_'+sifrq+'.fits'
    tsm_frq = tsm+'_IFRQ_'+sifrq
    tsm_frqf = tsm+'_IFRQ_'+sifrq+'.fits'
    out_uvfrq = out_root+'_IFRQ_'+sifrq+'.uv'
    out_msfrq = out_root+'_IFRQ_'+sifrq+'.ms'
    out_mswfrq = out_root+'_IFRQ_'+sifrq+'.wsc'
    out_msnfrq = out_root+'_IFRQ_'+sifrq+'.wsn'
    out_uvffrq = out_root+'_IFRQ_'+sifrq+'.uvf'
    # Extract the relevant slice from the low res diffuse sky model
    text = miriadbin+'imsub in='+gsmxxx+' out='+gsm_frq+' region="images('+str(ifrq+1)+')"'
    subprocess.run(text, shell='True')
    text = miriadbin+'fits op=xyout in='+gsm_frq+' out='+gsm_frqf
    subprocess.run(text, shell='True')
    # Use a large negative value for min_abs_val to make sure all negative values are used
    sky_diff0 = oskar.Sky.from_fits_file(gsm_frqf, min_abs_val=-1.e6)
    sky_diff0_array = sky_diff0.to_array()
    # set Ref Freq and spectral index columns of diffuse model to be correct
    sky_diff0_array[:,6] = freqs[ifrq]
    sky_diff0_array[:,7] = 0.    
    # Need to make sure that the component lists are fully sampled so
    # set BMAJ and BMIN columns of model to have FWHM=2*PIXEL
    # Amplitude may need to be re-normalised so that flux density is preserved
    sky_diff0_array[:,9] = 2. * gxypix_asec
    sky_diff0_array[:,10] = 2. * gxypix_asec
    sky_diff0_vec = np.ones(12)
    # It seems that OSKAR uses the Jy/pixel units correctly to preserve flux density, so no need to re-normalise
    sky_diff0_vec[2] = 1.
    sky_diffl = oskar.Sky.from_array(sky_diff0_array*sky_diff0_vec)
    sky.append(sky_diffl)
    LOG.info("Number of comps in low res diffuse sky model: %d", sky_diffl.num_sources)
    if (do_TRECS):
        # Extract the relevant slice from the TRECS sky model
        text = miriadbin+'imsub in='+tsm+' out='+tsm_frq+' region="images('+str(ifrq+1)+')"'
        subprocess.run(text, shell='True')
        text = miriadbin+'fits op=xyout in='+tsm_frq+' out='+tsm_frqf
        subprocess.run(text, shell='True')
        # Use a small positive value to restrict model to significant sources
        sky_diff1 = oskar.Sky.from_fits_file(tsm_frqf, min_abs_val=trecs_cut)
        sky_diff1_array = sky_diff1.to_array()
       # set Ref Freq and spectral index columns of bright T-RECS model to be correct
        sky_diff1_array[:,6] = freqs[ifrq]
        sky_diff1_array[:,7] = 0.    
         # set BMAJ and BMIN columns of model to have FWHM=2*PIXEL
        sky_diff1_array[:,9] = 2. * txypix_asec
        sky_diff1_array[:,10] = 2. * txypix_asec
        sky_diff1_vec = np.ones(12)
        # It seems that OSKAR uses the Jy/pixel units correctly to preserve flux density, so no need to re-normalise
        sky_diff1_vec[2] = 1.
        sky_diffh = oskar.Sky.from_array(sky_diff1_array*sky_diff1_vec)
        sky.append(sky_diffh)
        LOG.info("Number of comps in high res diffuse sky model: %d", sky_diffh.num_sources)

    LOG.info("Number of comps in total sky model: %d", sky.num_sources)
    settings['observation/start_frequency_hz'] = str(freqs[ifrq])
    settings['interferometer/ms_filename'] = out_msfrq
    gpu_id = ifrq - int(ifrq/num_gpu)*num_gpu
    settings['simulator/cuda_device_ids'] = str(gpu_id)
    LOG.info("Simulating %s", settings['interferometer/ms_filename'])
    # Make sure that the random number seeds are modifed for each channel
    settings['interferometer/noise/seed'] = str(int(rseed+ifrq*(3-5**0.5)*180.))
    tel = oskar.Telescope(settings=settings)
    sim = oskar.Interferometer(settings=settings)
    sim.set_sky_model(sky)
    sim.set_telescope_model(tel)
    sim.run()
    otext = 'Finished oskar sim for channel: '+sifrq
    return otext

# Note that this version is hard-wired to be a continuous track from the perspective of the ionospheric model,
# while for other purposes it is still broken down into a sequence of short cuts


# ANNA starts here
# Check if the filename argument is provided
if len(sys.argv) > 1:
    filename = sys.argv[1]  # First command-line argument (index 1)
    print(f"The filename is: {filename}")
else:
    print("No filename provided.")


config = configparser.ConfigParser()

config.read(filename)

#n_cores = int(config.getfloat("pipeline", "n_cores"))
#doplot = config.getboolean("pipeline", "doplot")

base_dir = config.get("pipeline", "base_dir")
out_dir = base_dir+'Outputs/'
out_root = config.get("pipeline", "out_root")
num_gpu = int(config.getfloat("pipeline","num_gpu"))
num_cpu = int(config.getfloat("pipeline","num_cpu"))
num_wsclean = int(config.getfloat("pipeline","num_wsclean"))

# Field Centre
ra0_deg = config.getfloat("field","field_ra") # in degrees
dec0_deg = config.getfloat("field","field_dec") # in degrees

# Visibility sampling in time
length_sec = config.getfloat("observation","telescope_time")*3600. # total length of observation
int_sec = config.getfloat("observation","int_sec") # integration time in seconds
cut_sec = length_sec # length of each cut in seconds
num_cuts = round(length_sec/cut_sec)  # number of cuts

freq_fov = config.getfloat("observation","freq_fov") # to set fsl and mhd scales
freq_min = config.getfloat("observation","freq_min") # in Hz
freq_max = config.getfloat("observation","freq_max") # in Hz
freq_inc = config.getfloat("observation","freq_inc") # in Hz
chan_inc = config.getfloat("observation","chan_inc") # channel BW for smearing calculation in Hz


gsm_freq_min = config.getfloat("observation","gsm_freq_min")
gsm_freq_max = config.getfloat("observation","gsm_freq_max")

# Image plane (FFT plus degridded) sky model spatial extent and sampling
# Base model consists of GSM2016 
# Options for MHD simulation of fine scale filaments, T-RECS faint sky, EoR signal
# Should cover out to first null of station beam at nu_min
null_mult = config.getfloat("pipeline","null_mult") # field diameter in units of first null diam at nu_min for Airy pattern of 38m station
gnxy = int(config.getfloat("pipeline","gnxy"))  # pixels in GSM, MHD and EoR sky model
do_GSM = config.getboolean("pipeline","do_GSM")
do_MHD = config.getboolean("pipeline","do_MHD")
do_ATeam = config.getboolean("pipeline","do_ATeam")
do_GLEAM = config.getboolean("pipeline","do_GLEAM")
do_TRECS = config.getboolean("pipeline","do_TRECS")
# Image and Compact foreground multiplier factor prior to adding in EoR
fgd_att = config.getfloat("pipeline","fgd_att")
do_EOR = config.getboolean("pipeline","do_EOR")
do_SIM = config.getboolean("pipeline","do_SIM")# do or skip the main simulation loop
do_restart_SIM = config.getboolean("pipeline","do_restart_SIM") # pickup from interrupt
num_sim_res = 0
do_restart_wsc = config.getboolean("pipeline","do_restart_wsc") # pickup from interrupt
num_wsc_res = 0
toss_models = config.getboolean("pipeline","toss_models")
do_uvf = config.getboolean("pipeline","do_uvf")


# Final output cube parameters
ouvtap = config.getfloat("output","ouvtap")  # 60. # Gaussian taper in arcsec (in conjuction with uniform weighting)
ocell = config.getfloat("output","ocell")  # 16. # cellsize in arcsec
opix = int(config.getfloat("output","opix"))  # 2048  # pixels
oniter = config.getfloat("output","oniter")  # 1.e6 # clean iterations nom=1e6
do_wsc = config.getboolean("output","do_wsc")#False # should we generate "pretty" clean images
do_wsn = config.getboolean("output","do_wsn")#True # should we generate natural images


# T-RECS faint sky (<100 mJy) model cube
trecs_cube0 = base_dir+config.get("input","trecs_cube")  #base_dir+'SkyModels/sky_continuum_sdc3_v4_1'
trecs_cut = config.getfloat("input","trecs_cut")  #1.e-5 # filter the tapered trecs model image at this brightness per pixel

# EoR signal cube to be used
#eor_dir = config.get("input","eor_dir")  #base_dir+'EoRModels/'
eor_cube = base_dir+config.get("input","eor_cube")  #eor_dir+'m33/obsall'

# Compact (DFT) sky model contains A-Team plus GLEAM with filtering as below
cross_flux = config.getfloat("input","cross_flux")  #0.1 # lower flux cutoff in Jy for inner compact sky model
flux_val = config.getfloat("input","flux_val")  #5. # lower flux cutoff in Jy for outer compact sky model 

# De-mixing error model
# Far sidelobe attenuation factor to compensate for use of AC (rather than N=2 CC) station beam
# and residual de-mixing/subtraction error 
fsl_att = config.getfloat("calibration","fsl_att")  #1.e-3

# Visibility DI (direction independent) calibration error model
phas_err = config.getfloat("calibration","phas_err")  #0.2 # resdidual phase error in degrees
amp_err = config.getfloat("calibration","amp_err")  #0.002 # residual amp error
time_psd_exp = config.getfloat("calibration","time_psd_exp")  #2 # power law exponent of time error PSD
bp_phas_err = config.getfloat("calibration","bp_phas_err")  #0.2 # resdidual band pass phase error in degrees
bp_amp_err = config.getfloat("calibration","bp_amp_err")  #0.002 # residual band pass amp error
freq_psd_exp = config.getfloat("calibration","freq_psd_exp")  #2 # power law exponent of freq error PSD

# Visibility DD (direction dependent) calibration error model
do_DD =  config.getboolean("ionosphere","do_DD")
ion_res = config.getfloat("ionosphere","ion_res")  #0.1 # residual fraction of direction dependent phase error
# Ionospheric model parameters for two assumed layers
screen_width_metres = config.getfloat("ionosphere","screen_width")*1.e3  #200e3 # Total screen size (200 km)
bmax = config.getfloat("ionosphere","bmax")*1.e3  # 20 km sub-aperture size
sampling = config.getfloat("ionosphere","sampling")  #100.0  # 100 m/pixel
r0_1 = config.getfloat("ionosphere","r0_1")*1.e3  #7e3  # Scale size (7 km)
speed1 = config.getfloat("ionosphere","speed1")*1.e3/3600  # 150 km/h in m/s
angle1 = config.getfloat("ionosphere","angle1")  #35. # angle in degrees
height1 = config.getfloat("ionosphere","height1")*1.e3  #290.e3 # height in m
r0_2 = config.getfloat("ionosphere","r0_2")*1.e3  #7e3  # Scale size (7 km)
speed2 = config.getfloat("ionosphere","speed2")*1.e3/3600  # 75 km/h in m/s
angle2 = config.getfloat("ionosphere","angle2")  #-69. # angle in degrees
height2 = config.getfloat("ionosphere","height2")*1.e3  # height in m
alpha_mag = config.getfloat("ionosphere","alpha_mag")  #0.999  # Evolve screen slowly

# Thermal Noise model
eff_tau = config.getfloat("noise","eff_tau")#1.e3 # effective integration time of observation in hours (for thermal noise)
ant_fact = config.getfloat("noise","ant_fact")#1. # scale up effective station number
pol_fact = config.getfloat("noise","pol_fact")#2. # scale up polarisation number
# Random number seed for thermal noise calculation and DI errors
rseed = int(config.getfloat("noise","rseed"))#9452136

# Thermal noise in Jy of each station defined in these files
freq_file = base_dir+config.get("models","freq_file")#'TelModels/noise_frequencies.txt'
nrms_file = base_dir+config.get("models","nrms_file")#'TelModels/rms_req.txt'

# Where to find sky and telescope models
sm_dir = base_dir+config.get("models","sm_dir")#+'SkyModels/'
# MHD simulation below has been extrapolated to cover 50.2 to 352.6 MHz in steps of 5.4 MHz using a 3d order poly in log(nu) and log(S)
tdfile = sm_dir+config.get("models","tdfile")#+'GS_1500_X_sw3.fits'
tel_dir = base_dir+config.get("models","tel_dir")#+'TelModels/V1/telescope.tm/'
num_tel = int(config.getfloat("models","num_tel"))#512


# Where to find key executables 
oskarbin = config.get("system","oskarbin")#'/usr/local/bin/'
casabin = config.get("system","casabin")#'/home/r.braun/casa/casa-6.4.4-31-py3.8/bin/'
miriadbin = config.get("system","miriadbin")#'/usr/local/miriad_test_new/miriad/linux64/bin/'
wsbin = config.get("system","wsbin")#'/usr/bin/'
cpbin = config.get("system","cpbin")#'/bin/'


print('reading done')
print(ion_res)
exit()



#ANNA ends here

cur_dir = os.getcwd()

os.chdir(out_dir)

# Start by establishing image sky model FoV
airy_rad_deg = 180.*(np.arcsin(1.12*3.e8/freq_fov/38.))/np.pi
fov_deg = null_mult*2.*airy_rad_deg
ra0_rad = ra0_deg*np.pi/180.
dec0_rad = dec0_deg*np.pi/180.

# The OSKAR thermal noise calculation wants the RMS visibility noise
# The RMS input file has the station SEFD so we need to rescale this
# We need to take account of at least channel bandwidth and integration time
# On top of that we might wish to take account of fewer stations and/or
# single versus dual polarisations, since we are now set-up for only generating XX.
# Finally we need to scale up to some effective total integration time
dbdt_fact = 2.*chan_inc*int_sec
etau_fact = (eff_tau * 3600.) / (cut_sec * num_cuts)
# read the nominal station sensitivities and write the tuned sensitivities to rms_file
rms_file = '_rms_file.txt'
sefd = np.loadtxt(nrms_file)
rms = sefd / (ant_fact * dbdt_fact**0.5 * etau_fact**0.5 * pol_fact**0.5)
np.savetxt(rms_file, rms)
ra0_hr = ra0_deg/15.
fovrad_deg = fov_deg/2.
ha_min = -(length_sec/3600.)/2.
ha_inc = (length_sec/3600.)/(num_cuts-1+1.e-12)
freqs = np.single(np.arange(freq_min,freq_max+freq_inc,freq_inc))
num_freq = int((freq_max-freq_min)/freq_inc+1)
freq_ghz = (freq_min+freq_max)/2./1.e9
freq_bw_mhz = freq_inc*num_freq/1.e6
gxypix_deg = fov_deg/gnxy
gxypix_asec = fov_deg/gnxy*3600.
gxyrefpix = gnxy/2
txypix_asec = 5.

num_int = round(cut_sec/int_sec)

if do_DD:
    # Some secondary ionospheric model parameters
    ionof_total = out_root+'_IONO_Track.fits'
    iono_total = out_root+'_IONO_Track'
    num_times = int(length_sec/int_sec)  # 
    m = int(bmax / sampling)  # Pixels per sub-aperture (200).
    n = int(screen_width_metres / bmax)  # Sub-apertures across the screen (10).
    num_pix = n * m
    rate = 1./int_sec  # The inverse frame rate (6 per minute).
    pscale = screen_width_metres / (n * m)  # Pixel scale (100 m/pixel).
    # Parameters for each layer.
    # (scale size [m], speed [m/s], direction [deg], layer height [m]).
    layer_params = np.array([(r0_1, speed1, angle1, height1),
                            (r0_2, speed2, angle2, height2)])

    my_screens = ArScreens.ArScreens(n, m, pscale, rate, layer_params, alpha_mag)
    # Ionosphere model creation
    my_screens.run(num_times)
    # Convert to TEC
    # phase = image[pixel] * -8.44797245e9 / frequency
    frequency = 1e8
    phase2tec = -frequency / 8.44797245e9    
    w = WCS(naxis=4)
    w.naxis = 4
    w.wcs.cdelt = [pscale, pscale, 1.0 / rate, 1.0]
    w.wcs.crpix = [num_pix // 2 + 1, num_pix // 2 + 1, num_times // 2 + 1, 1.0]
    w.wcs.ctype = ['XX', 'YY', 'TIME', 'FREQ']
    w.wcs.crval = [0.0, 0.0, 0.0, frequency]
    data = np.zeros([1, num_times, num_pix, num_pix])
    for layer in range(len(my_screens.screens)):
        for i, screen in enumerate(my_screens.screens[layer]):
            data[:, i, ...] += phase2tec * screen[np.newaxis, ...]

        
    fits.writeto(filename=ionof_total, data=data, header=w.to_header(), overwrite=True)
    text = miriadbin+'fits op=xyin in='+ionof_total+' out='+iono_total
    subprocess.run(text, shell='True')

    # Could use a standard 4 hour track ionospheric model to save time 
    # iono_total = 'Std4h_IONO_Track'
    iono_total_att = out_root+'_IONO_att'
    iono_total_attf = out_root+'_IONO_att.fits'

# determine orientation of output image relative to Galactic plane
ddec = 0.5
dec1_deg = dec0_deg + ddec
skypos0 = SkyCoord(ra=ra0_deg*u.degree, dec=dec0_deg*u.degree)
skypos1 = SkyCoord(ra=ra0_deg*u.degree, dec=dec1_deg*u.degree)
l0 = skypos0.galactic.l.degree
b0 = skypos0.galactic.b.degree
l1 = skypos1.galactic.l.degree
b1 = skypos1.galactic.b.degree
phi = np.arctan2((l1-l0),(b1-b0))

#timebase = '21SEP21.'
#timedt = str(int(1.e8*get_dt_days(ra0_deg)))
#ttime = timebase+timedt

# Define base diffuse sky model names 
gsm = out_root+'_GSM'
tsm = out_root+'_TSM'
gsms = out_root+'_GSMs'
gsmr = out_root+'_GSMr'
gsmx = out_root+'_GSMX'
gsmxx = out_root+'_GSMXX'
gsmxxx = out_root+'_GSMXXX'
gsmxxxf = out_root+'_GSMXXX.fits'
gsmf = out_root+'_GSM.fits'

# Define high resolution diffuse sky proxy model
gtd = out_root+'_GTD'
gtdf = out_root+'_GTD.fits'

# Define compact sky model name
cmpt = out_root+'_CMPT'

# Start by making a diffuse sky model  

# Basic image sky model is GSM2016
gsmfile = 'gsm_2016_cube.fits'
# model resolution is 56 arcmin FWHM
bmaj = 0.9333
bmin = 0.9333
bpa = 0.

# Set up header for low res cube
outhdr3d = Header({'Simple': True})
outhdr3d.set('NAXIS',3)
outhdr3d.set('NAXIS1',gnxy)
outhdr3d.set('NAXIS2',gnxy)
outhdr3d.set('NAXIS3',num_freq)
outhdr3d.set('CTYPE1','RA---SIN')
outhdr3d.set('CTYPE2','DEC--SIN')
outhdr3d.set('CTYPE3','FREQ    ')
outhdr3d.set('CRVAL1',ra0_deg)
outhdr3d.set('CRVAL2',dec0_deg)
outhdr3d.set('CRVAL3',freq_min)
outhdr3d.set('CRPIX1',gxyrefpix)
outhdr3d.set('CRPIX2',gxyrefpix)
outhdr3d.set('CRPIX3',1)
outhdr3d.set('CDELT1',-gxypix_deg)
outhdr3d.set('CDELT2',gxypix_deg)
outhdr3d.set('CDELT3',freq_inc)
outhdr3d.set('CUNIT1','deg     ')
outhdr3d.set('CUNIT2','deg     ')
outhdr3d.set('CUNIT3','Hz      ')
outhdr3d.set('BMAJ',bmaj)
outhdr3d.set('BMIN',bmin)
outhdr3d.set('BPA',bpa)
outhdr3d.set('BUNIT','KELVIN')

# Set up header for single low res frequency first
outhdr2d = Header({'Simple': True})
outhdr2d.set('BITPIX',-32)
outhdr2d.set('NAXIS',2)
outhdr2d.set('NAXIS1',gnxy)
outhdr2d.set('NAXIS2',gnxy)
outhdr2d.set('CTYPE1','RA---SIN')
outhdr2d.set('CTYPE2','DEC--SIN')
outhdr2d.set('CRVAL1',ra0_deg)
outhdr2d.set('CRVAL2',dec0_deg)
outhdr2d.set('CRPIX1',gxyrefpix)
outhdr2d.set('CRPIX2',gxyrefpix)
outhdr2d.set('CDELT1',-gxypix_deg)
outhdr2d.set('CDELT2',gxypix_deg)
outhdr2d.set('CUNIT1','deg     ')
outhdr2d.set('CUNIT2','deg     ')

# Get GSM2016 model

gsm_2016 = GlobalSkyModel16(freq_unit='Hz',data_unit='TRJ',resolution='hi',include_cmb=True)
#gsm_2016 = GlobalSkyModel16(freq_unit='Hz',data_unit='TRJ',resolution='hi')
gsmcube = np.zeros((num_freq, gnxy, gnxy),dtype=np.float32)
text = '/bin/rm -r '+gsmfile
subprocess.run(text, shell='True')

# change units from Kelvin to Jy/pixel at the same time
waves2 = (constants.c/freqs)**2
pix_sr = (gxypix_deg*constants.pi/180.)**2
KtoJyppix = 2.*constants.Boltzmann*pix_sr/waves2/1.e-26

if (do_GSM):
    for ichan in range(0,num_freq,1):
        gsm_2016.generate(freqs[ichan],gsm_freq_min,gsm_freq_max)
        # gsm_2016.generate(freqs[ichan])
        gsm_2016.write_fits(gsmfile)
        hdugsm = fits.open(gsmfile)
        hdugsm[1].header.set('COORDSYS','G')
        gsmcube[ichan,:,:], footprint = np.single(reproject_from_healpix(hdugsm[1],outhdr2d))*KtoJyppix[ichan]
        text = '/bin/rm -r '+gsmfile
        subprocess.run(text, shell='True')

sumgsmofnu = np.nansum(gsmcube,axis=(1, 2))
# write output fits cubes
gsmhdu = fits.PrimaryHDU(gsmcube)
gsmhdu.header.set('CTYPE1','RA---SIN')
gsmhdu.header.set('CTYPE2','DEC--SIN')
gsmhdu.header.set('CTYPE3','FREQ    ')
gsmhdu.header.set('CRVAL1',ra0_deg)
gsmhdu.header.set('CRVAL2',dec0_deg)
gsmhdu.header.set('CRVAL3',freq_min)
gsmhdu.header.set('CRPIX1',gxyrefpix)
gsmhdu.header.set('CRPIX2',gxyrefpix)
gsmhdu.header.set('CRPIX3',1)
gsmhdu.header.set('CDELT1',-gxypix_deg)
gsmhdu.header.set('CDELT2',gxypix_deg)
gsmhdu.header.set('CDELT3',freq_inc)
gsmhdu.header.set('CUNIT1','deg     ')
gsmhdu.header.set('CUNIT2','deg     ')
gsmhdu.header.set('CUNIT3','Hz      ')
gsmhdu.header.set('BMAJ',bmaj)
gsmhdu.header.set('BMIN',bmin)
gsmhdu.header.set('BPA',bpa)
gsmhdu.header.set('BUNIT','JY/PIXEL')
gsmhdu.writeto(gsmf,overwrite=True)
del gsmcube
text = '/bin/rm -r '+gsm+' '+gsmr+' '+gsms+' '+gtd+' '+gsmx+' '+gsmxx+' '+gsmxxx+' '+gsmxxxf
subprocess.run(text, shell='True')
text = miriadbin+'fits op="xyin" in='+gsmf+' out='+gsm
subprocess.run(text, shell='True')

if (do_MHD):
    # Extract simulated foreground filaments image
    # Base cube is 512x512 spatial pixels and covers 115 to 169 MHz in steps of 5.4 MHz
    # Will use regrid to extract the relevant frequency samples while fudging the spatial header
    # to make it appear identical in areal coverage to what is needed.
    # Actually the pixel size needs to be a little bigger, say by 1.05, so that we don't end up with blanked pixels at the spatial edges
    dra_rad = -fov_deg*np.pi/180/512.*1.05
    ddec_rad = fov_deg*np.pi/180/512.*1.05
    tdrot = '_Gsynchrot'
    tdrotf = '_Gsynchrot.fits'
    tdregrid = '_Gsynch'
    tdregridf = '_Gsynch.fits'
    text = '/bin/rm -r '+tdregrid+' '+tdregridf+' '+tdrot+' '+tdrotf
    subprocess.run(text, shell='True')
    # get local copy of filament cube
    hduim = fits.open(tdfile)
    imdat = np.squeeze(hduim[0].data)
    imhdr = hduim[0].header
    # This is an attempt to get a filament orientation perpendicular to the Galaxy
    phip = phi
    if (phi >= -3*np.pi/4 and phi < -np.pi/4):
        imdat = np.rot90(imdat,k=-1,axes=(1,2))
        phip = phi + np.pi/2
        
    if (phi >= np.pi/4 and phi < 3.*np.pi/4):
        imdat = np.rot90(imdat,k=1,axes=(1,2))
        phip = phi - np.pi/2

    if (phi < -3.*np.pi/4):
        imdat = np.rot90(imdat,k=2,axes=(1,2))
        phip = phi + np.pi

    if (phi >= 3.*np.pi/4):
        imdat = np.rot90(imdat,k=2,axes=(1,2))
        phip = phi - np.pi

    bord0 = int(imdat.shape[1]*4/16)
    bord1 = int(imdat.shape[1]*4/16)
    # pad first with reflected input image
    impad0 = np.pad(imdat,((0,0),(bord0,bord0),(bord0,bord0)),mode='reflect')
    # pad some more with zeros
    impad = np.pad(impad0,((0,0),(bord1,bord1),(bord1,bord1)))
    xlen = impad.shape[2]
    ylen = impad.shape[1]
    zlen = impad.shape[0]
    ufreq = np.fft.fftfreq(xlen)
    gx = np.zeros((zlen,ylen,xlen))
    gyx = np.zeros((zlen,ylen,xlen))
    gxyx = np.zeros((zlen,ylen,xlen))
    theta = phip
    a = np.tan(theta/2.)
    b = -np.sin(theta)
    # do FFT-shear-based rotation following Larkin et al. 1997 (only for rotations of -45 to +45 degree)
    # for larger rotations need to combine with a suitable transpose as done above
    for iplan in range(0,zlen,1):
        for icol in range(0,xlen,1):
            gx[iplan,icol,:] = np.real(np.fft.ifft(np.fft.fft(impad[iplan,icol,:])*np.exp(-2.j*np.pi*(icol-xlen/2)*ufreq*a)))

        for icol in range(0,xlen,1):
            gyx[iplan,:,icol] = np.real(np.fft.ifft(np.fft.fft(gx[iplan,:,icol])*np.exp(-2.j*np.pi*(icol-xlen/2)*ufreq*b)))

        for icol in range(0,xlen,1):
            gxyx[iplan,icol,:] = np.real(np.fft.ifft(np.fft.fft(gyx[iplan,icol,:])*np.exp(-2.j*np.pi*(icol-xlen/2)*ufreq*a)))
    

    gout = crop_frac(gxyx,0.5,0.5)
    hduimrot = fits.PrimaryHDU(gout,imhdr)
    hduimrot.writeto(tdrotf,overwrite=True)
    del gx
    del gyx
    del gxyx
    del gout
    text = miriadbin+'fits op=xyin in='+tdrotf+' out='+tdrot
    subprocess.run(text, shell='True')
    text = miriadbin+'puthd in="'+tdrot+'/crval1" value='+str(ra0_rad) 
    subprocess.run(text, shell='True')
    text = miriadbin+'puthd in="'+tdrot+'/crval2" value='+str(dec0_rad) 
    subprocess.run(text, shell='True')
    text = miriadbin+'puthd in="'+tdrot+'/cdelt1" value='+str(dra_rad) 
    subprocess.run(text, shell='True')
    text = miriadbin+'puthd in="'+tdrot+'/cdelt2" value='+str(ddec_rad) 
    subprocess.run(text, shell='True')
    text = miriadbin+'regrid in='+tdrot+' tin='+gsm+' tol=0.002 out='+tdregrid 
    subprocess.run(text, shell='True')
    text = miriadbin+'fits op=xyout in='+tdregrid+' out='+tdregridf 
    subprocess.run(text, shell='True')
    tdhdu = fits.open(tdregridf)
    tdcube = tdhdu[0].data
    sumtdofnu = np.nansum(tdcube,axis=(1, 2))
    if (do_GSM):
        tdcube =  (sumgsmofnu / sumtdofnu)[:, None, None] * tdcube
    # Resolution of simulation defined by original pixel sampling
    tdbmaj = fov_deg/512.
    tdbmin = fov_deg/512.
    tdbpa = 0.
    tdhdu = fits.PrimaryHDU(tdcube)
    tdhdu.header.set('CTYPE1','RA---SIN')
    tdhdu.header.set('CTYPE2','DEC--SIN')
    tdhdu.header.set('CTYPE3','FREQ    ')
    tdhdu.header.set('CRVAL1',ra0_deg)
    tdhdu.header.set('CRVAL2',dec0_deg)
    tdhdu.header.set('CRVAL3',freq_min)
    tdhdu.header.set('CRPIX1',gxyrefpix)
    tdhdu.header.set('CRPIX2',gxyrefpix)
    tdhdu.header.set('CRPIX3',1)
    tdhdu.header.set('CDELT1',-gxypix_deg)
    tdhdu.header.set('CDELT2',gxypix_deg)
    tdhdu.header.set('CDELT3',freq_inc)
    tdhdu.header.set('CUNIT1','deg     ')
    tdhdu.header.set('CUNIT2','deg     ')
    tdhdu.header.set('CUNIT3','Hz      ')
    tdhdu.header.set('BMAJ',tdbmaj)
    tdhdu.header.set('BMIN',tdbmin)
    tdhdu.header.set('BPA',tdbpa)
    tdhdu.header.set('BUNIT','JY/PIXEL')
    tdhdu.writeto(gtdf,overwrite=True)
    del tdcube
    text = miriadbin+'fits op="xyin" in='+gtdf+' out='+gtd
    subprocess.run(text, shell='True')
    xfactor = (bmaj/tdbmaj)**2
    # merge FFTs of high res simulation (normalised to same integral) and low res GSM, factor is unity given Jy/pixel units 
    text = miriadbin+'immerge in="'+gtd+','+gsm+'" options="notaper" factor=1. out='+gsmx
    subprocess.run(text, shell='True')
else:
    gsmx = gsm
    # End of (do_MHD) condition

if (do_TRECS):
    # Now add in the TRECS faint sky model 
    # Start by making the bright and faint subsets defined by trecs_split
    trecs_split = 2.e-4
    trecs_cube = '_trecs_cube'
    trecs_mask = '_trecs_mask'
    trecsb_cube = '_trecs_bright'
    trecsf_cube = '_trecs_faint'    
    trecsb_cubepr = '_trecs_brightpr'
    trecsf_cubepr = '_trecs_faintpr'
    trecsf_cubeprc67 = '_trecs_faintprc67'
    text = '/bin/rm -r '+trecs_cube+' '+trecs_mask+' '+trecsb_cube+' '+trecsb_cubepr+' '+trecsf_cube+' '+trecsf_cubepr+' '+trecsf_cubeprc67
    subprocess.run(text, shell='True')
    # get local copy of trecs cube
    text = cpbin+'cp -r '+trecs_cube0+' '+trecs_cube
    subprocess.run(text, shell='True')
    text = miriadbin+'imsub in='+trecs_cube+' region="images(1)" out='+trecs_mask
    subprocess.run(text, shell='True')
    text = miriadbin+'puthd in='+trecs_mask+'/crval3 value=0.08'
    subprocess.run(text, shell='True')
    text = miriadbin+'puthd in='+trecs_mask+'/cdelt3 value=0.005'
    subprocess.run(text, shell='True')
    text = miriadbin+'maths exp="<'+trecs_cube+'>" mask="<'+trecs_mask+'>.gt.'+str(trecs_split)+'" options="grow" out='+trecsb_cubepr
    subprocess.run(text, shell='True')
    text = miriadbin+'imblr in='+trecsb_cubepr+' out='+trecsb_cube
    subprocess.run(text, shell='True')    
    text = miriadbin+'puthd in='+trecsb_cube+'/crval3 value=0.08'
    subprocess.run(text, shell='True')
    text = miriadbin+'puthd in='+trecsb_cube+'/cdelt3 value=0.005'
    subprocess.run(text, shell='True')
    text = miriadbin+'maths exp="<'+trecs_cube+'>" mask="<'+trecs_mask+'>.le.'+str(trecs_split)+'" options="grow" out='+trecsf_cubepr
    subprocess.run(text, shell='True')
    text = miriadbin+'imblr in='+trecsf_cubepr+' out='+trecsf_cubeprc67
    subprocess.run(text, shell='True')    
    text = miriadbin+'puthd in='+trecsf_cubeprc67+'/crval3 value=0.08'
    subprocess.run(text, shell='True')
    text = miriadbin+'puthd in='+trecsf_cubeprc67+'/cdelt3 value=0.005'
    subprocess.run(text, shell='True')
    text = miriadbin+'convol map='+trecsf_cubeprc67+' fwhm=67 scale=4.91785e-3 out='+trecsf_cube
    subprocess.run(text, shell='True')    
    text = miriadbin+'puthd in='+trecsf_cube+'/bmaj value=3.24825e-4'
    subprocess.run(text, shell='True')
    text = miriadbin+'puthd in='+trecsf_cube+'/bmin value=3.24825e-4'
    subprocess.run(text, shell='True')    
    trecsb_reg = '_TRecsb_regrid'
    trecsf_reg = '_TRecsf_regrid'
    trecsf_regr = '_TRecsf_regridr'
    trecsb_norm = '_TRecsb_norm'
    trecsf_norm = '_TRecsf_norm'
    text = '/bin/rm -r '+trecsb_reg+' '+trecsf_reg+' '+trecsf_regr+' '+trecsb_norm+' '+trecsf_norm
    subprocess.run(text, shell='True')
    # Redefine the TRECS header to have the requested central coordinates
    text = miriadbin+'puthd in='+trecsb_cube+'/crval1 value='+str(ra0_rad)
    subprocess.run(text, shell='True')
    text = miriadbin+'puthd in='+trecsb_cube+'/crval2 value='+str(dec0_rad)
    subprocess.run(text, shell='True')
    # Regrid bright in frequency only while preserving xy
    mir_desc = str(freq_min/1.e9)+','+str(1)+','+str(freq_inc/1.e9)+','+str(int(num_freq))
    text = miriadbin+'regrid in='+trecsb_cube+' tin='+gsm+' axes="3" tol=0.002 out='+trecsb_reg
    subprocess.run(text, shell='True')
    txypix_asec = 5.
    # Make sure that units are Jy/pixel by defining unorm appropriately, while applying any fgd attenuation
    unorm = 1. * fgd_att
    text = miriadbin+'maths exp="<'+trecsb_reg+'>*'+str(unorm)+'" out='+trecsb_norm 
    subprocess.run(text,shell='True')
    # replace any blanks with zero
    text = miriadbin+'imblr in="'+trecsb_norm+'" out='+tsm 
    subprocess.run(text, shell='True')
    # Now deal with the faint TRECS sky
    # Regrid to GSM cube sampling
    text = miriadbin+'regrid in='+trecsf_cube+' tin='+gsm+' tol=0.002 out='+trecsf_reg 
    subprocess.run(text, shell='True')
    text = miriadbin+'imblr in='+trecsf_reg+' out='+trecsf_regr 
    subprocess.run(text, shell='True')
    # now combine with rest of diffuse sky model while correcting for Jy/pixel with new pixel size
    unorm = (gxypix_asec/txypix_asec)**2
    text = miriadbin+'maths exp="<'+gsmx+'>+(<'+trecsf_regr+'>*'+str(unorm)+')" out='+gsmxx 
    subprocess.run(text, shell='True')
    
else:
    gsmxx = gsmx
    # End of (do_TRECS) condition
    
if (do_EOR):
    # Finally add in the EoR signal model
    # first copy across and resample to low res grid, current units are K, sky area may or may not be adequate, pad with zero as needed
    eor_reg = '_EoR_regrid'
    eor_norm = '_EoR_norm'
    eor_normr = '_EoR_normr'
    text = '/bin/rm -r '+eor_reg+' '+eor_norm+' '+eor_normr
    subprocess.run(text, shell='True')
    text = miriadbin+'regrid in='+eor_cube+' tin='+gsm+' tol=0.002 out='+eor_reg 
    subprocess.run(text, shell='True')
    # change units from K to Jy/pixel which varies as frequency**2 (assume freq increases with channel number)
    waves2 = (constants.c/freqs)**2
    pix_sr = (gxypix_deg*constants.pi/180.)**2
    KtoJyppix = 2.*constants.Boltzmann*pix_sr/waves2/1.e-26
    zmin = (np.amin(KtoJyppix))**0.5
    zmax = (np.amax(KtoJyppix))**0.5
    text = miriadbin+'maths exp="<'+eor_reg+'>*z*z" zrange="'+str(zmin)+','+str(zmax)+'" out='+eor_norm 
    subprocess.run(text, shell='True')
    # replace any blanks with zero
    text = miriadbin+'imblr in="'+eor_norm+'" out='+eor_normr 
    subprocess.run(text, shell='True')
    # Now add in EoR while attenuating the image foreground model (if fgd_att not = 1)
    text = miriadbin+'maths exp="'+str(fgd_att)+'*<'+gsmxx+'>+<'+eor_normr+'>" out='+gsmxxx 
    subprocess.run(text, shell='True')
    text = miriadbin+'fits op="xyout" in='+gsmxxx+' out='+gsmxxxf 
    subprocess.run(text, shell='True')

else:
    text = miriadbin+'maths exp="'+str(fgd_att)+'*<'+gsmxx+'>" out='+gsmxxx 
    subprocess.run(text, shell='True')
    text = miriadbin+'fits op="xyout" in='+gsmxxx+' out='+gsmxxxf 
    subprocess.run(text, shell='True')

# Now set up for the OSKAR sky model 

# Load base telescope model.
settings = oskar.SettingsTree('oskar_sim_interferometer')

settings['simulator/double_precision'] = 'true'
settings['simulator/use_gpus'] = 'true'
settings['simulator/max_sources_per_chunk'] = '600000'
settings['simulator/cuda_device_ids'] = '0'

settings['observation/phase_centre_ra_deg'] = str(ra0_deg)
settings['observation/phase_centre_dec_deg'] = str(dec0_deg)
settings['observation/num_channels'] = str(int(1))
settings['observation/frequency_inc_hz'] = str(freq_inc)
settings['observation/length'] = str(cut_sec)
settings['observation/num_time_steps'] = str(int(num_int))
settings['observation/start_time_utc'] = get_start_time(ra0_deg, length_sec)

settings['telescope/input_directory'] = tel_dir
settings['telescope/allow_station_beam_duplication'] = 'true'
settings['telescope/pol_mode'] = 'Scalar'

if (do_DD):
    # Add in ionospheric screen model
    settings['telescope/ionosphere_screen_type'] = 'External'
else:
    settings['telescope/ionosphere_screen_type'] = 'None'

    
settings['interferometer/channel_bandwidth_hz'] = str(chan_inc)
settings['interferometer/time_average_sec'] = str(int_sec)
settings['interferometer/ignore_w_components'] = 'false'

# Add in Telescope noise model via files where rms has been tuned
settings['interferometer/noise/enable'] = 'true'
settings['interferometer/noise/seed'] = str(int(rseed))
settings['interferometer/noise/freq'] = 'Data'
settings['interferometer/noise/freq/file'] = freq_file
settings['interferometer/noise/rms'] = 'Data'
settings['interferometer/noise/rms/file'] = rms_file

# Attenuate the bright sky list to reflect both AC versus CC side-lobe differences as well as de-mixing accuracy
fsl_att_vec = np.ones(12)
fsl_att_vec[2] = fsl_att
# Also introduce possibility for a foreground attenuation to simulate subtraction success
fgd_att_vec = np.ones(12)
fgd_att_vec[2] = fgd_att

# Load base sky model for both actual and attenuated skies
sky0 = oskar.Sky()
sky1 = oskar.Sky()
sky_c = oskar.Sky()
if (do_ATeam):
    sky_bright = oskar.Sky.from_array((bright_sources()*fgd_att_vec))
    sky_bright_att = oskar.Sky.from_array((bright_sources()*fsl_att_vec*fgd_att_vec))
    sky0.append(sky_bright)
    sky1.append(sky_bright_att)
    
if (do_GLEAM):
    # Load GLEAM plus LoBES v2 osm (from Trott et al. June 2022)
    sky_init = oskar.Sky.load(sm_dir+'GLM_LoBESv2.osm')
    sky_array = sky_init.to_array()
    sky_gleam = oskar.Sky.from_array(sky_array*fgd_att_vec)
    sky_gleam_att = oskar.Sky.from_array(sky_array*fsl_att_vec*fgd_att_vec)
    sky0.append(sky_gleam)
    sky1.append(sky_gleam_att)
    
# Filter sky models for pointing centre and outer cutoff value
sky_c = make_sky_model(sky0, sky1, fgd_att, fsl_att, settings, fovrad_deg, cross_flux, flux_val)
# Now save comnpact sky model to file which can be loaded inside channel loop
sky_c_name = out_root+'_c.osm'
sky_c.save(sky_c_name)

if (do_DD):
    # Attenuate ionospheric effects and write as fits for oskar to use
    text = miriadbin+'maths exp="<'+iono_total+'>*'+str(ion_res)+'" out='+iono_total_att
    subprocess.run(text, shell='True')
    text = miriadbin+'fits op="xyout" in='+iono_total_att+' out='+iono_total_attf
    subprocess.run(text, shell='True')
    settings['telescope/external_tec_screen/input_fits_file'] = iono_total_attf


# Save nominal settings as dictionary so we can pass them to function in channel loop
settings_dict = settings.to_dict(include_defaults=True)

# For each frequency:
# Extract the relevant slice from the diffuse image sky models (low and high res) and append these componets to the compact sky model
# Do the oskar simulation 

out_ms = out_root+'.ms'
out_msm = out_root+'.msm'
out_msw = out_root+'.msw'
out_msn = out_root+'.msn'
out_mswf_ima = out_root+'.msw_image.fits'
out_mswf_psf = out_root+'.msw_psf.fits'
out_msnf_ima = out_root+'.msn_image.fits'
out_msnf_psf = out_root+'.msn_psf.fits'

# Create the gain errors to use in the simulation

data_shape = (num_int, num_freq, num_tel)
gains = np.empty(data_shape,dtype=np.csingle)

# Distribution parameters.

t_beta = time_psd_exp # time psd exponent
t_mean_amp = 1.0
t_mean_phase = 0.0  # radians
t_std_amp = amp_err
t_std_phase = phas_err*np.pi/180.  # radians

f_beta = freq_psd_exp # freq psd exponent
f_mean_amp = 1.0
f_mean_phase = 0.0  # radians
f_std_amp = bp_amp_err
f_std_phase = bp_phas_err*np.pi/180.  # radians

for itel in range(0, num_tel):
    rseedi = rseed + itel*num_tel
    t_a0 = cn.powerlaw_psd_gaussian(t_beta, num_int, random_state=rseedi)
    t_amp = (t_a0-np.average(t_a0))*t_std_amp/np.std(t_a0)+t_mean_amp
    t_p0 = cn.powerlaw_psd_gaussian(t_beta, num_int, random_state=rseedi+1)
    t_phase = (t_p0-np.average(t_p0))*t_std_phase/np.std(t_p0)+t_mean_phase
    t_gains = t_amp * np.exp(1j * t_phase)

    f_a0 = cn.powerlaw_psd_gaussian(f_beta, num_freq, random_state=rseedi+2)
    f_amp = (f_a0-np.average(f_a0))*f_std_amp/np.std(f_a0)+f_mean_amp
    f_p0 = cn.powerlaw_psd_gaussian(f_beta, num_freq, random_state=rseedi+3)
    f_phase = (f_p0-np.average(f_p0))*f_std_phase/np.std(f_p0)+f_mean_phase
    f_gains = f_amp * np.exp(1j * f_phase)

    tf_gains = np.outer(t_gains,f_gains)
    gains[:,:,itel] = tf_gains

# Write HDF5 file with recognised dataset names.
with h5py.File(tel_dir+"gain_model.h5", "w") as hdf_file:
    hdf_file.create_dataset("freq (Hz)", data=freqs)
    hdf_file.create_dataset("gain_xpol", data=gains)

if (do_SIM):    
    num_0 = 0
    if (do_restart_SIM):
        num_0 = num_sim_res
        
    # Fix pool number in this loop to match the number of GPUs available and set the GPU ids appropriately
    pool = mp.Pool(processes=num_gpu)
    
    for ifrq in range(num_0,num_freq,1):
        pool.apply_async(oskar_proc, args=(settings_dict,ifrq,out_root,do_TRECS,gxypix_asec,txypix_asec,num_gpu,rseed,sky_c_name), callback = log_result)

    pool.close()
    pool.join()
    print (result_list)

    # Fix pool number in this loop to match the number of CPUs since it can safely run in parallel
    pool = mp.Pool(processes=num_cpu)

    for ifrq in range(num_0,num_freq,1):
        sifrq = str(ifrq).zfill(len(str(num_freq))+1)
        out_uvfrq = out_root+'_IFRQ_'+sifrq+'.uv'
        out_msfrq = out_root+'_IFRQ_'+sifrq+'.ms'
        out_uvffrq = out_root+'_IFRQ_'+sifrq+'.uvf'
        script_line0 = 'vishead(vis="'+out_msfrq+'", mode="put", hdkey="telescope", hdvalue="SKA1-LOW")\n'
        script_line1 = 'exportuvfits(vis="'+out_msfrq+'", fitsfile="'+out_uvffrq+'", datacolumn="data", multisource=False, writestation=False, overwrite=True)\n'
        if (do_uvf):
            with open('_casa_script'+sifrq+'.py', 'w') as f:
                f.write(script_line0)
                f.write(script_line1)
        else:
            with open('_casa_script'+sifrq+'.py', 'w') as f:
                f.write(script_line0)
            
        ctext = casabin+'casa --nologger -c _casa_script'+sifrq+'.py'
        pool.apply_async(casa_proc, args=(ctext,sifrq), callback = log_result)

    pool.close()
    pool.join()
    print (result_list)

if (do_wsc):
    num_0 = 0
    mem_frac = int(100/num_wsclean)
    if (do_restart_wsc):
        num_0 = num_wsc_res
        
    # Number of wscleans that we can safely run in parallel
    pool = mp.Pool(processes=num_wsclean)
    for ifrq in range(num_0,num_freq,1):
        sifrq = str(ifrq).zfill(len(str(num_freq))+1)
        out_uvfrq = out_root+'_IFRQ_'+sifrq+'.uv'
        out_msfrq = out_root+'_IFRQ_'+sifrq+'.ms'
        out_mswfrq = out_root+'_IFRQ_'+sifrq+'.wsc'
        out_msnfrq = out_root+'_IFRQ_'+sifrq+'.wsn'
        out_uvffrq = out_root+'_IFRQ_'+sifrq+'.uvf'
        text =wsbin+'wsclean -no-update-model-required -reorder -mem '+str(mem_frac)+' -use-wgridder -multiscale -parallel-gridding 10 -weight uniform -oversampling 4095 -kernel-size 15 -nwlayers 1000 -grid-mode kb -taper-edge 100 -padding 2 -taper-gaussian '+str(int(ouvtap))+' -super-weight 4 -name '+out_mswfrq+' -size '+str(int(opix))+' '+str(int(opix))+' -scale '+str(int(ocell))+'asec -niter '+str(int(oniter))+' -auto-threshold 4 -mgain 0.8 -pol xx -make-psf '+out_msfrq
        pool.apply_async(subp_proc, args=(text,sifrq), callback = log_result)

    pool.close()
    pool.join()
    print (result_list)


if (do_wsn):
    num_0 = 0
    mem_frac = int(100/num_wsclean)
    if (do_restart_wsc):
        num_0 = num_wsc_res
        
    # Number of wscleans that we can safely run in parallel
    pool = mp.Pool(processes=num_wsclean)
    for ifrq in range(num_0,num_freq,1):
        sifrq = str(ifrq).zfill(len(str(num_freq))+1)
        out_uvfrq = out_root+'_IFRQ_'+sifrq+'.uv'
        out_msfrq = out_root+'_IFRQ_'+sifrq+'.ms'
        out_mswfrq = out_root+'_IFRQ_'+sifrq+'.wsc'
        out_msnfrq = out_root+'_IFRQ_'+sifrq+'.wsn'
        out_uvffrq = out_root+'_IFRQ_'+sifrq+'.uvf'
        text = wsbin+'wsclean -reorder -mem '+str(mem_frac)+' -use-wgridder -parallel-gridding 10 -weight natural -oversampling 4095 -kernel-size 15 -nwlayers 1000 -grid-mode kb -taper-edge 100 -padding 2 -name '+out_msnfrq+' -size '+str(int(opix))+' '+str(int(opix))+' -scale '+str(int(ocell))+'asec -niter 0 -pol xx -make-psf '+out_msfrq
        pool.apply_async(subp_proc, args=(text,sifrq), callback = log_result)

    pool.close()
    pool.join()
    print (result_list)

# We could concatenate all of the MSes into one, but this is probably OTT, so never mind    
# mscut_list = glob.glob(out_root+'_IFRQ_*.ms')
# script_line0 = 'concat(vis='+str(mscut_list)+', concatvis="'+out_ms+'", timesort=True)\n'
# with open('_casa_script.py', 'w') as f:
#    f.write(script_line0)



if do_wsc:
    ima_cube = np.single(np.zeros((int(num_freq),int(opix),int(opix))))
    psf_cube = np.single(np.zeros((int(num_freq),int(opix),int(opix))))

    # No point in parallising this since I/O bound
    for ifrq in range(0,num_freq,1):
        sifrq = str(ifrq).zfill(len(str(num_freq))+1)
        # define some names
        out_msfrq = out_root+'_IFRQ_'+sifrq+'.ms'
        out_mswfrq = out_root+'_IFRQ_'+sifrq+'.wsc'
        psf_name = out_mswfrq+'-psf.fits'
        dir_name = out_mswfrq+'-dirty.fits'
        mod_name = out_mswfrq+'-model.fits'
        res_name = out_mswfrq+'-residual.fits'
        ima_name = out_mswfrq+'-image.fits'
        ima_hdu = fits.open(ima_name)
        data = np.single(ima_hdu[0].data[:,:])
        ima_cube[ifrq,:,:] = data
        psf_hdu = fits.open(psf_name)
        data = np.single(psf_hdu[0].data[:,:])
        psf_cube[ifrq,:,:] = data
        if (ifrq==0):
            ima0_hdr = ima_hdu[0].header
            psf0_hdr = psf_hdu[0].header

        text = '/bin/rm -r '+psf_name+' '+dir_name+' '+mod_name+' '+res_name+' '+ima_name+' '
        subprocess.run(text, shell='True')

    hdu_ima_cube = fits.PrimaryHDU(ima_cube)
    hdu_ima_cube.header = ima0_hdr
    hdu_ima_cube.writeto(out_mswf_ima)

    hdu_psf_cube = fits.PrimaryHDU(psf_cube)
    hdu_psf_cube.header = psf0_hdr
    hdu_psf_cube.writeto(out_mswf_psf)


if do_wsn:
    iman_cube = np.single(np.zeros((int(num_freq),int(opix),int(opix))))
    psfn_cube = np.single(np.zeros((int(num_freq),int(opix),int(opix))))

    # No point in parallising this since I/O bound
    for ifrq in range(0,num_freq,1):
        sifrq = str(ifrq).zfill(len(str(num_freq))+1)
        # define some names
        out_msfrq = out_root+'_IFRQ_'+sifrq+'.ms'
        out_msnfrq = out_root+'_IFRQ_'+sifrq+'.wsn'
        psfn_name = out_msnfrq+'-psf.fits'
        dirn_name = out_msnfrq+'-dirty.fits'
        iman_name = out_msnfrq+'-image.fits'
        iman_hdu = fits.open(iman_name)
        data = np.single(iman_hdu[0].data[:,:])
        iman_cube[ifrq,:,:] = data
        psfn_hdu = fits.open(psfn_name)
        data = np.single(psfn_hdu[0].data[:,:])
        psfn_cube[ifrq,:,:] = data

        if (ifrq==0):
            ima0n_hdr = iman_hdu[0].header
            psf0n_hdr = psfn_hdu[0].header

        text = '/bin/rm -r '+psfn_name+' '+dirn_name+' '+iman_name+' '
        subprocess.run(text, shell='True')

    hdu_iman_cube = fits.PrimaryHDU(iman_cube)
    hdu_iman_cube.header = ima0n_hdr
    hdu_iman_cube.writeto(out_msnf_ima)
    hdu_psfn_cube = fits.PrimaryHDU(psfn_cube)
    hdu_psfn_cube.header = psf0n_hdr
    hdu_psfn_cube.writeto(out_msnf_psf)

# Tidy up all temporary files
for ifrq in range(0,num_freq,1):
    sifrq = str(ifrq).zfill(len(str(num_freq))+1)
    # define some names
    gsm_frq = gsm+'_IFRQ_'+sifrq
    gsm_frqf = gsm+'_IFRQ_'+sifrq+'.fits'
    text = '/bin/rm -r '+gsm_frq+' '+gsm_frqf+' '
    subprocess.run(text, shell='True')
    if (do_SIM):
        xfile='_casa_script'+sifrq+'.py'
        text = '/bin/rm -r '+xfile 
        subprocess.run(text, shell='True')

    if (do_TRECS):
        tsm_frq = tsm+'_IFRQ_'+sifrq
        tsm_frqf = tsm+'_IFRQ_'+sifrq+'.fits'
        text = '/bin/rm -r '+tsm_frq+' '+tsm_frqf+' '
        subprocess.run(text, shell='True')
    
if (toss_models):
    text = '/bin/rm -r '+gsm+' '+gsmx+' '+gsmxx+' '+gsmxxx+' '+gsmxxxf+' '+gsmf
    subprocess.run(text, shell='True')

if (do_DD):
    text = '/bin/rm -r '+iono_total_att+' '+iono_total_attf
    subprocess.run(text, shell='True')

if (do_MHD):
    text = '/bin/rm -r '+tdrot+' '+tdrotf+' '+gtd+' '+gtdf+' '
    subprocess.run(text, shell='True')

if (do_EOR):
    text = '/bin/rm -r '+eor_reg+' '+eor_norm+' '+eor_normr
    subprocess.run(text, shell='True')
    
if (do_TRECS):
    text = '/bin/rm -r '+trecs_mask+' '+trecsb_cube+' '+trecsb_cubepr+' '+trecsf_cube+' '+trecsf_cubepr+' '+trecsf_cubeprc67
    subprocess.run(text, shell='True')
    text = '/bin/rm -r '+trecsb_reg+' '+trecsf_reg+' '+trecsf_regr+' '+trecsb_norm+' '
    subprocess.run(text, shell='True')
    if (toss_models):
        text = '/bin/rm -r '+tsm
        subprocess.run(text, shell='True')

    
os.chdir(cur_dir)