#!/usr/bin/python3 -Wignore
"""Starbug II - JWST PSF photometry
usage: starbug2 [-ABCDfhMPv] [-b bgdfile] [-d apfile] [-o directory] [-p file.param] [-s opt=val] image.fits ...
   -A  --artific         : run artificial star tests
   -B  --background      : run background estimation
   -b  --bgdfile         : load background (-bgd.fits) file
   -C  --clean           : run source cleaning before photometry 
   -d  --apfile  ap.fits : load a source detection (-ap.fits) file to skip the source detection step
   -D  --detect          : run source detection
   -f  --find            : attempt to find associated -ap -bgd files
   -h  --help            : display uasage information
   -M  --match           : match outputs from all input image files
   -n  --ncores      num : number of CPU cores to split process between
   -o  --output      dir : output directory
   -p  --param   a.param : load parameter file
   -P  --photom          : run psf photometry
   -s  --set      option : set value in parameter file at runtime (-s SIGSKY=3)
   -v  --verbose         : display verbose outputs

   --> Single run commands
       --init                     : Initialise Starbug (post install)
       --generate-psf             : Generate ALL the PSF files to "PSFDIR"
       --local-param              : Make a local copy of the default parameter file
       --generate-region   a.fits : Make a ds9 region file with a detection file
       --clean-table       a.fits : Clean up an individual table
       --generate-run      *.fits : Generate a simple run script
       --version                  : Print starbug2 version
"""
import os,sys,getopt
sys.stdout.write("loading..\r")
import pkg_resources
import starbug2 
from starbug2.utils import *
from starbug2.misc import *
from starbug2.matching import dither_match
import multiprocessing

VERBOSE= 0x01
KILLPROC=0x02
STOPPROC=0x04
FINDFILE=0x08

DODETECT=0x100
DOBGDEST=0x200
DOPHOTOM=0x400
DOCLEAN= 0x800 
DOARTIFL=0x1000
DOMATCH= 0x2000

GENRATPSF =0x10000
GENRATRUN =0x20000
INITSB    =0x40000

options=0

pfile=None
setdict={}
output=None
ncores=1

def usage():
    if not (options&VERBOSE): quit(__doc__.split('\n')[1])
    else: quit(__doc__)

opts,args=getopt.getopt(sys.argv[1:],"ABCDfhMPvb:d:n:o:p:s:", ("artific","background", "clean", "detect", "find", "help", "match", "photom", "verbose", "xtest",
                                                      "bgdfile=", "apfile=", "ncores=", "output=", "param=", "set=",
                                                      "init", "generate-psf", "local-param", "generate-region=", "clean-table", "version", "generate-run"))

for opt,optarg in opts:
    if opt in ("-h","--help"):    usage()
    if opt in ("-p","--param"):   pfile= optarg
    if opt in ("-v","--verbose"):
        options|=VERBOSE
        #setdict["VERBOSE"]=1

    if opt in ("-A","--artific"): options |= DOARTIFL
    if opt in ("-B","--background"): options|=DOBGDEST
    if opt in ("-C","--clean"):  options |=DOCLEAN
    if opt in ("-D","--detect"): options |= DODETECT
    if opt in ("-M","--match"):  options |= DOMATCH
    if opt in ("-P","--photom"): options |= DOPHOTOM

    if opt in ("-d","--apfile"):
        if os.path.exists(optarg): setdict["AP_FILE"]=optarg
        else: perror("AP_FILE does not exist\n")

    if opt in ("-b","--bgdfile"):
        if os.path.exists(optarg): setdict["BGD_FILE"]=optarg
        else: perror("BGD_FILE does not exist\n")

    if opt in ("-f","--find"): options|=FINDFILE
    if opt in ("-n","--ncores"): ncores=int(optarg)
        #options &= ~VERBOSE

    if opt in ("-o","--output"):
        if os.path.exists(optarg): 
            output=optarg
            setdict["OUTDIR"]=optarg
        else:
            perror("output directory '%s' does not exist\n"%optarg)
            options|=KILLPROC

    if opt in ("-s","--set"): 
        if '=' in optarg:
            key,val=optarg.split('=')
            try: val=float(val)
            except: pass
            setdict[key]=val
        else:
            perror("unable to set parameter, use syntax -s KEY=VALUE\n")
            options|=KILLPROC

    if opt=="--init": options|=(INITSB|STOPPROC)
    if opt=="--generate-psf": options|=(GENRATPSF|STOPPROC)

    if opt=="--local-param":
        os.system("cp %s/default.param starbug.param"%pkg_resources.resource_filename("starbug2","param/"))
        printf("--> generating starbug.param\n")
        options|=STOPPROC

    if opt=="--generate-region":
        table=Table.read(optarg,format="fits")#load_table(optarg)
        _,fname,_=split_fname(optarg)
        export_region(table, fname="%s.reg"%fname)
        printf("generating region --> %s.reg\n"%fname)
        options|=STOPPROC

    if opt=="--clean-table":
        printf("cleaning table <-- %s\n"%optarg)
        table=Table.read(optarg,format="fits")#load_table(optarg)
        options|=STOPPROC

    if opt=="--version": quit("starbug2-v%s"%pkg_resources.get_distribution("starbug2").version)

    if opt=="--generate-run": 
        options|=(GENRATRUN|STOPPROC)

