import os

# File and target configuration
vis_file = 'uid___A002_X1003af4_Xa540.ms'  # 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)
You will need to manually set this to specific the continuum channels for each spectral window
Note that SPWs 3 and 4 are line-free and can be used as-is. Only SPWs 1 and 2 channels need to be specified.
[You should have already done this for yesterday's continuum imaging exercise, so you can reuse the same values here]
"""
# 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:240.9498861665~240.9567221040GHz;241.1676596040~241.3034017915GHz;241.3991049165~241.4098471040GHz;241.5182455415~241.6334799165GHz;241.7633627290~241.8004721040GHz;241.9303549165~241.9967611665GHz;242.0905111665~242.1666830415GHz;242.3834799165~242.4996908540GHz;242.5993002290~242.7106283540GHz,'
                 '3:243.9977299165~244.0231205415GHz;244.1656986665~244.2057377290GHz;244.2233158540~244.2467533540GHz;244.3033939790~244.3658939790GHz;244.4069096040~244.4615971040GHz;244.5074955415~244.5846439790GHz;244.6246830415~244.6715580415GHz;244.6949955415~244.7496830415GHz;244.7819096040~244.7897221040GHz;244.7994877290~244.8043705415GHz;244.8307377290~244.8385502290GHz;244.8492924165~244.8551517915GHz;244.8990971040~244.9108158540GHz;244.9879642915~245.0768314790GHz;245.1110111665~245.1530033540GHz;245.2487064790~245.2731205415GHz;245.3453861665~245.3893314790GHz;245.4274174165~245.5621830415GHz;245.6285892915~245.6823002290GHz;245.7545658540~245.8151127290GHz;245.8561283540~245.8619877290GHz'
                 )
NITER = 0  ### <- Update this for full clean ###
THRESHOLD = '0Jy'  ### <- Update this for full clean ###
ROBUST = 0.5  # Feel free to play with the robust parameter to see how it affects the image
INTERACTIVE = False  # Set to True if you want to run the imaging interactively
MASKTYPE = 'auto-multithresh'

"""
To image specific lines, you will need to specify the start channel, channel width, and number of channels for each chunk.
You can also use units of frequency or velocity instead of channel numbers (see tclean documentation).
The below numbers are just placeholders and will need to be updated based on the spectral line(s) in the data.
"""
LINE_CHUNKS = [
    {'start': 1, 'width': 1, 'nchan': 1},
    {'start': 1, 'width': 1, 'nchan': 1}
]

# Imaging parameters
tclean_params = {
    'imsize': [320, 300],
    'cell': ['0.22arcsec'],
    'phasecenter': 'ICRS 17:47:56.2008 -029.59.39.588',
    'gridder': 'mosaic',
    'deconvolver': 'hogbom',
    'robust': ROBUST,
    'pbcor': True,
    'niter': NITER,
    'usemask': MASKTYPE,
    'interactive': INTERACTIVE,
    'specmode': 'cube',
    'spw': '0',
    'threshold': THRESHOLD,
    'weighting': 'briggsbwtaper',
    'restoringbeam': 'common'
}

# --- Step-by-step walkthrough ---

## 1. Split out the target data and science spectral windows
"""
You should have already done this for yesterday's continuum imaging exercise, so you can skip this step.
I just left it in for completeness.
"""

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. Subtract continuum from the split MS

if not os.path.exists(contsub_vis):
    print(f"Performing continuum subtraction")
    
    uvcontsub(vis=split_vis,
              outputvis=contsub_vis,
              fitspec=CONT_CHANNELS,
              fitorder=0,
              datacolumn='data')

## 3. Make dirty cube of full spectral window

dirty_line_name = f"{image_basename}.spw0.dirty"

if not os.path.isdir(f"{dirty_line_name}.image"):
    print(f"Creating dirty line cube: {dirty_line_name}")
    
    # Create a modified parameter dict for dirty cube
    dirty_params = tclean_params.copy()
    dirty_params.update({
        'niter': 0,
        'usemask': None,
        'threshold': None,
        'interactive': False
    })
    
    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")
    
    tclean(vis=contsub_vis,
           imagename=dirty_line_name,
           selectdata=True,
           datacolumn=column,
           **dirty_params)

## 4. Make clean line images for each chunk

for chunk_idx, chunk in enumerate(LINE_CHUNKS):
    line_name = f"{image_basename}.spw0.chunk{chunk_idx}"
    
    if not os.path.isdir(f"{line_name}.image"):
        print(f"Creating line cube for chunk {chunk_idx}")
        
        # Create a modified parameter dict for this chunk
        chunk_params = tclean_params.copy()
        chunk_params.update({
            'start': chunk['start'],
            'width': chunk['width'],
            'nchan': chunk['nchan']
        })
        
        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")
        
        tclean(vis=contsub_vis,
               imagename=line_name,
               selectdata=True,
               datacolumn=column,
               **chunk_params)