#!/usr/bin/env python

import sys
import os
import subprocess
from ipp.Gpc1Manager import Gpc1Manager as Gpc1Manager
from ipp.StackManager import StackManager as StackManager
from ipp.helpers import mkdir
from ipp.IppException import IppException as IppException

class StackBuilder(Gpc1Manager):
    def __init__(self,
                 label,
                 recipient,
                 host='ippdb01.ifa.hawaii.edu', 
                 user='ipp', 
                 passwd='ipp'):
        """Create a connector to the gpc1 database
        """
        Gpc1Manager.__init__(self, recipient=recipient, host=host, user=user, passwd=passwd)
        self.label = label
        if self.isLabelUsed():
            print 'Label [%s] already used' % label
            print 'IF YOU KNOW WHAT YOU ARE DOING, delete all chips with label %s in chipRun:' % label
            print "  DELETE FROM chipImfile USING chipImfile, chipRun WHERE chipImfile.chip_id=chipRun.chip_id AND label = '%s';" % label
            print "  DELETE FROM chipRun WHERE label = '%s';" % label
            print 'OR'
            import datetime
            print 'Use another label (e.g. %s_%s)' % (label, datetime.datetime.now().strftime("%Y%m%dT%H%M%S"))
            raise IppException('Label [%s] already used' % label)
        self.rootDirectory = '%s/%s' % (Constants.ippdorLocalTmp, label)
        mkdir(self.rootDirectory)

    def isLabelUsed(self):
        query = """SELECT label from chipRun WHERE label = '%s' LIMIT 1""" % self.label
        values = self.multipleQuery(query)
        return values is not None and len(values) != 0

    def isExposureNameValid(self, exposureName):
        query = """SELECT exp_id FROM rawExp WHERE exp_name='%s'""" % exposureName
        return int(self.simpleQuery(query)[0])

    def buildCondorFiles(self, exposureNames):
        # Create DAGs for each exposure
        for exposureName in exposureNames:
            print '\tCreating DAG for exposure [%s]' % (exposureName)
            command = "/home/panstarrs/ippdor/local/bin/generate_exposure_dag.py %s %s %s" % (self.label, 
                                                                                              self.recipient, 
                                                                                              exposureName)
            try:
                subprocess.call(command.split(' '))
            except OSError:
                print 'Cannot find /home/panstarrs/ippdor/local/bin/generate_exposure_dag.py, Trying with anonymous generate_exposure_dag.py'
                command = "generate_exposure_dag.py %s %s %s" % (self.label, 
                                                                 self.recipient,
                                                                 exposureName)
                subprocess.call(command.split(' '))
        # Generate stack generation job file
        self.stackGenerationJob()
        # Generate configuration file
        self.generateConfiguration()
        # Generate DAG of DAGs
        dagOfdagsName = self.generateDAG(exposureNames)
        return dagOfdagsName

    def generateConfiguration(self):
        # Generate configuration file
        output = open('%s/configuration' % self.rootDirectory, 'w')
        output.write('# DAGMAN_ABORT_ON_SCARY_SUBMIT = FALSE\n')
        output.close()

    def stackGenerationJob(self):
        output = open('%s/stack_generation.job' % self.rootDirectory, 'w')
        output.write('Executable = %s\n' % StackManager.scriptName)
        output.write('Universe = vanilla\n')
        output.write('Output = /home/panstarrs/ippdor/tmp/02test/out.%s.stack_generation\n' % (self.label))
        output.write('Error = /home/panstarrs/ippdor/tmp/02test/err.%s.stack\n' % (self.label))
        output.write('Log = /tmp/stack_generation.log\n')
        output.write('Requirements = %s\n' % Constants.REQUIREMENTS)
        output.write('\n')
        output.write('getenv = True\n')
        output.write('\n')
        output.write('Arguments = DAG %s %s\n' % (self.label, self.recipient))
        output.write('Queue\n')
        output.write('\n')
        output.close()

    def generateDAG(self, exposureNames):
        dagName = '%s/ch2st.dag' % self.rootDirectory
        output = open(dagName, 'w')
        output.write('# Chip to warp DAGs\n')
        chip2warpNodesList = ''
        for exposureName in exposureNames:
            output.write('SUBDAG EXTERNAL CHIP2WARP_%s %s/%s/%s.dag\n' % (exposureName, 
                                                                          self.rootDirectory,
                                                                          exposureName,
                                                                          exposureName))
            chip2warpNodesList = '%s CHIP2WARP_%s' % (chip2warpNodesList, exposureName)
        output.write('\n')
        output.write('# Generate Stack DAGs\n')
        output.write('JOB GENERATE_STACK %s/stack_generation.job\n' % self.rootDirectory)
        output.write('\n')
        output.write('# Stack DAGs\n')
        output.write('SUBDAG EXTERNAL STACK %s/stack/stack.dag\n' % (self.rootDirectory))
        output.write('\n')
        output.write('# Node Dependencies\n')
        output.write('PARENT %s CHILD GENERATE_STACK\n' % chip2warpNodesList)
        output.write('PARENT GENERATE_STACK CHILD STACK\n')
        output.write('\n')
        output.write('# Retries\n')
        for exposureName in exposureNames:
            output.write('# RETRY CHIP2WARP_%s 3\n' % exposureName)
        output.write('# RETRY GENERATE_STACK 3\n')
        output.write('# RETRY STACK 3\n')        
        output.write('\n')
        output.close()
        return dagName