##############################################
# Options set, verify/run one time functions #
##############################################

if ncores>1: options&= ~VERBOSE


if pfile==None:
    if os.path.exists("./starbug.param"):pfile="starbug.param"
    else:
        perror("no local parameter file found, using default.param\n")
        pfile="%s/default.param"%pkg_resources.resource_filename("starbug2","param/")

if options&INITSB:
    init_starbug(load_params(pfile)["PSFDIR"])

if options&GENRATPSF:
    generate_psfs(load_params(pfile)["PSFDIR"])

if options&GENRATRUN:
    cmd="starbug2 -vf "#"+ " ".join( ["-vf"])
    starbug2.misc.generate_runscript(args, cmd)


if options&STOPPROC: quit(0) ## quiet ending the process if required

if len(args)==0: 
    perror("must include fits image file\n")
    options|=KILLPROC

if options&KILLPROC:
    perror("..quitting :(\n\n")
    usage()

############################################################################
# All the options set and one-time run functions done, Time to run Starbug #
############################################################################

from starbug2.starbug import StarbugBase ## Ive put this here because it takes some time

#def fn(fname, sb_list):
def fn(fname):
    printf("%s\n"%fname)
    sb=None
    if os.path.exists(fname):
        dname,bname,ext=split_fname(fname)

        if options&FINDFILE:
            ap="%s/%s-ap.fits"%(dname,bname)
            bgd="%s/%s-bgd.fits"%(dname,bname)
            if os.path.exists(ap): setdict["AP_FILE"]=ap
            if os.path.exists(bgd): setdict["BGD_FILE"]=bgd

        if ext==".fits":
            if options&VERBOSE: setdict["VERBOSE"]=1
            sb=StarbugBase(fname, pfile=pfile, options=setdict)
            
            if sb.verify(): return#quit("..quitting :(")
            #if options&VERBOSE: printf("fname:%s\n"%fname)
            
            if options & DODETECT: sb.detect()
            if options & DOCLEAN: sb.cleanup()
            if options & DOBGDEST: sb.bgd_estimate()
            if options & DOPHOTOM:
                sb.photometry()
                if options & DOCLEAN: sb.cleanup()
            if options & DOARTIFL: sb.artificial_stars()
            sb.export(outdir=output)

            #sb_list.append(sb)
        else: perror("file must be type '.fits' not %s\n"%ext)
    else: perror("can't access %s\n"%fname)
    return sb

if __name__=="__main__":

    printf("\x1b[1mlaunching \x1b[36mstarbug\x1b[0m\n")
    #starbugs=multiprocessing.Manager().list()
    #results=pool.starmap(fn, args)#[(fname,starbugs) for fname in args])
    #pool.map(fn, args)

    pool=multiprocessing.Pool(processes=ncores)
    starbugs=pool.map(fn,args)

    for n,sb in enumerate(starbugs): 
        if not sb: perror("FAILED: %s\n"%args[n])


    ###########################
    # Match Output Catalogues #
    ###########################

    if options|DOMATCH and len(starbugs)>1:
        fname=combine_fnames( [sb.fname for sb in starbugs] ) 
        fname="%s/%s"%(split_fname(fname)[:2]) if fname else "out"
        param=load_params(pfile)

        if options&VERBOSE: printf("Matching outputs\n")

        if options&DODETECT:
            colnames=starbug2.match_cols
            colnames+=[ name for name in param["MATCH_COLS"].split() if name not in colnames]

            av,full=dither_match( [sb.detections for sb in starbugs] , threshold=param["MATCH_THRESH"], colnames=colnames)

            N=len(starbugs)-param["RM_MATCH"]
            mask=( (av["NUM"]>N) if param["RM_MATCH"]>=0 else (av["NUM"]>0))

            printf("-> %s-ap*...\n"%fname)
            header={"FILTER":starbugs[0].filter}
            av.meta.update(header)
            #header.update(param)
            export_table(hstack((av,full)), fname="%s-apfull.fits"%fname, header=header)
            export_table(av[mask], fname="%s-apmatch.fits"%fname, header=header)

        if options&DOPHOTOM:
            perror("em sorry not implemented PSF tab matching yet\n")

