#!/usr/bin/env jython

import stilts
import datetime
import re
import sys
import os
import logging
import glob
from subprocess import call, PIPE, Popen

from pslogger import PSLogger
from scratchdb import ScratchDb

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

'''
Abstract base-class for DVO classes

Provides generic functionality to pull relevant stuff from a DVO database into a scratch MySQL database while also providing access queries
'''
class Dvo(object):

    '''
    Constructor

    - sets up logging
    - checks we are using the correct DVO
    - checks we have up-to-date Images.dat
    - checks we have up-to-date SkyTable.fits

    '''
    def __init__(self, logger, config, skychunk, ippToPspsDb, scratchDbName=None):

        # set up logging
        self.logger = logger
        self.config = config
        self.skychunk = skychunk
        self.ippToPspsDb = ippToPspsDb

        # set up empty lists
        self.regionsToIngest = []
        self.regionsToPurge = []
        self.allPopulatedRegions = []
        self.regionsAlreadyIngested = []
        self.regionsIngestedButOutOfDate = []
       
        # connect to the specified scratch database, if a name is given, otherwise use the first version
        try:
            self.scratchDb = ScratchDb(self.logger, self.config, '1', scratchDbName)
        except: raise

        # we have (still) 3 modes of ingest from dvo cpt/cps/cpm files 
        # 1) dvograbber (useFullTables == FALSE), 
        # 2) jython/stilts (useStilts == TRUE, useFullTables == TRUE)
        # 3) dvopsps (useStilts == FALSE, useFullTables == TRUE)
        # my goal (EAM, 2013/01/31) is to migrate to only dvopsps

        self.useStilts = 0
        self.scratchDb.setUseFullTables(True)

        # decide if we are using the right DVO for this scratchDb
        self.correctDvo = self.scratchDb.isCorrectDvo(self.skychunk.dvoLocation)

        # this scratch Db either has no DVO info ingested, or a different DVO ingested. Prompt user
        if not self.correctDvo:
            print "*******************************************************************************"
            self.logger.warning("Wrong or no DVO in use. Do you want to reset and use '" + self.skychunk.dvoLocation + "'?")
            response = raw_input("(y/n) ")
            if response == "y":
                self.logger.warning("Are you ABSOLUTELY sure you want to do this?")
                response = raw_input("(y/n) ")
                if response == "y":
                    self.resetAllTables()


    '''
    Destructor
    '''
    def __del__(self):
        self.logger.debug("Dvo destructor")

    '''
    Check we are importing the DVO currenly ingested
    '''
    def isThisDvoCurrentlyIngested(self):
        return self.correctDvo

    '''
    Resets all dvo tables. Be very careful....
    '''
    def resetAllTables(self):

        self.logger.infoPair("Resetting", "all tables")
        if not self.scratchDb.resetAllDvoTables(): return
        self.correctDvo = True
        if not self.loadSkyTable():
            self.correctDvo = False
            return

    '''
    Loads the DVO Photcodes.dat FITS file
    '''
    def loadPhotcodes(self):

        # go no further if we've already partly ingested a different DVO
        if not self.correctDvo: return

        # check if we have up-to-date version
        path = self.skychunk.dvoLocation + "/Photcodes.dat"
        if self.scratchDb.alreadyImportedThisDvoTable(path): 
            self.logger.debugPair("DVO Photcodes.dat file", "up-to-date")
            return
        
        self.logger.infoSeparator()
        self.logger.infoPair("DVO Photcodes.dat file", "NOT up-to-date")

        self.importPhotcodes(
                path,
                self.scratchDb.dvoPhotcodesTable)
        self.logger.infoPair("Adding primary key to", self.scratchDb.dvoPhotcodesTable)
        self.scratchDb.makeColumnPrimaryKey(self.scratchDb.dvoPhotcodesTable, "CODE")

        self.scratchDb.setImportedThisDvoTable(path)

    '''
    Loads the DVO Images.dat FITS file
    '''
    def loadImages(self):

        # go no further if we've already partly ingested a different DVO
        if not self.correctDvo: return

        # check if we have up-to-date version
        path = self.skychunk.dvoLocation + "/Images.dat"
        if self.scratchDb.alreadyImportedThisDvoTable(path): 
            self.logger.debugPair("DVO Images.dat file", "up-to-date")
            return
        
        self.logger.infoSeparator()
        self.logger.infoPair("DVO Images.dat file", "NOT up-to-date")

        self.importFits(
                path,
                "SOURCE_ID IMAGE_ID EXTERN_ID FLAGS PHOTCODE",
                self.scratchDb.dvoImagesTable)
        self.logger.infoPair("Adding primary key to", self.scratchDb.dvoImagesTable)
        self.scratchDb.execute("ALTER TABLE " + self.scratchDb.dvoImagesTable + " ADD PRIMARY KEY (IMAGE_ID, EXTERN_ID)")

        self.scratchDb.setImportedThisDvoTable(path)

    '''
    Loads the DVO SkyTable into MySQL and then applies an index
    '''
    def loadSkyTable(self):
       
        # go no further if we've already partly ingested a different DVO
        if not self.correctDvo: return False

        path =  self.skychunk.dvoLocation + "/SkyTable.fits"
        if self.scratchDb.alreadyImportedThisDvoTable(path): 
            self.logger.debugPair("DVO SkyTable.fits file", "up-to-date")        
            return False

        self.logger.infoSeparator()
        self.logger.infoPair("DVO SkyTable.fits file", "NOT up-to-date")        

        if self.useStilts:
            self.importFits(
                path,
                "R_MIN R_MAX D_MIN D_MAX INDEX NAME",
                self.scratchDb.dvoSkyTable)
        else:
            
            # create dvoSkyTable 
            sql = "drop TABLE dvoSkyTable"
            try:
                self.scratchDb.execute(sql)
            except: pass

            # create dvoSkyTable 
            sql = "CREATE TABLE dvoSkyTable (R_MIN REAL, R_MAX REAL, D_MIN REAL, D_MAX REAL, INDEX_ INT, NAME CHAR(18))"
            self.scratchDb.execute(sql)

            # TODO path to DVO prog hardcoded temporarily
            cmd = "dvopsps skytable"
            cmd += " -dbhost " + self.scratchDb.dbHost
            cmd += " -dbname " + self.scratchDb.dbName
            cmd += " -dbuser " + self.scratchDb.dbUser
            cmd += " -dbpass " + self.scratchDb.dbPass
            cmd += " -D CATDIR " + self.skychunk.dvoLocation

            self.logger.infoPair("Running dvopsps", cmd)
            p = Popen(cmd, shell=True, stdout=PIPE)
            p.wait()

        self.logger.infoPair("Adding index to", self.scratchDb.dvoSkyTable)

        self.scratchDb.createIndex(self.scratchDb.dvoSkyTable, "INDEX")
        
        self.scratchDb.setImportedThisDvoTable(path)
        self.logger.infoPair("Finished importing SkyTable at", self.skychunk.dvoLocation)

        return True

    '''
    Sets a bunch of region files from DVO for a stripe starting at this RA and with this width
    '''
    def setSkyAreaAsStripe(self, ra=0., width=4.):

        self.setSkyArea(ra, ra+width, -91., 91.)

    '''
    Sets a bunch of region files from DVO for a RA/Dec box centered with these coords and with this side length
    '''
    def setSkyAreaAsBox(self, ra=180., dec=0., side=181.0):

        halfSide = side/2.0
        self.setSkyArea(ra-halfSide, ra+halfSide, dec-halfSide, dec+halfSide)

    '''
    Determines what DVO files need to be ingested to the defines region
    '''
    def setSkyArea(self, minRa=-1., maxRa=361., minDec=-91., maxDec=91):
    
        # go no further if we've already partly ingested a different DVO
        if not self.correctDvo: return

        if (maxRa - minRa <= 0) or (maxDec - minDec) <= 0:
            self.logger.errorPair("Zero range in either RA or Dec", 
                    "RA: %.2f -> %.2f, Dec: %.2f -> %.2f" % (minRa, maxRa, minDec, maxDec))
            return False        

        # update the two meta tables, Images and SkyTable
        self.loadSkyTable()

        # set bounding-box coords
        self.minRa = minRa
        self.maxRa = maxRa
        self.minDec = minDec
        self.maxDec = maxDec

        # reset all lists
        allRegions = self.scratchDb.getDvoRegions(minRa, maxRa, minDec, maxDec)
        allIngestedRegions = self.scratchDb.getIngestedDvoRegions()
        self.regionsToIngest = []
        self.regionsToPurge = []
        self.allPopulatedRegions = []
        self.allPopulatedRegionInfo = []
        self.regionsAlreadyIngested = []
        self.regionsIngestedButOutOfDate = []
        tablesToCleanup = []

        # loop through all regions in DVO inside this bounding-box
        for region in allRegions:

           regionId = region[0]
           regionRa = region[1]
           regionDec = region[2]
           regionPath = region[3]

           # get paths for this region for the file types we are interested in
           paths = []
           for fileType in self.ingestFileTypes:
               paths.append(self.skychunk.dvoLocation + "/" + regionPath + "." + fileType)

           # check for the existence of all interested file types
           skipThisRegion = False
           for path in paths:
               if not os.path.isfile(path):
                   skipThisRegion = True
                   break

           if skipThisRegion: continue

           # ok, have real, populated files so add this region to the list
           self.allPopulatedRegions.append(regionPath)
           self.allPopulatedRegionInfo.append([regionId, regionRa, regionDec])

           # if we have already imported up-to-date versions of all interested files, then we skip this region
           alreadyIngested = True
           for path in paths:
               if not self.scratchDb.alreadyImportedThisDvoTable(path): 
                   alreadyIngested = False
                   break
    
           if alreadyIngested:
               self.regionsAlreadyIngested.append(regionPath)
               continue

           # we have no record of ingesting this region, but does the table exist? 
           # May hev been left behind after a crash, and therefore needs to be cleaned-up
           for fileType in self.ingestFileTypes:
               tableName = self.scratchDb.getDbFriendlyTableName(regionPath + "." + fileType)
               if self.scratchDb.tableExists(tableName):
                   tablesToCleanup.append(tableName)

           # check if we have out-of-date versions of any interested files, if so, add region to the purge list
           outOfDate = False
           for path in paths:
               if self.scratchDb.haveOutOfDateVersionOfThisDvoTable(path):
                   outOfDate = True
                   break

           if outOfDate: self.regionsIngestedButOutOfDate.append(regionPath)

           self.regionsToIngest.append(regionPath)

        # add out-of-range regions for purge list
        items = list(set(allIngestedRegions) - set(self.allPopulatedRegions))
        for item in items: self.regionsToPurge.append(item)

        # now add out-of-date items to purge list
        for item in self.regionsIngestedButOutOfDate: self.regionsToPurge.append(item)

        self.sizeOfAllPopulatedRegions = self.getDiskSizeOfRegions(self.allPopulatedRegions)
        self.sizeOfRegionsAlreadyIngested = self.getDiskSizeOfRegions(self.regionsAlreadyIngested)
        self.sizeOfRegionsIngestedButOutOfDate = self.getDiskSizeOfRegions(self.regionsIngestedButOutOfDate)
        self.sizeOfRegionsToPurge = self.getDiskSizeOfRegions(self.regionsToPurge)
        self.sizeOfRegionsToIngest = self.getDiskSizeOfRegions(self.regionsToIngest)

        # any tables need cleaning up?
        if len(tablesToCleanup) > 0:
            self.logger.infoPair("Found tables to be cleaned up", "%d" % len(tablesToCleanup))
            self.scratchDb.dropTables(tablesToCleanup)
            

    '''
    Prints out a summary of the currently set regions
    '''
    def printSummary(self):

        self.logger.infoTitle("Summary for currently set DVO region") 
        self.logger.infoPair("DVO ingest type", self.__class__.__name__) 

        # go no further if we've already partly ingested a different DVO
        if not self.correctDvo:
            self.logger.infoPair("Wrong DVO in use", self.skychunk.dvoLocation)
            return

        for fileType in self.ingestFileTypes:
            self.logger.infoPair("Ingesting files of type", fileType)

        self.logger.infoPair("RA range", "%.2f -> %.2f" % (self.minRa, self.maxRa))
        self.logger.infoPair("Dec range", "%.2f -> %.2f" % (self.minDec, self.maxDec))

        self.logger.info("+----------------+----------------+----------------+----------------+----------------+")
        self.logger.info("|   populated    |   ingested     |  out-of-date   |    to-purge    |   to ingest    |")
        self.logger.info("+----------------+----------------+----------------+----------------+----------------+")
        self.logger.info("| %5d (%6s) | %5d (%6s) | %5d (%6s) | %5d (%6s) | %5d (%6s) |" % (
                    len(self.allPopulatedRegions), self.getFileSizeStr(self.sizeOfAllPopulatedRegions), 
                    len(self.regionsAlreadyIngested), self.getFileSizeStr(self.sizeOfRegionsAlreadyIngested), 
                    len(self.regionsIngestedButOutOfDate), self.getFileSizeStr(self.sizeOfRegionsIngestedButOutOfDate),
                    len(self.regionsToPurge), self.getFileSizeStr(self.sizeOfRegionsToPurge),
                    len(self.regionsToIngest), self.getFileSizeStr(self.sizeOfRegionsToIngest))) 
        self.logger.info("+----------------+----------------+----------------+----------------+----------------+")

    '''
    Returns the size on disk of the regions that need to be ingested in GB
    '''
    def getDiskSizeOfRegionsToBeIngested(self):

        try:
            return self.sizeOfRegionsToIngest/1073741824.0
        except:
            return 0.0

    '''
    Purges a single region from the database
    '''
    def purgeRegion(self, region):

        regions = []
        regions.append(region)
        self.purgeRegions(regions)

    '''
    Purges the provided regions from the database
    '''
    def purgeRegions(self, regions):

        # go no further if we've already partly ingested a different DVO
        if not self.correctDvo: return False

        if len(regions) < 1: return False

        self.logger.infoPair("Purging", "%d regions" % len(regions)) 

        detectionDeleteCount = self.scratchDb.purgeTheseDvoRegions(regions, self.ingestFileTypes)
        if detectionDeleteCount > 0: self.logger.infoPair("Deleted", "%d detections" % detectionDeleteCount)

        return True

    '''
    Formats a nice file size string
    Accepts a size in bytes
    '''
    def getFileSizeStr(self, byteSize):

        # format a string for the user -smaller than one MB
        if byteSize < 1048576: return "%.1f K" % (byteSize/1024.0)
        # smaller than one GB
        elif byteSize < 1073741824: return "%.1fM" % (byteSize/1048576.0)
        # more than a GB
        else: return "%.1fG" % (byteSize/1073741824.0)

    '''
    Returns the disk size of this file
    '''
    def getDiskSize(self, path):

        return os.stat(path).st_size

    '''
    Returns the disk size of these DVO regions, i.e. the size on disk of all relevant DVO FITS files
    '''
    def getDiskSizeOfRegions(self, regions):

        size = 0.0
        for region in regions:
             self.logger.infoPair("sizes for region", region)
             # get combined size of all interested files
             for fileType in self.ingestFileTypes:
                 size = size + self.getDiskSize(self.skychunk.dvoLocation + "/" + region + "." + fileType)
                 # EAM TEST I/O
                 self.logger.infoPair(fileType, size)

        return size


    '''
    Ingests object files for this region, given an INDEX from the SkyTable
    '''
    def ingestRegionUsingIndex(self, index):
        return self.ingestRegion(self.scratchDb.getRegionNameFromThisDvoIndex(index))

    '''
    Abstract method for ingesting a particular reegion. Must be implemented in all subclasses
    '''
    def ingestRegion(self, region):
        return self.logger.errorPair("ingestRegion()", "not implemented for this class") 

    '''
    Synchronises MySQL database with the sky area currently set from DVO
    Populates dvoDetections and dvoDone with FITS tables from DVO for the defined sky area, this
    includes purging detections outside sky area
    '''
    def sync(self):

        # go no further if we've already partly ingested a different DVO
        if not self.correctDvo: return False

        # check we have something to do
        if len(self.regionsToPurge) < 1 and len(self.regionsToIngest) < 1:
            self.logger.infoPair("Nothing to ingest or purge", "finished ingest")
            return True

        # make sure we have an up-to-date Images table
        self.loadImages()

        # purge stuff
        self.purgeRegions(self.regionsToPurge)

        totalRegions = len(self.regionsToIngest)
        if totalRegions > 0: self.logger.infoSeparator()

        # now loop through all regions and ingest to MySQL
        ingestedRegionsCount = 0
        regionCount = 1
        for region in self.regionsToIngest:

           self.logger.infoTitle("Region %d of %d: %s" % (regionCount, totalRegions, region))
           if self.ingestRegion(region): ingestedRegionsCount += 1
           regionCount += 1

        self.logger.infoSeparator()
        self.logger.infoPair("Ingested", "%d DVO regions (out of %d)" % (ingestedRegionsCount, len(self.regionsToIngest)))

        # update lists after attempted sync
        self.setSkyArea(self.minRa, self.maxRa, self.minDec, self.maxDec)
        self.printSummary()

        return True

    '''
    Imports a FITS file. A regex filter lets you choose which tables to pull from the file
    An optional final argument lets you name the resultant database table
    '''
    def importFits(self, path, columns, tableName):

      self.logger.infoPair("Importing file", "%s (%s)" % (path, self.getFileSizeStr(self.getDiskSize(path))))
      self.logger.debugPair("Writing to database table", tableName)

      tables = stilts.treads(path)

      count = 0
      for table in tables:

          self.logger.debugPair("Reading IPP table", table.name)
          table = stilts.tpipe(table, cmd='explodeall')
     
          # IPP FITS files are littered with infinity values. Remove them
          self.logger.debugPair("Removing", "infinity values")
          table = stilts.tpipe(table, cmd='replaceval -Infinity null *')
          table = stilts.tpipe(table, cmd='replaceval Infinity null *')

          attempts = 0
          self.logger.debugPair("Writing table to database",  table.name)
          while True:
    
              try:
                  if attempts > 0: self.logger.infoPair("Attempt %d to write" % attempts, table.name)
                  table.cmd_keepcols(columns).write(self.scratchDb.url + '#' + tableName)
                  break
              except:
                  # wait a few secs then try again
                  time.sleep(2)
                  attempts = attempts + 1

          count = count + 1

      self.logger.debugPair("Finished importing", tableName)

      return tableName

    '''
    Imports the Photcode file
    '''
    def importPhotcodes(self, path, tableName):

      self.logger.infoPair("Importing file", "%s (%s)" % (path, self.getFileSizeStr(self.getDiskSize(path))))
      self.logger.debugPair("Writing to database table", tableName)

      tables = stilts.treads(path)

      # complain if len(tables) > 1

      self.logger.debugPair("Reading IPP table", tables[0].name)
      table = stilts.tpipe(tables[0], cmd='explodeall')
     
      # IPP FITS files are littered with infinity values. Remove them (probably not needed for photcodes)
      self.logger.debugPair("Removing", "infinity values")
      table = stilts.tpipe(table, cmd='replaceval -Infinity null *')
      table = stilts.tpipe(table, cmd='replaceval Infinity null *')

      # convert the 1-char field TYPE to an int for querying ease (probably more sophisticated was to do this...)
      table = stilts.tpipe(table, cmd='addcol -after TYPE TYPE_AS_INT "TYPE + 0"')

      attempts = 0
      self.logger.debugPair("Writing table to database",  table.name)
      while True:
    
          try:
              if attempts > 0: self.logger.infoPair("Attempt %d to write" % attempts, table.name)
              table.cmd_keepcols("CODE NAME TYPE_AS_INT C_LAM K").write(self.scratchDb.url + '#' + tableName)
              break
          except:
              # wait a few secs then try again
              time.sleep(2)
              attempts = attempts + 1
              
              count = count + 1
              
      self.logger.debugPair("Finished importing", tableName)

      return tableName


    '''
    ingest skyregion into MySQL database using the native DVO program dvopsps
    Populates dvoDetections and dvoDone with FITS tables from DVO for the defined sky area, this
    includes purging detections outside sky area
    '''
    def nativeIngestDetections(self, boxId, raCenter, decCenter, boxSize):

        # XXX put the chunk below in a separate method 
        # blow away existing dvoDetection table

        # clear the 'ingested' field for all boxes
        self.ippToPspsDb.clearIngestedBoxes()
        
        # drop detections table
        self.logger.infoPair("Dropping table", self.scratchDb.dvoDetectionTable)
        self.scratchDb.dropTable(self.scratchDb.dvoDetectionTable)
        
        # create detections table
        self.logger.infoPair("Creating table", self.scratchDb.dvoDetectionTable)
        sql = "CREATE TABLE " + self.scratchDb.dvoDetectionTable + " LIKE dvoDetection"
        try: self.scratchDb.execute(sql)
        except:
            self.logger.errorPair("Unable to create table", self.scratchDb.dvoDetectionTable)
            return False
        self.scratchDb.changeEngineToInnoDB(self.scratchDb.dvoDetectionTable)
        
        # add fileID column
        sql = "ALTER TABLE " + self.scratchDb.dvoDetectionTable + " ADD fileID INT NOT NULL"
        try: self.scratchDb.execute(sql)
        except: 
            self.logger.errorPair("Unable to add fileID column to ", self.scratchDb.dvoDetectionTable)
            return False

        # add an index to the fileID column
        self.scratchDb.createIndex(self.scratchDb.dvoDetectionTable, "fileID")

        # make sure we have an up-to-date Images table
        self.loadImages()

        # dvopsps -D catdir CATDIR -region .... -dbhost xx -dbname xx -dbuser xx -dbpass xx
        halfSize = boxSize / 2.0
        # -region raCenter-halfSize raCenter+halfSize decCenter-halfSize decCenter+halfSize

        # TODO path to DVO prog hardcoded temporarily
        cmd = "dvopsps detections"
        cmd += " -dbhost " + self.scratchDb.dbHost
        cmd += " -dbname " + self.scratchDb.dbName
        cmd += " -dbuser " + self.scratchDb.dbUser
        cmd += " -dbpass " + self.scratchDb.dbPass
        cmd += " -D CATDIR " + self.skychunk.dvoLocation
        cmd += " -region "
        cmd += " " + str(raCenter-halfSize)
        cmd += " " + str(raCenter+halfSize)
        cmd += " " + str(decCenter-halfSize)
        cmd += " " + str(decCenter+halfSize)

        if self.skychunk.parallel:
            cmd += " -parallel"

        self.logger.infoPair("Running dvopsps", cmd)
        p = Popen(cmd, shell=True, stdout=PIPE)
        p.wait()

        self.ippToPspsDb.setIngestedBox(boxId)

        # add fileID column
        sql = "ANALYZE TABLE " + self.scratchDb.dvoDetectionTable
        try: self.scratchDb.execute(sql)
        except: 
            self.logger.errorPair("Unable to analyze mysql table", self.scratchDb.dvoDetectionTable)
            return False

        # update lists after attempted sync
        # self.printSummary()

        return True
