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'
}

line_params = {
    'specmode': 'cube', # Must use 'cube' for spectral line imaging
    'spw': '0', # We're using the first SPW here, but feel free to change this (+ channel selection and threshold) if you want to image other SPWs
    'weighting': 'briggsbwtaper',
    'restoringbeam': 'common' # If not used, each channel will have a slighly different beam, as it is frequency-dependent
}

"""
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)

"""
LINE IMAGING
"""

## 3. Perform continuum subtraction

if not os.path.isdir(contsub_vis):
    print(f"Performing continuum subtraction from {split_vis}")
    
    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 uvcontsub")
    
    uvcontsub(vis=split_vis,
              outputvis=contsub_vis,
              fitspec=CONT_CHANNELS,
              fitorder=0, # Polynomial order for continuum fitting, typically 0 or 1
              writemodel=False,
              datacolumn=column)

## 4. Make a dirty image of the full spectral cube

spw = line_params['spw']
full_line_name = f"{image_basename}.spw_{spw}.full"
dirty_full_line_name = f"{full_line_name}.dirty"

if not os.path.isdir(f"{dirty_full_line_name}.image"):
    print(f"Creating dirty full line image: {dirty_full_line_name}")

    tb.open(contsub_vis)
    colnames = tb.colnames()
    tb.close()
    column = 'corrected' if 'CORRECTED_DATA' in colnames else 'data'
    print(f"Using {column.upper()} column for tclean")
    
    dirty_line_params = {**common_params, **line_params, 'pbcor': False}

    tclean(vis=contsub_vis,
           imagename=dirty_full_line_name,
           selectdata=True,
           datacolumn=column,
           **dirty_line_params)

## 5. Process spectral line chunks
# This roughly covers the channels where we see line emission in the dirty cube, 
# but feel free to change these based on what you see in the dirty image and what you want to explore. 
# I've pre-filled one chunk here, but you can add more by following the format in the LINE_CHUNKS list below.
# Try adding in the second line that we see in the dirty cube

LINE_CHUNKS = [
    {'start': 400, 'width': 1, 'nchan': 245},
    # {'start': LINE_2_START, 'width': 1, 'nchan': LINE_2_NCHAN}
]

print("Processing individual line chunks...")
for chunk in LINE_CHUNKS:
    start_chan = int(chunk['start'])
    end_chan = start_chan + int(chunk['nchan']) - 1
    chunk_name = f"{image_basename}.spw_{spw}.chan_{start_chan}-{end_chan}"

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

        clean_line_params = {**common_params, **line_params,
                             'niter': 1000, 
                             'usemask': 'auto-multithresh', 
                             'threshold': '0.0Jy'} # <--- MEASURE THE RMS IN THE DIRTY IMAGE TO SET THIS THRESHOLD!

        tclean(vis=contsub_vis,
               imagename=chunk_name,
               selectdata=True,
               datacolumn=column,
               **clean_line_params,
               start=chunk['start'],
               width=chunk['width'],
               nchan=chunk['nchan'])
        
    # Export the cleaned line image to FITS format for later use in analysis
    exportfits(imagename=chunk_name+'.image',
                fitsimage=chunk_name+'.fits',
                dropdeg=True,
                overwrite=True)