"""
task UWpipeline Interface to the ISOC PipelineII
"""
import os, argparse,  datetime
import numpy as np
from uw.like2 import (process, tools, )
from uw.like2.pipeline import (check_data, pipeline_job, stagedict, check_converge, stream)
from uw.like2.analyze import app


# Instances of these classes are in the procnames dictionary below. They must implement a main function which will
# by run by the Proc
class StartStream(object):
    """ setup, start a stream """
    def main(self, args):
        ps  = stream.PipelineStream()
        for the_stage in args.stage:
            # note set job_list to the name of a local file in the pointlike folder -- needs better encapsulation
            job_list = args.job_list
            if job_list is None or job_list=='None':
                job_list = stagedict.stagenames[the_stage]['job_list']
            ps(the_stage, job_list, test=args.test)
            
class Summary(object):
    """ runs a sumamry job, via uw.like2.analyze.app"""
    def get_stage(self, args):
        thedict = args.stage[0] 
        t = thedict.split(':',1)
        if len(t)==2:
            stage, nextstage = t 
        else: stage,nextstage = t[0], None
        return stage
    def main(self, args):
        stage = self.get_stage(args)
        kw = stagedict.stagenames[stage].get('sum', None)
        if kw is not None:
            app.main(kw.split())
            
class JobProc(Summary):
    """ process args for running pipeline jobs"""
    def main(self, args):
        stage = self.get_stage(args)
        proc, jobargs = stagedict.stagenames[stage].setup()
        if args.rois!='':
            print '-->passing roi list:', args.rois
            jobargs['roi_list']=args.rois
        print '--> starting pipeline job: running class %s with args %s' % (proc, jobargs)
        pipeline_job.main(proc, **jobargs)
  
class CheckJobs(Summary):
    """ process args for the check_jobs step"""
    def main(self, args):
        from uw.like2.pub import healpix_map
        stage = self.get_stage(args)
        # do something instead
        check_converge.main(args)

class CheckData(Summary):
    def main(self, args):
        stage = self.get_stage(args)
        if os.path.exists('failed_rois.txt'):
           os.remove('failed_rois.txt')
           print 'Removed file with failed rois'
        
        check_data.main(args)
 
class Proc(dict):
    def __init__(self, run, help='', **kwargs):
        """ run: class or module -- must have main function """
        super(Proc, self).__init__(self,  help=help, **kwargs)
        self.run = run
        self['help']=help

    def __call__(self, args):
        self.run.main(args)
 
procnames = dict(
    # proc names (except for start) are generated by the UWpipeline task as stream executes
    # start actually launches a stream
    start      = Proc(StartStream(), help='start a stream: assumed if no proc is specified'),
    check_data = Proc(CheckData(),   help='check that required data files are present'),
    job        = Proc(JobProc(),     help='run a parallel pipeline job'),
    job_proc   = Proc(JobProc(),     help='run a parallel pipeline job'),
    check_jobs = Proc(CheckJobs(),   help='check for convergence, combine results, possibly submit new stream'),
    summary_plots= Proc(Summary(),   help='Generate summary plots, depending on stage name'),
    )
proc_help = '\nproc names\n\t' \
    +'\n\t'.join(['%-15s: %s' % (key, procnames[key]['help'])  for key in sorted(procnames.keys())])
    

#keys = stagedict.stagenames.keys()
#stage_help = '\nstage name, or sequential stages separated by "y:" names are\n\t' \
#    +  '\n\t'.join(['%-15s: %s' % (key,stagedict.stagenames[key]['help']) \
#        for key in sorted(stagedict.stagenames.keys())])

def find_script_folder(cwd):
    """ look for folder with scripts: expect to be in folder containing skymodels or scripts
    cwd : string, current folder
        """
    m = cwd.find('skymodels')
    if m<0:
        print 'WARNING: did not find "skymodels" in path to cwd, which is %s' %cwd
        script_folder=cwd
    else: script_folder=cwd[:m]
    if not os.path.exists(os.path.join(script_folder, 'configure.sh')):
        # try scripts folder
        test = os.path.join(script_folder, 'scripts')
        if os.path.exists(test):
            if not os.path.exists(os.path.join(test, 'configure.sh')):
                raise Exception('Script "configure.sh" not  found in %s or %s' % (script_folder, test) )
            script_folder = test
        else:
            raise Exception('Script "configure.sh" not  found in %s ' % (script_folder,) )
    return script_folder  
  
