#!/usr/bin/env python

"""
This is the script that generates:
- DAG for each stage from chip to warp for each exposure 
- DAG of all these DAGs
"""

import sys
import subprocess
import datetime

from ipp.Gpc1Manager import Gpc1Manager as Gpc1Manager
from ipp.helpers import mkdir
from ipp.constants.Globals import Globals as GlobalConstants
from ipp.exceptions.IppException import IppException
from ipp.exceptions.UninstantiableIppException import UninstantiableIppException

class Constants:
    """Constants which are used only in this script
    """
    def __init__(self):
        raise UninstantiableIppException()

    # In LABEL_ALREADY_USED, replace:
    # - @LABEL@ with label value 
    LABEL_ALREADY_USED = """
Error: Label [@LABEL@] already used.

To fix the problem:
1) Either use another label (e.g. @LABEL@_%s)
2) Or __IF YOU KNOW WHAT YOU ARE DOING__, delete all chips with label [@LABEL@] in chipRun:
  DELETE FROM chipImfile USING chipImfile, chipRun WHERE chipImfile.chip_id=chipRun.chip_id AND label = '@LABEL@';
  DELETE FROM chipRun WHERE label = '@LABEL@';
""" % (datetime.datetime.now().strftime("%Y%m%dT%H%M%S"))
    # Format string for the DAG of DAGs:
    # - #1: target directory
    DAG_OF_DAGS_FMT = "%s/ch2wp.dag"

class ChipWarpBuilder(Gpc1Manager):
    """Main class for DAG/Condor job files building
    """
    def __init__(self,
                 label,
                 recipient):
        """
        The constructor:
        - creates a connector to the gpc1 database
        - checks if the label has already been used
        - creates a local directory
        """
        Gpc1Manager.__init__(self, recipient=recipient)
        self.label = label
        if self.is_label_used():
            print Constants.LABEL_ALREADY_USED.replace('@LABEL@', self.label)
            raise IppException('Label [%s] already used' % label)
        self.root_directory = '%s/%s' % (GlobalConstants.LOCAL_TMP_DIR, label)
        mkdir(self.root_directory)

    def build_condor_files(self, exposures_names):
        """Generate all condor subdag files
        """
        # Create DAGs for each exposure
        for exposure_name in exposures_names:
            print '\tCreating DAG for exposure [%s]' % (exposure_name)
            command = GlobalConstants.GENERATE_EXPOSURE_DAG % (self.label, 
                                                               self.recipient, 
                                                               exposure_name)
            try:
                subprocess.call(command.split(' '))
            except OSError:
                print 'Cannot find %s' % GlobalConstants.GENERATE_EXPOSURE_DAG
                command = "generate_exposure_dag.py %s %s %s" % (self.label, 
                                                                 self.recipient,
                                                                 exposure_name)
                subprocess.call(command.split(' '))
        # Generate completion job file
        self.completion_job()
        # Generate configuration file
        self.generate_configuration_file()
        # Generate DAG of DAGs
        dag_of_dags_name = self.generate_dag(exposures_names)
        return dag_of_dags_name

    def generate_configuration_file(self):
        """Method generating the condor configuration file (empty now)
        """
        output = open('%s/configuration' % self.root_directory, 'w')
        output.write('# File empty for now\n')
        output.close()

    def completion_job(self):
        """Method generating the condor job file for completion
        """
        output = open('%s/chip_to_warp_completion.job' % self.root_directory, 'w')
        output.write('Executable = %s/send_email.py\n' % GlobalConstants.IPPDOR_BINARY_PATH)
        output.write('Universe = vanilla\n')
        output.write('Output = %s/%s.out.ch2wp.completion\n' % (GlobalConstants.NFS_TMP_DIR, 
                                                                self.label))
        output.write('Error = %s/%s.err.ch2wp.completion\n' % (GlobalConstants.NFS_TMP_DIR, 
                                                               self.label))
        output.write('Log = %s/%s_ch2wp.completion.log\n' % (GlobalConstants.LOCAL_TMP_DIR, 
                                                             self.label))
        output.write('Requirements = %s\n' % (GlobalConstants.REQUIREMENTS))
        output.write('Notification = Error\n')
        output.write('\n')
        output.write('getenv = True\n')
        output.write('\n')
        output.write('Arguments =  %s %s %s/%s_final_email.txt\n' % (self.label, 
                                                                     self.recipient,
                                                                     GlobalConstants.NFS_TMP_DIR,
                                                                     self.label))
        output.write('Queue\n')
        output.write('\n')
        output.close()
        # E-mail content
        output = open('%s/%s_final_email.txt' % (GlobalConstants.NFS_TMP_DIR,
                                                 self.label), 'w')
        output.write('Condor DAG complete: label [%s]\n' % self.label)
        output.close()

    def generate_dag(self, exposures_names):
        """Method generating the DAG of DAGs
        """
        dag_name = Constants.DAG_OF_DAGS_FMT % self.root_directory
        output = open(dag_name, 'w')
        output.write('# Chip to warp DAGs\n')
        chip2warp_nodes_list = ''
        for exposure_name in exposures_names:
            output.write('SUBDAG EXTERNAL CHIP2WARP_%s %s/%s/%s.dag\n' % (exposure_name, 
                                                                          self.root_directory,
                                                                          exposure_name,
                                                                          exposure_name))
            chip2warp_nodes_list = '%s CHIP2WARP_%s' % (chip2warp_nodes_list, exposure_name)
        output.write('\n')
        output.write('# Send an e-mail when done\n')
        output.write('JOB COMPLETION_EMAIL %s/chip_to_warp_completion.job\n' % self.root_directory)
        output.write('\n')
        output.write('# Node Dependencies\n')
        output.write('PARENT %s CHILD COMPLETION_EMAIL\n' % chip2warp_nodes_list)
        output.write('\n')
        output.close()
        return dag_name

