#!/usr/bin/env jython

import stilts
import datetime
import re

from java.lang import *
from java.sql import *


class Batch(object):

    driverName="com.mysql.jdbc.Driver"

    '''
    Constructor
    '''
    def __init__(self, batchType, inputFitsPath, outputFitsPath, dbHost, dbName, dbUser, dbPass, survey=""):

        # set up class variables
        self.pspsVoTableFilePath = "../config/" + batchType + "/tables.vot"
        self.inputFitsPath = inputFitsPath
        self.outputFitsPath = outputFitsPath
        self.dbHost = dbHost
        self.dbName = dbName
        self.dbUser = dbUser
        self.dbPass = dbPass
        self.survey = survey

        # set up JDBC connection
        self.url = "jdbc:mysql://"+self.dbHost+"/"+self.dbName+"?user="+self.dbUser+"&password="+self.dbPass
        self.con = DriverManager.getConnection(self.url)
        self.stmt = self.con.createStatement()

        # get survey ID from init table
        sql = "SELECT surveyID from Survey WHERE name = '" + survey + "'"
        try:
            rs = self.stmt.executeQuery(sql)  
            rs.first()
            self.surveyID = rs.getInt(1)
        except:
            self.log("No survey ID found for this survey: '" + survey + "'")
            self.surveyID = -1;
    
        # store today's date
        now = datetime.datetime.now();
        self.dateStr = now.strftime("%Y-%m-%d")

        if self.inputFitsPath != "": 
            self.parseFitsHeader()

    '''
    Destructor
    '''
    def __del__(self):

        self.log("Batch destructor")
        self.stmt.close()
        self.con.close()

    '''
    Prints a log message with the current time
    '''
    def log(self, msg):

        print datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + " | " + msg

    '''
    Updates a table with surveyID
    '''
    def updateSurveyID(self, table):

        sql = "UPDATE " + table + "  SET surveyID=%d" % self.surveyID
        self.stmt.execute(sql)

    '''
    Updates a table with filterID grabbed from Filter init table
    '''
    def updateFilterID(self, table):

        sql = "UPDATE "+table+" AS a, Filter AS b SET a.filterID=b.filterID WHERE b.filterType = '" + self.filter + "'"
        self.stmt.execute(sql)

    '''
    Reads FITS header and stores all fields in a dictionary object
    '''
    def parseFitsHeader(self):

        fitsFile = open(self.inputFitsPath)

        self.header = {}

        while (True):
           record = fitsFile.read(80)

           # quit when we reach 'END'
           if record.startswith("END"): break

           # ignore comments
           if record.startswith("COMMENT"): continue
           match = re.match('(.*)=(.*)', record)
           if match:

               # remove HIERARCH prefix
               param = match.group(1).replace("HIERARCH", "")
               param = param.strip()

               value = match.group(2)
               # remove trailing comment after / char, if there is one
               index = value.find("/")
               if index != -1: value = value[0:index]

               # remove ' chars around content
               value = value.replace("'", "")

               # remove leading and trailing whitespace
               value = value.strip()

               # store in out dictionary object
               self.header[param] = value

               #print param + "|" + value + "|"

    '''
    Creates PSPS tables from VOTable descriptions
    '''   
    def createEmptyPspsTables(self):

         self.pspsTables = stilts.treads(self.pspsVoTableFilePath)
         for table in self.pspsTables:
             self.log("Creating PSPS table: " + table.name)
             table.write(self.url + '#' + table.name)

         self.indexPspsTables()

    '''
    Subclass should implement this to index IPP tables
    '''   
    def indexIppTables(self):
        self.log("indexIppTables not implemented")


    '''
    Adds an index to the supplied table and column
    '''
    def createIndex(self, table, column):

        self.log("Creating index on column '"+column+"' for table '"+table+"'")

        sql = "CREATE INDEX "+table+"_index ON "+table+" ("+column+")"
        try:
            self.stmt.execute(sql)
        except:
            self.log("Index already in place on '" + column + "' for table '" + table + "'")


    '''
    Subclass should implement this to index PSPS tables
    '''   
    def indexPspsTables(self):
        self.log("indexPspsTables not implemented")

    '''
    Imports IPP tables from FITS file

    Accepts a regular expression filter so not all tabls need to be imported
    '''
    def importIppTables(self, filter):

      tables = stilts.treads(self.inputFitsPath)

      count = 0
      for table in tables:

          match = re.match(filter, table.name)
          if not match: continue
          self.log("Creating IPP table " + table.name)
          table = stilts.tpipe(table, cmd='explodeall')
          try:
              table.write(self.url + '#' + table.name)
          except:
              self.log("ERROR problem writing table '" + table.name + "' to the database")

          count = count + 1

      self.log("Imported %d tables from '%s' " % (count, self.inputFitsPath))

      self.indexIppTables()

    '''
    Exports PSPS tables from the database to FITS format
    '''    
    def exportPspsTablesToFits(self):

        self.log("Exporting all PSPS tables to FITS")
        _tables = []

        self.log("    Selecting database tables")
        for table in self.pspsTables:
           _table = stilts.tread(self.url + '#SELECT * FROM ' + table.name)
           _table = stilts.tpipe(_table, cmd='tablename ' + table.name)
           _tables.append(_table)

        self.log("    Writing to FITS file " + self.outputFitsPath)
        stilts.twrites(_tables, self.outputFitsPath, fmt='fits')


    '''
    Main processing method - must be implemented by subclasses
    '''
    def populatePspsTables(self):
        self.log("Not implemented yet")

