"""
The purpose of this class is to:
- Implement the stack.pro pantasks script tasks
- Generate the associated DAG
"""

#
# Generate the stack DAG for a label and (a set of exposures?)
#
from ipp.Gpc1Manager import Gpc1Manager as Gpc1Manager
from ipp.helpers import mkdir
from ipp.constants.Globals import Globals as GlobalConstants
from ipp.constants.Stack import Stack as StackConstants

#################################################################
class StackManager(Gpc1Manager):
    """
    Implementation of the stack.pro pantasks script tasks
    """
    scriptName = '/home/panstarrs/ippdor/local/bin/generate_stack_jobs.py'

    def __init__(self,
                 label,
                 recipient):
        """
        Creates a connector to the gpc1 database as well as local and
        nfs directories (their name is stack)
        """
        Gpc1Manager.__init__(self, recipient=recipient)
        self.label = label
        # Retrieve exposures names (even if they are known) and their warp_id
        self.exposures_names = []
        self.warp_ids = []
        query = """
SELECT rawExp.exp_name, warpRun.warp_id, warpRun.data_group
FROM warpRun
  JOIN fakeRun USING(fake_id)
  JOIN camRun USING(cam_id)
  JOIN chipRun USING(chip_id)
  JOIN rawExp USING(exp_id)
WHERE
  warpRun.state = 'full'
  AND warpRun.label = '%s'""" % (self.label)
        print
        print query
        print
        data = self.multiple_query(query)
        for datum in data:
            (exposure_name, warp_id, data_group) = datum
            self.exposures_names.append(exposure_name)
            self.warp_ids.append(str(warp_id))
            self.data_group = data_group
        self.stack_directory = '%s/%s/stack' % (GlobalConstants.LOCAL_TMP_DIR, self.label)
        mkdir(self.stack_directory)

    def build_create_stack_entries_job(self, stage_name = 'E010'):
        """
        Stage E010:
        stacktool -dbname gpc1 -definebyquery -set_label ${label} 
                  -set_data_group $data_group  -set_workdir $workdir
                  -warp_id <w1> -warp_id <w2> [...]
        """
        argument_warp_ids =  '-warp_id %s' % ' -warp_id '.join(self.warp_ids)
        jobs_filenames = '%s/%s.jobs' % (self.stack_directory, stage_name)
        print "Generating Condor jobs file [%s] for stacking of %d exposures (stage %s)" % (jobs_filenames, 
                                                                                            len(self.exposures_names), 
                                                                                            stage_name)
        output = open(jobs_filenames, 'w')
        output.write('Executable = /home/panstarrs/ippdor/local/bin/encapsulate_stack.sh\n')
        output.write('Universe = vanilla\n')
        output.write('Output = %s/out.%s.stack_%s\n' % (GlobalConstants.NFS_TMP_DIR,
                                                        self.label, 
                                                        stage_name))
        output.write('Error = %s/err.%s.stack_%s\n' % (GlobalConstants.NFS_TMP_DIR,
                                                       self.label,
                                                       stage_name))
        output.write('Log = %s/stack_%s.log\n' % (GlobalConstants.LOCAL_TMP_DIR,
                                                  stage_name))
        output.write('Requirements = %s\n' % StackConstants.REQUIREMENTS)
        output.write('Notification = Error\n')
        output.write('\n')
        output.write('getenv = True\n')
        output.write('\n')
        arguments = "%s -definebyquery -dbname gpc1" % GlobalConstants.STACKTOOL
        arguments = "%s -set_dist_group NODIST -min_num 2" % arguments
        arguments = "%s -set_label %s" % (arguments, self.label)
        arguments = "%s -set_data_group %s" % (arguments, self.data_group)
        arguments = "%s -set_workdir neb://@HOST.0@/gpc1/%s/stack" % (arguments, self.label)
        arguments = "%s %s -min_num 2 " % (arguments, argument_warp_ids)
        output.write('Arguments = %s\n' % arguments)
        output.write('Queue\n')
        output.write('\n')
        output.close()

    def build_new_stack_jobs(self, stage_name = 'E020'):
        """Equivalent of stacktool -tosum
        """
        query = """
SELECT
    stackRun.stack_id,
    stackRun.tess_id,
    stackRun.skycell_id,
    stackRun.workdir,
    stackRun.reduction,
    stackRun.state,
    stackSumSkyfile.path_base
FROM stackRun
JOIN stackInputSkyfile USING(stack_id)
JOIN warpRun USING(warp_id)
LEFT JOIN stackSumSkyfile USING(stack_id)
LEFT JOIN Label ON Label.label = stackRun.label
WHERE
    ((stackRun.state = 'new' AND stackSumSkyfile.stack_id IS NULL)
    OR (stackRun.state = 'update' AND stackSumSkyfile.fault = 0 AND stackSumSkyfile.quality = 0))
    AND (Label.active OR Label.active IS NULL)
    AND warpRun.label='%s'
GROUP BY stack_id
HAVING SUM(IF(warpRun.state = 'full', 1, 0)) = COUNT(stackInputSkyfile.warp_id)
""" % (self.label)
        skycells_to_be_processed = self.multiple_query(query)
        print 'Got %d jobs to generate' % len(skycells_to_be_processed)
        jobs_filenames = '%s/%s.jobs' % (self.stack_directory, stage_name)
        message = 'Generating Condor jobs file [%s] ' % (jobs_filenames)
        message = '%s for stacking of %d skycells (stage %s)' % (message,
                                                                 len(skycells_to_be_processed), 
                                                                 stage_name)
        print message
        output = open(jobs_filenames, 'w')
        output.write('Executable = /home/panstarrs/ippdor/local/bin/encapsulate_stack.sh\n')
        output.write('Universe = vanilla\n')
        output.write('Output = %s/out.%s.stack_%s\n' % (GlobalConstants.NFS_TMP_DIR,
                                                        self.label, 
                                                        stage_name))
        output.write('Error = %s/err.%s.stack_%s\n' % (GlobalConstants.NFS_TMP_DIR,
                                                       self.label, 
                                                       stage_name))
        output.write('Log = %s/stack_%s.log\n' % (GlobalConstants.LOCAL_TMP_DIR,
                                                  stage_name))
        output.write('Requirements = %s\n' % StackConstants.REQUIREMENTS)
        output.write('Notification = Error\n')
        output.write('\n')
        output.write('getenv = True\n')
        output.write('\n')
        for skycell_info in skycells_to_be_processed:
            (stack_id, tess_id, skycell_id, workdir, reduction, state, path_base) = skycell_info
            outroot = '%s/%s/%s/%s.%s.stk.%d' % (workdir, tess_id, skycell_id, tess_id, skycell_id, stack_id)
            arguments = "Arguments = %s/stack_skycell.pl --dbname gpc1 --threads 1 " % GlobalConstants.IPP_BIN_PATH
            arguments = "%s --stack_id %d" % (arguments, stack_id)
            arguments = "%s --outroot %s --redirect-output " % (arguments, outroot)
            arguments = "%s --run-state %s" % (arguments, state)
            output.write('%s\n' % arguments)
            output.write('Queue\n')
            output.write('\n')
            if reduction is not None:
                print 'reduction defined (%s) but not used' % reduction
            if path_base is not None:
                print 'path_base defined (%s) but not used' % path_base
        if len(skycells_to_be_processed) == 0:
            arguments = "Arguments = echo No skycell to be processed for label [%s]" % self.label
            output.write('%s\n' % arguments)
            output.write('Queue\n')
            output.write('\n')
        output.close()
        return

    def build_revert_stack_jobs(self, stage_name):
        """-revertsumskyfile
        """
        query = """
DELETE FROM stackSumSkyfile
USING stackSumSkyfile, stackRun
WHERE stackSumSkyfile.stack_id = stackRun.stack_id
    AND stackRun.state = 'new'
    AND fault != 0
    AND stackRun.label = '%s'
""" % (self.label)
        self.multiple_query(query)
        self.build_new_stack_jobs(stage_name)
        return

    def build_stack_summary_jobs(self):
        """E060:
        skycell_jpeg.pl --stage stack --stage_id $SASS_ID --camera $CAMERA --outroot $outroot
        """
        query = """
SELECT DISTINCT sass_id,rawExp.camera,stackRun.workdir,stackRun.tess_id,stackRun.state
       FROM stackRun
       JOIN stackInputSkyfile ON stackRun.stack_id = stackInputSkyfile.stack_id
       JOIN stackAssociationMap ON stackRun.stack_id = stackAssociationMap.stack_id
       JOIN stackAssociation USING(sass_id)
       JOIN warpRun USING(warp_id)
       JOIN fakeRun USING(fake_id)
       JOIN camRun USING(cam_id)
       JOIN chipRun USING(chip_id)
       JOIN rawExp USING(exp_id)
       LEFT JOIN stackSummary USING(sass_id)
WHERE stackRun.state = 'full' AND
      stackSummary.projection_cell IS NULL AND
      stackRun.label = '%s'
""" % (self.label)
        entries = self.multiple_query(query)
        print 'Got %d jobs to generate' % len(entries)
        stage_name = 'E060'
        jobs_filenames = '%s/%s.jobs' % (self.stack_directory, stage_name)
        print "Generating Condor jobs file [%s] for jpeging of %d skycells (stage %s)" % (jobs_filenames, 
                                                                                          len(entries), 
                                                                                          stage_name)
        output = open(jobs_filenames, 'w')
        output.write('Executable = /home/panstarrs/ippdor/local/bin/encapsulate_stack.sh\n')
        output.write('Universe = vanilla\n')
        output.write('Output = %s/out.%s.stack_%s\n' % (GlobalConstants.NFS_TMP_DIR,
                                                        self.label, 
                                                        stage_name))
        output.write('Error = %s/err.%s.stack_%s\n' % (GlobalConstants.NFS_TMP_DIR,
                                                       self.label, 
                                                       stage_name))
        output.write('Log = %s/stack_%s.log\n' % (GlobalConstants.LOCAL_TMP_DIR,
                                                  stage_name))
        output.write('Requirements = %s\n' % StackConstants.REQUIREMENTS)
        output.write('Notification = Error\n')
        output.write('\n')
        output.write('getenv = True\n')
        output.write('\n')
        for entry in entries:
            (sass_id, camera, workdir, tess_id, state) = entry
            outroot = '%s/%s/%s.stk.%d.summary' % (workdir, tess_id, tess_id, sass_id)
            arguments = 'Arguments = skycell_jpeg.pl --stage stack '
            arguments = '%s --stage_id %d --camera %s --outroot %s' % (arguments,
                                                                       sass_id,
                                                                       camera, 
                                                                       outroot)
            output.write('%s\n' % arguments)
            output.write('Queue\n')
            output.write('\n')
            if state is not None:
                print 'state defined (%s) but not used' % state
        if len(entries) == 0:
            arguments = "Arguments = echo No skycell to be jpegged for label [%s]" % self.label
            output.write('%s\n' % arguments)
            output.write('Queue\n')
            output.write('\n')
        output.close()
        return

    def build_dag(self):
        """
        DAG construction
        """
        dag_name = '%s/stack.dag' % (self.stack_directory)
        print "Generating Stack DAG file [%s] for label %s" % (dag_name,
                                                               self.label)
        output = open(dag_name, 'w')
        output.write('# Generated DAG\n')
        output.write('\n')
        output.write('JOB\tE010\t%s/E010.jobs\n' % self.stack_directory)
        output.write('SCRIPT POST E010 %s E020 %s %s\n' % (StackManager.scriptName, self.label, self.recipient))
        output.write('JOB\tE020\t%s/E020.jobs\n' % self.stack_directory)
        output.write('SCRIPT POST E020 %s E030 %s %s\n' % (StackManager.scriptName, self.label, self.recipient))
        output.write('JOB\tE030\t%s/E030.jobs\n' % self.stack_directory)
        output.write('SCRIPT POST E030 %s E040 %s %s\n' % (StackManager.scriptName, self.label, self.recipient))
        output.write('JOB\tE040\t%s/E040.jobs\n' % self.stack_directory)
        output.write('SCRIPT POST E040 %s E050 %s %s\n' % (StackManager.scriptName, self.label, self.recipient))
        output.write('JOB\tE050\t%s/E050.jobs\n' % self.stack_directory)
        output.write('SCRIPT POST E050 %s E060 %s %s\n' % (StackManager.scriptName, self.label, self.recipient))
        output.write('JOB\tE060\t%s/E060.jobs\n' % self.stack_directory)
        output.write('\n')
        output.write('PARENT E010 CHILD E020\n')
        output.write('PARENT E020 CHILD E030\n')
        output.write('PARENT E030 CHILD E040\n')
        output.write('PARENT E040 CHILD E050\n')
        output.write('PARENT E050 CHILD E060\n')
        output.write('\n')
        output.close()
        return dag_name

#################################################################
if __name__ == '__main__':
    import sys
    print '<Error>: Not intended to be used as an independent program'
    print 'Use %s' % StackManager.scriptName
    sys.exit(1)
