"""
Generate the full warp DAG for one exposure
Generate the condor job files for the warp substages
"""
from ipp.Gpc1Manager import Gpc1Manager as Gpc1Manager
from ipp.helpers import mkdir
from ipp.constants.Globals import Globals as GlobalConstants
from ipp.constants.Warp import Warp as WarpConstants

#################################################################
class WarpManager(Gpc1Manager):
    """
    Warp stage in Condor
    """
    SCRIPT_NAME = '/home/panstarrs/ippdor/local/bin/generate_warp_jobs.py'

    def __init__(self,
                 label,
                 exposure_name,
                 recipient):
        """
        Creates a connector to the gpc1 database as well as local and
        nfs directories (their name is warp)
        """
        Gpc1Manager.__init__(self, recipient=recipient)
        self.label = label
        self.exposure_name = exposure_name
        self.nfs_warp_directory = '%s' % GlobalConstants.NFS_TMP_DIR
        self.warp_directory = '%s/%s/%s/warp' % (GlobalConstants.LOCAL_TMP_DIR, 
                                                label, exposure_name)
        mkdir(self.warp_directory)

    def build_new_exp_warp_jobs(self, stage_name = 'D010'):
        """Equivalent of warptool -tooverlap and warptool -revert
        """
        query = """
SELECT
    warpRun.warp_id,
    warpRun.workdir,
    warpRun.tess_id,
    rawExp.camera,
    exp_tag
FROM warpRun
JOIN fakeRun
    USING(fake_id)
JOIN camRun
    USING(cam_id)
JOIN chipRun
    USING(chip_id)
JOIN rawExp
    USING(exp_id)
LEFT JOIN warpSkyCellMap
    USING(warp_id)
LEFT JOIN warpMask
    ON warpRun.label = warpMask.label
WHERE
    warpRun.state = 'new'
    AND fakeRun.state = 'full'
    AND camRun.state = 'full'
    AND chipRun.state = 'full'
    AND warpSkyCellMap.warp_id IS NULL
    AND warpMask.label IS NULL
    AND warpRun.label = '%s'
    AND exp_name = '%s'
""" % (self.label, self.exposure_name)
        warps = self.multiple_query(query)
        jobs_filename = '%s/%s.jobs' % (self.warp_directory, stage_name)
        message = 'Generating Condor jobs file [%s] for processing' % jobs_filename
        message = "%s %d warps (stage %s)" % (message, len(warps), stage_name)
        print message
        output = open(jobs_filename, 'w')
        output.write('Executable = /home/panstarrs/ippdor/local/bin/encapsulate_warp.sh\n')
        output.write('Universe = vanilla\n')
        output.write('Output = %s/out.%s.%s.warp_%s\n' % (self.nfs_warp_directory,
                                                          self.label, 
                                                          self.exposure_name, 
                                                          stage_name))
        output.write('Error = %s/err.%s.%s.warp_%s\n' % (self.nfs_warp_directory,
                                                         self.label, 
                                                         self.exposure_name, 
                                                         stage_name))
        output.write('Log = %s/%s_warp_%s.log\n' % (GlobalConstants.LOCAL_TMP_DIR,
                                                    self.exposure_name, 
                                                    stage_name))
        output.write('Requirements = %s\n' % WarpConstants.REQUIREMENTS)
        output.write('Notification = Error\n')
        output.write('\n')
        output.write('getenv = True\n')
        output.write('\n')
        for warp in warps:
            (warp_id, workdir, tess_dir, camera, exp_tag) = warp
            outlogfile = "%s/%s/%s.wrp.%s.log" % (workdir, exp_tag, exp_tag, warp_id)
            arguments = "Arguments = %s --warp_id %d" % (GlobalConstants.WARP_OVERLAP, warp_id)
            arguments = "%s --camera %s --tess_dir %s --logfile %s" % (arguments, camera, 
                                                                       tess_dir, outlogfile)
            output.write('%s\n' % arguments)
            output.write('Queue\n')
            output.write('\n')
        if len(warps) == 0:
            arguments = "Arguments = echo Nothing to do for stage %s" % stage_name
            output.write('%s\n' % arguments)
            output.write('Queue\n')
            output.write('\n')
        output.close()

    def build_warp_exp_revert_jobs(self, stage_name):
        """revert the -overlap exp stage
        """
        query = """
DELETE FROM warpSkyCellMap
USING warpSkyCellMap, warpRun, fakeRun, camRun, chipRun, rawExp
WHERE warpSkyCellMap.warp_id = warpRun.warp_id
    AND warpRun.fake_id = fakeRun.fake_id
    AND fakeRun.cam_id = camRun.cam_id
    AND camRun.chip_id = chipRun.chip_id
    AND chipRun.exp_id = rawExp.exp_id
    AND warpRun.state = 'new'
    AND warpSkyCellMap.fault != 0
    AND warpRun.label = '%s'
    AND rawExp.exp_name = '%s'
""" % (self.label, self.exposure_name)
        self.multiple_query(query)
        self.build_new_exp_warp_jobs(stage_name)

    def build_new_skycell_warp_jobs(self, stage_name = 'D110'):
        """Equivalent of warp.skycell.[load|run|revert] (warptool -towarp)
        Dropped the label part. We are not in the general case
        """
        query = """
SELECT DISTINCT
        warpSkyCellMap.warp_id,
        warpImfile.warp_skyfile_id,
        warpSkyCellMap.skycell_id,
        warpSkyCellMap.tess_id,
        rawExp.camera,
        warpRun.workdir,
        rawExp.exp_tag,
        warpRun.reduction,
        chipRun.magicked,
        warpRun.state
    FROM warpRun
    JOIN warpSkyCellMap
        USING(warp_id)
    JOIN warpImfile
        ON warpSkyCellMap.warp_id = warpImfile.warp_id
        AND warpSkyCellMap.skycell_id = warpImfile.skycell_id
    JOIN fakeRun
        USING(fake_id)
    JOIN camRun
        USING(cam_id)
    JOIN chipRun
        USING(chip_id)
    JOIN chipProcessedImfile
        USING(chip_id)
    JOIN rawExp
        ON chipRun.exp_id = rawExp.exp_id
    LEFT JOIN warpSkyfile
        ON warpRun.warp_id = warpSkyfile.warp_id
        AND warpSkyCellMap.skycell_id = warpSkyfile.skycell_id
        AND warpSkyCellMap.tess_id = warpSkyfile.tess_id
    LEFT JOIN warpMask
        ON warpRun.label = warpMask.label
    LEFT JOIN Label ON warpRun.label = Label.label
    WHERE
        warpRun.state = 'new'
        AND warpSkyfile.warp_id IS NULL
        AND warpSkyfile.skycell_id IS NULL
        AND warpSkyfile.tess_id IS NULL
        AND fakeRun.state = 'full'
        AND camRun.state = 'full'
        AND chipRun.state = 'full'
        AND warpMask.label IS NULL
        AND warpSkyCellMap.fault = 0
        AND (Label.active OR Label.active IS NULL)
        AND warpRun.label = '%s'
        AND rawExp.exp_name = '%s'
""" % (self.label, self.exposure_name)
        warps = self.multiple_query(query)
        jobs_filename = '%s/%s.jobs' % (self.warp_directory, stage_name)
        message ="Generating Condor jobs file [%s] for processing" % jobs_filename
        message = "%s of %d warps (stage %s)" % (message, len(warps), stage_name)
        print message
        output = open(jobs_filename, 'w')
        output.write('Executable = /home/panstarrs/ippdor/local/bin/encapsulate_warp.sh\n')
        output.write('Universe = vanilla\n')
        output.write('Output = %s/out.%s.%s.warp_%s\n' % (self.nfs_warp_directory,
                                                          self.label, 
                                                          self.exposure_name, 
                                                          stage_name))
        output.write('Error = %s/err.%s.%s.warp_%s\n' % (self.nfs_warp_directory,
                                                         self.label, 
                                                         self.exposure_name, 
                                                         stage_name))
        output.write('Log = %s/%s_warp_%s.log\n' % (self.warp_directory,
                                                    self.exposure_name, stage_name))
        output.write('Requirements = %s\n' % WarpConstants.REQUIREMENTS)
        output.write('Notification = Error\n')
        output.write('\n')
        output.write('getenv = True\n')
        output.write('\n')
        for warp in warps:
            ( warp_id, warp_skyfile_id, skycell_id, tess_dir, camera, workdir, 
              exp_tag, reduction, raw_magicked, state ) = warp
            outroot = '%s/%s/%s.wrp.%d.%s' % (workdir, exp_tag, exp_tag, warp_id, skycell_id)
            arguments = "Arguments = %s" % GlobalConstants.WARP_SKYCELL
            arguments = "%s --warp_id %d" % (arguments, warp_id)
            arguments = "%s --warp_skyfile_id %d" % (arguments, warp_skyfile_id) 
            arguments = "%s --skycell_id %s" % (arguments, skycell_id)
            arguments = "%s --tess_dir %s" % (arguments, tess_dir)
            arguments = "%s --camera %s" % (arguments,  camera)
            arguments = "%s --outroot %s" % (arguments, outroot)
            arguments = "%s --run-state %s" % (arguments, state)
            if raw_magicked is not None and raw_magicked > 0:
                arguments = '%s --magicked %d' % (arguments, raw_magicked)
            if reduction is not None and reduction != 'NULL':
                arguments = '%s --reduction %s' % (arguments, reduction)
            output.write('%s\n' % arguments)
            output.write('Queue\n')
            output.write('\n')
        if len(warps) == 0:
            arguments = 'Arguments = echo Nothing to do for stage %s' % stage_name
            output.write('%s\n' % arguments)
            output.write('Queue\n')
            output.write('\n')
        output.close()

    def build_warp_skycell_revert_jobs(self, stage_name):
        """Revert stage D110
        """
        query = """
DELETE FROM warpSkyfile
USING warpSkyfile, warpRun, fakeRun, camRun, chipRun, rawExp
WHERE warpSkyfile.warp_id = warpRun.warp_id
    AND warpRun.fake_id = fakeRun.fake_id
    AND fakeRun.cam_id = camRun.cam_id
    AND camRun.chip_id = chipRun.chip_id
    AND chipRun.exp_id = rawExp.exp_id
    AND warpRun.state = 'new'
    AND warpSkyfile.fault != 0
    AND warpRun.label = '%s'
    AND rawExp.exp_name = '%s'
""" % (self.label, self.exposure_name)
        self.multiple_query(query)
        self.build_new_skycell_warp_jobs(stage_name)

    def build_new_summary_warp_jobs(self):
        """Stage D210
        """
        query = """
SELECT DISTINCT warp_id,rawExp.camera,warpRun.workdir,warpRun.tess_id,rawExp.exp_tag,warpRun.state
       FROM warpRun 
       JOIN fakeRun USING(fake_id) JOIN camRun USING(cam_id)
       JOIN chipRun USING(chip_id) JOIN rawExp USING(exp_id)
       LEFT JOIN warpSummary USING(warp_id)
WHERE warpRun.state = 'full' 
  AND warpSummary.projection_cell IS NULL
  AND warpRun.label = '%s'
  AND rawExp.exp_name = '%s'
""" % (self.label, self.exposure_name)
        summaries = self.multiple_query(query)
        jobs_filename = '%s/%s.jobs' % (self.warp_directory, 'D210')
        message = "Generating Condor jobs file [%s] for processing of" % jobs_filename
        message = "%s %d warp summaries (stage D210)" % (message, len(summaries))
        print message
        output = open(jobs_filename, 'w')
        output.write('Executable = /home/panstarrs/ippdor/local/bin/encapsulate_warp.sh\n')
        output.write('Universe = vanilla\n')
        output.write('Output = %s/out.%s.%s.warp_D210\n' % (self.nfs_warp_directory,
                                                            self.label, self.exposure_name))
        output.write('Error = %s/err.%s.%s.warp_D210\n' % (self.nfs_warp_directory,
                                                           self.label, self.exposure_name))
        output.write('Log = %s/%s_warp_D210.log\n' % (self.warp_directory,
                                                      self.exposure_name))
        output.write('Requirements = %s\n' % WarpConstants.REQUIREMENTS)
        output.write('Notification = Error\n')
        output.write('\n')
        output.write('getenv = True\n')
        output.write('\n')
        for summary in summaries:
            (warp_id, camera, workdir, tess_id, exp_tag, state) = summary
            outroot = '%s/%s/%s.wrp.%d.%s' % (workdir, exp_tag, exp_tag, warp_id, tess_id)
            arguments = 'Arguments = %s -stage_id %d -camera %s -outroot %s' % (GlobalConstants.WARP_JPEG,
                                                                                warp_id, camera, outroot)
            output.write('%s\n' % arguments)
            output.write('Queue\n')
            output.write('\n')
            if state is not None:
                print 'state is defined but not used: %s' % state
        if len(summaries) == 0:
            arguments = 'Arguments = echo Nothing to do for stage D210'
            output.write('%s\n' % arguments)
            output.write('Queue\n')
            output.write('\n')
        output.close()

    def advance_warp(self):
        """Stage 310 (last). Create the job to 'warptool -advancerun -warp_id -label'
        """
        jobs_filename = '%s/%s.jobs' % (self.warp_directory, 'D310')
        print "Generating Condor jobs file [%s] for advancing of warp related to %s (stage D310)" % (jobs_filename, 
                                                                                                     self.exposure_name)
        output = open(jobs_filename, 'w')
        output.write('Executable = /home/panstarrs/ippdor/local/bin/encapsulate_warp.sh\n')
        output.write('Universe = vanilla\n')
        output.write('Output = %s/out.%s.%s.warp_D210\n' % (self.nfs_warp_directory,
                                                            self.label, self.exposure_name))
        output.write('Error = %s/err.%s.%s.warp_D210\n' % (self.nfs_warp_directory,
                                                           self.label, self.exposure_name))
        output.write('Log = %s/%s_warp_D210.log\n' % (self.warp_directory,
                                                      self.exposure_name))
        output.write('Requirements = %s\n' % WarpConstants.REQUIREMENTS)
        output.write('Notification = Error\n')
        output.write('\n')
        output.write('getenv = True\n')
        output.write('\n')
        arguments = '%s -advancerun -label %s -warp_id %d' % (GlobalConstants.WARPTOOL, 
                                                              self.label, 
                                                              self.get_warp_id(self.exposure_name, 
                                                                               self.label))
        output.write('Arguments = %s\n' % arguments)
        output.write('Queue\n')
        output.write('\n')
        output.close()

    def build_warp_dag(self):
        """Build the warp stage DAG
        """
        dag_name = '%s/warp.dag' % (self.warp_directory)
        print "Generating Warp DAG file [%s] for exposure %s (label %s)" % (dag_name,
                                                                            self.exposure_name,
                                                                            self.label)
        output = open(dag_name, 'w')
        output.write('# Generated DAG\n')
        output.write('\n')
        output.write('# Exp\n')
        output.write('JOB\tD010\t%s/D010.jobs\n' % self.warp_directory)
        output.write('SCRIPT POST D010 %s D020 %s %s\n' % (WarpManager.SCRIPT_NAME,
                                                           self.label, self.exposure_name))
        output.write('JOB\tD020\t%s/D020.jobs\n' % self.warp_directory)
        output.write('SCRIPT POST D020 %s D030 %s %s\n' % (WarpManager.SCRIPT_NAME, 
                                                           self.label, self.exposure_name))
        output.write('JOB\tD030\t%s/D030.jobs\n' % self.warp_directory)
        output.write('SCRIPT POST D030 %s D040 %s %s\n' % (WarpManager.SCRIPT_NAME, 
                                                           self.label, self.exposure_name))
        output.write('JOB\tD040\t%s/D040.jobs\n' % self.warp_directory)
        output.write('SCRIPT POST D040 %s D110 %s %s\n' % (WarpManager.SCRIPT_NAME, 
                                                           self.label, self.exposure_name))
        output.write('\n')
        output.write('# Skycell\n')
        output.write('JOB\tD110\t%s/D110.jobs\n' % self.warp_directory)
        output.write('SCRIPT POST D110 %s D120 %s %s\n' % (WarpManager.SCRIPT_NAME, 
                                                           self.label, self.exposure_name))
        output.write('JOB\tD120\t%s/D120.jobs\n' % self.warp_directory)
        output.write('SCRIPT POST D120 %s D130 %s %s\n' % (WarpManager.SCRIPT_NAME, 
                                                           self.label, self.exposure_name))
        output.write('JOB\tD130\t%s/D130.jobs\n' % self.warp_directory)
        output.write('SCRIPT POST D130 %s D140 %s %s\n' % (WarpManager.SCRIPT_NAME, 
                                                           self.label, self.exposure_name))
        output.write('JOB\tD140\t%s/D140.jobs\n' % self.warp_directory)
        output.write('SCRIPT POST D140 %s D210 %s %s\n' % (WarpManager.SCRIPT_NAME, 
                                                           self.label, self.exposure_name))
        output.write('\n')
        output.write('# Summary\n')
        output.write('JOB\tD210\t%s/D210.jobs\n' % self.warp_directory)
        output.write('SCRIPT POST D210 %s D310 %s %s\n' % (WarpManager.SCRIPT_NAME, 
                                                           self.label, self.exposure_name))
        output.write('\n')
        output.write('# Advance\n')
        output.write('JOB\tD310\t%s/D310.jobs\n' % self.warp_directory)
        output.write('\n')
        output.write('PARENT D010 CHILD D020\n')
        output.write('PARENT D020 CHILD D030\n')
        output.write('PARENT D030 CHILD D040\n')
        output.write('PARENT D040 CHILD D110\n')
        output.write('\n')
        output.write('PARENT D110 CHILD D120\n')
        output.write('PARENT D120 CHILD D130\n')
        output.write('PARENT D130 CHILD D140\n')
        output.write('PARENT D140 CHILD D210\n')
        output.write('\n')
        output.write('PARENT D210 CHILD D310\n')
        output.write('\n')
        output.close()
        return dag_name

#################################################################
if __name__ == '__main__':
    import sys
    print 'Error: Use %s' % WarpManager.SCRIPT_NAME
    sys.exit(0)