def load_input_file(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
    exposures_names = []
    input_file = open(filename)
    label_not_defined = True
    tessellation_not_defined = True
    recipient_not_defined = True
    for original_line in input_file:
        line = original_line[:-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]' % original_line[:-1]
                print 'Check your file format:'
                print load_input_file.__doc__
                raise IppException('Error: Input file format [%s]' % filename)
            else:
                if elements[0] == 'label':
                    if label_not_defined:
                        label = elements[1]
                        label_not_defined = False
                    else:
                        raise IppException('label defined twice in [%s]' % filename)
                elif elements[0] == 'tessId':
                    if tessellation_not_defined:
                        tessellation_id = elements[1]
                        tessellation_not_defined = False
                    else:
                        raise IppException('tessId defined twice in [%s]' % filename)
                elif elements[0] == 'exposure':
                    exposures_names.append(elements[1])
                elif elements[0] == 'recipient':
                    if recipient_not_defined:
                        recipient = elements[1]
                        recipient_not_defined = False
                    else:
                        raise IppException('recipient defined twice in [%s]' % filename)
                else:
                    print 'Unknown key [%s] in [%s]' % (elements[0], filename)
                    print load_input_file.__doc__
                    raise IppException('Unknown key [%s] in [%s]' % (elements[0], filename))
    input_file.close()
    return (label, tessellation_id, exposures_names, recipient)

#####################################################################################################
if __name__ == '__main__':
    #Implement proper arguments management?
    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 load_input_file.__doc__
        sys.exit(1)
    if sys.argv[1] == '-f':
        (LABEL, TESSELLATION_ID, EXPOSURES_NAMES, RECIPIENT) = load_input_file(sys.argv[2])
    else:
        ARG_START = 1
        LABEL = sys.argv[ARG_START]
        ARG_START += 1
        TESSELLATION_ID = sys.argv[ARG_START]
        ARG_START += 1
        RECIPIENT = sys.argv[ARG_START]
        ARG_START += 1
        EXPOSURES_NAMES = sys.argv[ARG_START:]
        ARG_START += 1
    if len(EXPOSURES_NAMES) < 2:
        print 'You need more than 2 exposures to build a stack'
        print 'Aborting'
        sys.exit(4)
    BUILDER = ChipWarpBuilder(LABEL, RECIPIENT)
    # Check exposures
    EXPOSURES_IDS = dict()
    for _exposure_name in EXPOSURES_NAMES:
        print 'Checking exposure [%s]' % _exposure_name
        try:
            EXPOSURES_IDS[_exposure_name] = BUILDER.is_exposure_name_valid(_exposure_name)
        except IndexError:
            print 'Exposure [%s] is not valid' % _exposure_name
            print 'Aborting'
            sys.exit(3)
    # Generate Condor jobs, subdags and dag
    DAG_OF_DAGS_NAME = BUILDER.build_condor_files(EXPOSURES_NAMES)
    # Create entries in chiptool
    print 'Stacking %d exposures with label [%s]' % (len(EXPOSURES_NAMES), LABEL)
    for _exposure_name in EXPOSURES_NAMES:
        expId = EXPOSURES_IDS[_exposure_name]
        print '\tCreating chipRun entry for exposure [%s]/%d' % (_exposure_name, expId)
        COMMAND = "%s -definebyquery -exp_id %d" % (GlobalConstants.CHIPTOOL, 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_ID)
        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.' % COMMAND
            print '... Continuing though (assuming you know what you are doing)'
    # Show the command to be run
    print 'Now run:'
    print 'condor_submit_dag -maxidle 1000 -maxjobs 500 %s' % DAG_OF_DAGS_NAME
