import os

# File and target configuration
vis_file = 'uid___A002_X1003af4_Xa540.ms.split.cal'  # Change name if necessary
target_name = 'PN_Hb_5'
target_spw = '25,27,29,31'

# Define paths
split_vis = f"{vis_file}.target"
contsub_vis = f"{vis_file}.target.contsub"
image_basename = target_name

# Continuum channels (line-free)
CONT_CHANNELS = (
    '0:226.2010990572~226.2665287447GHz;226.3939701509~226.5687748384GHz;226.9725834322~227.1312748384GHz,'
    '1:230.0764655102~230.1931647289GHz;230.2170905102~230.3958014477GHz;230.6809576977~230.8245123852GHz;230.8484381664~231.0046881664GHz,'
    '2,'
    '3'
)

# Imaging parameters
common_params = {
    'imsize': [300, 300], # Set this to cover the full FOV, which you can estimate via FOV ~ 1.22 * lambda / D, where D is the dish diameter(= 12m here)
    'cell': ['0.22arcsec'], # Typically want the cell size to be ~1/5 of the synthesised beam size for proper sampling
    'phasecenter': 'ICRS 17:47:56.2008 -029.59.39.588',
    'gridder': 'mosaic',
    'deconvolver': 'hogbom', # Consider using 'multiscale' for extended emission, but start with 'hogbom' for simplicity
    'robust': 0.5, # Feel free to experiment with different robust values (e.g., 0.0, 1.0, -1.0) to see how it affects the image fidelity and resolution
    'pbcor': True,
    'niter': 0,
    'usemask': None,
    'interactive': False
}

continuum_params = {
    'specmode': 'mfs', # Multi-frequency synthesis ('mfs') for continuum imaging
    'spw': CONT_CHANNELS,
    'threshold': None,
    'weighting': 'briggs'
}

"""
CONTINUUM IMAGING
"""

## 1. Split out the target data and science spectral windows

if not os.path.isdir(split_vis):
    print(f"Splitting {vis_file} to {split_vis}")
    
    tb.open(vis_file)
    colnames = tb.colnames()
    tb.close()
    column = 'corrected' if 'CORRECTED_DATA' in colnames else 'data'
    print(f"Using {column.upper()} column for split")

    split(vis=vis_file,
          outputvis=split_vis,
          field=target_name,
          spw=target_spw,
          datacolumn=column)

## 2. Make continuum images

# 2.1 Make dirty continuum image

continuum_name = f"{image_basename}.continuum"
dirty_continuum_name = f"{continuum_name}.dirty"

if not os.path.isdir(f"{dirty_continuum_name}.image"):
    print(f"Creating dirty continuum image: {dirty_continuum_name}")
    
    tb.open(split_vis)
    colnames = tb.colnames()
    tb.close()
    column = 'corrected' if 'CORRECTED_DATA' in colnames else 'data'
    print(f"Using {column.upper()} column for tclean")
    
    dirty_cont_params = {**common_params, **continuum_params, 'pbcor': False}
    
    tclean(vis=split_vis,
           imagename=dirty_continuum_name,
           selectdata=True,
           datacolumn=column,
           **dirty_cont_params)

# 2.2 Make clean continuum image

if not os.path.isdir(f"{continuum_name}.image"):
    print(f"Creating clean continuum image: {continuum_name}")
    
    tb.open(split_vis)
    colnames = tb.colnames()
    tb.close()
    column = 'corrected' if 'CORRECTED_DATA' in colnames else 'data'
    print(f"Using {column.upper()} column for tclean")

    clean_cont_params = {**common_params, **continuum_params,
                         'niter': 1000, # This can be arbitrarily high if your threshold is well-defined
                         'usemask': 'auto-multithresh', # We'll use auto-masking, but you can also experiment with manual masks
                         'threshold': '0.0Jy'} # <--- MEASURE THE RMS IN THE DIRTY IMAGE TO SET THIS THRESHOLD!

    tclean(vis=split_vis,
           imagename=continuum_name,
           selectdata=True,
           datacolumn=column,
           **clean_cont_params)
    
# 2.3 Export the cleaned continuum image to FITS format for later use in analysis
exportfits(imagename=continuum_name+'.image',
           fitsimage=continuum_name+'.fits',
           dropdeg=True,
           overwrite=True)