def check_environment(args):
    if 'SKYMODEL_SUBDIR' not in os.environ:
        os.environ['SKYMODEL_SUBDIR'] = os.getcwd()
    else:
        skymodel = os.environ['SKYMODEL_SUBDIR']
        print 'skymodel:' , skymodel
        assert os.path.exists(skymodel), 'Bad path for skymodel folder: %s' %skymodel
        os.chdir(skymodel)
    cwd = os.getcwd()

    def checkit(f):
        return os.path.exists(f) or os.path.exists('../'+f)
    assert checkit('config.txt') or checkit('config.yaml'), \
        'expect this folder (%s), or its parent, to have a file config.txt or config.yaml'%cwd
    
    if args.scripts is None:
        script_folder = find_script_folder(cwd)
    else:
        if not os.path.exists(args.scripts) :
            raise Exception( 'SCRIPT folder %s not found' % args.scripts) 
        if not os.path.exists(os.path.join(args.scripts, 'configure.sh')):
            raise Exception('File "configure.sh" not found in folder %s' % args.scripts)
        script_folder = args.scripts
    
    if args.stage[0] is not None :
        os.environ['stage']=args.stage[0]

    # add these to the Namespace object 
    args.__dict__.update(skymodel=cwd, pointlike_dir=script_folder, script_folder=script_folder)

def check_names(stage, proc):
    keys = stagedict.keys
    if len(stage)==0:
        if proc is  None:
            raise Exception('No proc or stage argement specified')
        if proc not in procnames:
            raise Exception('proc name "%s" not in list %s' % (proc, procnames))
        return
    if stage[0] is None: 
        raise Exception('no stage specified')
    for s in stage:
        for t in s.split(':'):
            if t not in keys:
                raise Exception('stage "%s" not recognized: expect one of %s' %(t, sorted(keys)))

def main( args ):
    check_environment(args)
    check_names(args.stage, args.proc)
    proc = args.proc
    #tee = tools.OutputTee('summary_log.txt')
    print '\n'+ str(datetime.datetime.today())[:16]
    print '--> %s for %s'%(proc, args.stage)
    if proc not in procnames:
        print 'proc name "%s" not recognized: expect one of %s' % (proc, sorted(procnames))
        return
    try:
        procnames[proc](args)
    except Exception, msg:
        print 'Exception trying to execute procnames[%s](%s):\n\t%s' % ( proc,args, msg)
        raise
    #tee.close()

if __name__=='__main__':
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
            description=""" start a UWpipeline stream, or run a UWpipeline proc; \n
    """+stagedict.help+proc_help+"""
    \nExamples: 
        uwpipeline create
        uwpipeline finish -p summary_plots
    """)
    parser.add_argument('stage', nargs='*', default=[os.environ.get('stage', None)], help='stage name, default "%(default)s"')
    parser.add_argument('-p', '--proc', default=os.environ.get('PIPELINE_PROCESS', 'start'), 
            help='proc name,  default: "%(default)s"')
    parser.add_argument('--job_list', default=os.environ.get('job_list', None), help='file used to allocate jobs, default "%(default)s"')
    parser.add_argument('--scripts', default=None,    help='script folder for batch, must be writeable default %(default)s')
    parser.add_argument('--rois', default='', help='allow setting of list for special job'),
    parser.add_argument('--stream', default=os.environ.get('PIPELINE_STREAM', -1), help='pipeline stream number, default %(default)s')
    parser.add_argument('--test', action='store_true', help='Do not run' )
    parser.add_argument('--next', default=None, help='Command to run after successful completion')
    #parser.add_argument('--processor',  help='specify the processor' )
    args = parser.parse_args()
    main(args)
    