def loadInputFile(fileName):
    """File format is the following:
label    = PreferredLabel
tessId   = MD08.V2
exposure = exposure1
exposure = exposure2
exposure = exposure3
exposure = exposure4
# Comment line
exposure = exposure5    # Followed by a comment
"""
    import re
    exposureNames = []
    inFile = open(fileName)
    labelNotDefined = True
    tessIdNotDefined = True
    recipientNotDefined = True
    for originalLine in inFile:
        line = originalLine[:-1]
        line = re.sub('#.*', '', line)
        line = re.sub('\s+', '', line)
        #print 'Line after replacements: [%s]' % line
        if line != '':
            elements = line.split('=')
            if len(elements) != 2:
                print 'Cannot interpret line [%s]' % originalLine[:-1]
                print 'Check your file format:'
                print loadInputFile.__doc__
                raise IppException('Error: Input file format [%s]' % fileName)
            else:
                if elements[0] == 'label':
                    if labelNotDefined:
                        label = elements[1]
                        labelNotDefined = False
                    else:
                        raise IppException('label defined twice in [%s]' % fileName)
                elif elements[0] == 'tessId':
                    if tessIdNotDefined:
                        tessId = elements[1]
                        tessIdNotDefined = False
                    else:
                        raise IppException('tessId defined twice in [%s]' % fileName)
                elif elements[0] == 'exposure':
                    exposureNames.append(elements[1])
                elif elements[0] == 'recipient':
                    if recipientNotDefined:
                        recipient = elements[1]
                        recipientNotDefined = False
                    else:
                        raise IppException('recipient defined twice in [%s]' % fileName)
                else:
                    print 'Unknown key [%s] in [%s]' % (elements[0], fileName)
                    print loadInputFile.__doc__
                    raise IppException('Unknown key [%s] in [%s]' % (elements[0], fileName))
    return (label, tessId, exposureNames, recipient)

if __name__ == '__main__':
    if len(sys.argv)<2 or sys.argv[1].startswith('-h') or sys.argv[1].startswith('--h'):
        print '<Usage>: %s <label> <tessellation> <email recipient> [exposures list]' % sys.argv[0]
        print '         %s -f <input_file>' % sys.argv[0]
        print '   If <label> has already been used, the program will abort'
        print '   [exposures list] must have more than 2 exposures'
        print
        print loadInputFile.__doc__
        sys.exit(1)
    #TODO: Implement proper arguments management
    if sys.argv[1] == '-f':
        (label, tessellation, exposureNames, recipient) = loadInputFile(sys.argv[2])
    else:
        argStart = 1
        label = sys.argv[argStart]
        argStart += 1
        tessellation = sys.argv[argStart]
        argStart += 1
        recipient = sys.argv[argStart]
        argStart += 1
        exposureNames = sys.argv[argStart:]
        argStart += 1
    if len(exposureNames) < 2:
        print 'You need more than 2 exposures to build a stack'
        print 'Aborting'
        sys.exit(4)
    if os.getenv('USER') == 'ippdor':
        builder = StackBuilder(label, recipient)
    else:
        builder = StackBuilder(label, recipient, host='ipp0012.ifa.hawaii.edu')
    # Check exposures        
    expIds = dict()
    for exposureName in exposureNames:
        print 'Checking exposure [%s]' % exposureName
        try:
            expIds[exposureName] = builder.isExposureNameValid(exposureName)
        except IndexError:
            print 'Exposure [%s] is not valid' % exposureName
            print 'Aborting'
            sys.exit(3)
    # Create entries in chiptool
    print 'Stacking %d exposures with label [%s]' % (len(exposureNames), label)
    for exposureName in exposureNames:
        expId = expIds[exposureName]
        print '\tCreating chipRun entry for exposure [%s]/%d' % (exposureName, expId)
        command = "/home/panstarrs/ippdor/psconfig/20110920/20110920.lin64/bin/chiptool -dbname gpc1 -definebyquery -exp_id %d" % expId
        command = "%s -set_label %s" % (command, label)
        command = "%s -set_workdir neb://\@HOST\@.0/gpc1/%s" % (command, label)
        command = "%s -set_tess_id %s" % (command, tessellation)
        command = "%s -set_data_group %s_dg" % (command, label)
        command = "%s -set_note \'chip/stack_reprocessing\' -set_end_stage warp" % (command)
        try:
            subprocess.call(command.split(' '))
        except OSError:
            print 'System call to [%s] failed. Continuing though...' % command
    # Generate Condor jobs, subdags and dag
    dagOfDagsName = builder.buildCondorFiles(exposureNames)
    # Show the command to be run
    print 'Now run:'
    print 'condor_submit_dag -maxidle 1000 -maxjobs 500 %s' % dagOfDagsName
