#!/usr/bin/env python

"""

Driver script for the Berk package

Given a MeerKAT captureBlockId:
    - fetch dataset from archive (mocked up at the moment)
    - unpack it set-it up in a staging directory
    - set-up an oxkat run
    - run oxkat pipeline stages, checking output after each stage
    - generate catalogs with PyBDSF
    - tidy up (keep images, catalogs, possibly MS for re-imaging)

"""

import os
import sys
import argparse
import subprocess
import glob
import time
from berk import startup, archive, jobs, tasks

#------------------------------------------------------------------------------------------------------------
def makeParser():
    parser=argparse.ArgumentParser("berk")
    parser.add_argument("task", help="Task to run - one of:\
                        'list' (list downloaded observations available for processing by captureBlockId),\
                        'fetch' (retrieve data from archive, store under $BERK_MSCACHE location),\
                        'process' (link measurement set under 'staging' directory, run calibration and imaging with Oxkat up to 2GC),\
                        'analyse' (apply primary beam correction and make object catalogs with PyBDSF),\
                        'collect' (collect data products for all processed observations - i.e., primary beam corrected images and catalogs),\
                        'store' (store a more extensive set of processed data products for an observations for possible further processing in future),\
                        'builddb'(build survey-wide database and related products).",
                        choices = ['list', 'fetch', 'process', 'analyse', 'collect', 'store', 'builddb'])
    parser.add_argument("-o", "--observation", dest="captureBlockId", default = None, help="MeerKAT observation to fetch and process.\
                        For the 'fetch' task, this must be a link to an .rdb file on the archive, of the form\
                        https://archive-gw-1.kat.ac.za/captureBlockId/captureBlockId_sdp_l0.full.rdb?token=longTokenString.\
                        For the other tasks ('process', 'analyse'), this should be the captureBlockId itself.")

    return parser

#------------------------------------------------------------------------------------------------------------
if __name__ == '__main__':

    parser=makeParser()
    args=parser.parse_args()
    captureBlockId=args.captureBlockId

    topDir=os.getcwd()

    needsCaptureBlockId=['fetch', 'process', 'analyse', 'store']
    if captureBlockId is None and args.task in needsCaptureBlockId:
        print("No observation provided - this should be either a captureBlockId (for the 'process', 'analyse', 'store' tasks) or a link to an .rdb file on the archive (for the 'fetch' task).")
        sys.exit()

    if args.task == 'fetch':
        tasks.fetch(captureBlockId)

    elif args.task == 'list':
        msList=glob.glob(os.environ['BERK_MSCACHE']+os.path.sep+"*_sdp_l0.ms")
        print("Downloaded observations available for processing [captureBlockId]:")
        # print("captureBlockId\tprocess\tanalyse\n")
        # print("
        for ms in msList:
            captureBlockId=os.path.split(ms)[-1].split("_")[0]
            print("   %s" % (captureBlockId))

    elif args.task == 'store':
        tasks.store(captureBlockId)

    elif args.task == 'collect':
        tasks.collect()

    elif args.task == 'process':
        tasks.process(captureBlockId)

    elif args.task == 'analyse':
        tasks.analyse(captureBlockId)

    elif args.task == 'builddb':
        tasks.builddb()





