#!/usr/bin/python3.6
'''
Copyright (C) 2015 - Swedish Meteorological and Hydrological Institute (SMHI)

This file is part of RAVE.

RAVE is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

RAVE is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
See the GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with RAVE.  If not, see <http://www.gnu.org/licenses/>.
'''
## Command-line tool for creating acrr compoisites

## @file
## @author Anders Henja, SMHI
## @date 2015-04-13

import sys,string,os,datetime, math
import odim_source
import rave_dom_db
import rave_pgf_logger
import rave_tile_registry
import _raveio, _acrr, _cartesianvolume, _cartesian, _rave, _gra
from gadjust import gra, obsmatcher
from gadjust.gra import gra_coefficient
from compositing import compositing
from tiled_compositing import tiled_compositing
from rave_dom import observation
from gadjust import grapoint

logger = rave_pgf_logger.rave_pgf_stdout_client()

from rave_defines import CENTER_ID, GAIN, OFFSET, MERGETERMS, RAVE_PGF_QUALITY_FIELD_REPROCESSING

ravebdb = None
try:
  import rave_bdb
  ravebdb = rave_bdb.rave_bdb()
except:
  pass

nodomdb=False
try:
  import rave_dom_db
except:
  nodomdb=True

def strToNumber(sval):
  try:
    return int(sval)
  except ValueError as e:
    return float(sval)

class nrd_obsmatcher(obsmatcher.obsmatcher):
  def __init__(self, filepath):
    self.filepath = filepath
    super(nrd_obsmatcher, self).__init__(None)
  
  def read_nrd_synops(self, filename):
    fp = open(filename)
    lines = fp.readlines()
    result = []  
    year, month, day, hour = None, None, None, None
    for l in lines:
      toks = l.split()
      if len(toks) > 1 and toks[0].strip() == "SYNOP-BRDC":
        year = toks[1]
        month = toks[2]
        day = toks[3]
        hour = toks[4]
      elif len(toks) == 9 and toks[0].strip().isdigit():
        if strToNumber(toks[6].strip()) >= 0 and strToNumber(toks[7].strip()) >= 0.0:
          result.append(observation("%05d"%int(toks[0].strip()), "unknown", 
                                    observation.SYNOP, "%4d%02d%02d"%(int(year),int(month),int(day)), 
                                    "%02d0000"%int(hour),  strToNumber(toks[2].strip()), strToNumber(toks[1].strip()), 
                                    strToNumber(toks[7].strip()), 12)) 
    #print "NR SYNOP (%d)"%len(result)       
    return result

  def filter_synop_by_box(self, synops, ul, lr):
    result = []
    for s in synops:
      #print "s.longitude = %f, s.latitude = %f (%f,%f, %f,%f)"%(s.longitude, s.latitude, ul[0], ul[1], lr[0], lr[1])
      if s.longitude >= ul[0] and s.longitude <= lr[0] and s.latitude >= lr[1] and s.latitude <= ul[1]:
        result.append(s)
    return result

  def match(self, image, acc_period=12, quantity="ACRR", how_task="se.smhi.composite.distance.radar", offset_hours=0):
    ul,lr=image.getExtremeLonLatBoundaries()
    d = image.date
    t = image.time
    image.defaultParameter=quantity
    
    filename = "%s/SYN%04d%02d%02d%02d.DAT"%(self.filepath, int(d[:4]), int(d[4:6]), int(d[6:8]), int(t[0:2]))#SYN2014102006.DAT
    
    nrd_synops = self.filter_synop_by_box(self.read_nrd_synops(filename), (ul[0]*180.0/math.pi, ul[1]*180.0/math.pi), (lr[0]*180.0/math.pi, lr[1]*180.0/math.pi))
    #print "Number of N2 synop initially %d"%len(nrd_synops)
    xpts = 0
    xptst = 0
    xptsq = 0
    xptsv = 0

    result = []
    for obs in nrd_synops:
      if obs.accumulation_period == acc_period:
        xpts = xpts + 1
        t, v = image.getConvertedValueAtLonLat((obs.longitude*math.pi/180.0, obs.latitude*math.pi/180.0))
        if t in [_rave.RaveValueType_DATA, _rave.RaveValueType_UNDETECT]:
          xptst = xptst + 1
          d = image.getQualityValueAtLonLat((obs.longitude*math.pi/180.0, obs.latitude*math.pi/180.0), how_task)
          if d != None:
            xptsq = xptsq + 1
            if obs.liquid_precipitation >= grapoint.MIN_GMM and v >= grapoint.MIN_RMM:
              xptsv = xptsv + 1
              result.append(grapoint.grapoint.from_observation(t, v, d, obs))

    logger.info("obses = %d, xpts=%d, xptst = %d, xptsq = %d, xptsv = %d"%(len(nrd_synops), xpts, xptst, xptsq, xptsv))
    
    return result     

def generate_comp(filepath, pattern, fdate, ftime, areaid="allbltgmaps_4000", opath="/tmp", 
                  detectors="ropo,beamb,overshooting,distance", ignore_malfunc=True, 
                  quantity="DBZH", product="pcappi", selectionMethod="NEAREST_RADAR", 
                  height=500.0, range=200000.0, elangle=0.0, ctfilter=True, 
                  qitotal_field=None, prodpar=None, zr_a=200.0, zr_b=1.5):
  import glob
  filename = "%s/%s_comp_%s%s.h5"%(opath,areaid,fdate,ftime)
  if os.path.exists(filename):
    return filename
  files = glob.glob("%s/%s"%(filepath, pattern)) #se*_pvol*.h5
  if len(files) == 0:
    logger.warn("No files found for %s/%s"%(filepath, pattern))
    return filename
  comp = compositing(ravebdb)
  comp.height = height
  comp.elangle = elangle
  comp.range = range
  comp.nodata = 255
  comp.undetect = 0
  comp.filenames = files
  comp.detectors=string.split(detectors,",")
  comp.ignore_malfunc=ignore_malfunc
  comp.quantity = quantity
  comp.gain = GAIN
  comp.offset = OFFSET
  comp.set_product_from_string(product)
  comp.set_method_from_string(selectionMethod)
  comp.applyctfilter = ctfilter
  comp.applygra = False 
  comp.zr_A = zr_a
  comp.zr_b = zr_b
  comp.qitotal_field=qitotal_field
  comp.prodpar = prodpar
  
  comp.reprocess_quality_field = False

  comp = tiled_compositing(comp,  mp_process_qc=True)

  result = comp.generate(fdate, ftime, areaid)
  
  rio = _raveio.new()
  rio.object = result
  rio.filename = filename
  rio.save()
  
  return filename

def apply_gra_correction(gra_corrected_composite_path, comp, areaid, quantity, distancefield="se.smhi.composite.distance.radar", zr_a=200.0, zr_b=1.5):
  try:
    d = comp.date
    t = comp.time
    db = rave_dom_db.create_db_from_conf()
    dt = datetime.datetime(int(d[:4]), int(d[4:6]), int(d[6:]), int(t[:2]), int(t[2:4]), 0)
    dt = dt - datetime.timedelta(seconds=3600*12) # 12 hours back in time for now..
    grac = db.get_gra_coefficient(dt)
    if gra_corrected_composite_path != None and not os.path.exists(gra_corrected_composite_path):
      logger.info("Creating path for storing corrected composites: %s"%gra_corrected_composite_path)
      os.makedirs(gra_corrected_composite_path)

    if grac != None:
      logger.debug("Applying gra coefficients, quantity: %s"%quantity)
      gra = _gra.new()
      gra.A = grac.a
      gra.B = grac.b
      gra.C = grac.c
      gra.zrA = zr_a
      gra.zrb = zr_b
      #logger.info("GRA: A=%f, B=%f, C=%f, ZRA=%f, ZRB=%f"%(gra.A, gra.B, gra.C, gra.zrA, gra.zrb))
      param = comp.getParameter(quantity)
      dfield = param.getQualityFieldByHowTask(distancefield)
      gra_field = gra.apply(dfield, param)
      gra_field.quantity = quantity + "_CORR"
      comp.addParameter(gra_field)
      
      if gra_corrected_composite_path != None:
        rio = _raveio.new()
        rio.object = comp
        filename = "%s/%s_gra_corrected_comp_%s%s.h5"%(gra_corrected_composite_path,areaid,comp.date,comp.time)
        rio.save(filename)
        logger.debug("Created gra corrected composite %s"%filename)
      return comp, gra_field.quantity
    else:
      logger.info("No gra coefficients found for given date/time, ignoring gra adjustment")
  except Exception as e:
    import traceback
    traceback.print_exc()
    logger.error("Failed to apply gra coefficients", exc_info=1)
  return comp, quantity

def generate_acrr(files, opath, areaid, etime, edate, startdt, enddt, zr_a=200.0, zr_b=1.5, 
                  quantity="DBZH", accept=0.0, distancefield="se.smhi.composite.distance.radar", 
                  interval=12, N=13, applygra=False):
  acrrproduct = None
  
  acrr = _acrr.new()
  acrr.nodata = -1.0
  acrr.undetect = 0.0
  acrr.quality_field_name = distancefield

  acrr_quantity = quantity

  qualityfield = None

  filename = "%s/acrr_%dH_%s_%s%s.h5"%(opath, interval, areaid, edate, etime[:4])
  if os.path.exists(filename):
    return

  for fname in files:
    obj = None
    try:
      rio = _raveio.open(fname)
      obj = rio.object
    except Exception as e:
      continue

    if _cartesianvolume.isCartesianVolume(obj):
      obj = obj.getImage(0)
    
    if not _cartesian.isCartesian(obj):
      raise AttributeError("Must call plugin with cartesian products")

    if acrrproduct == None:
      acrrproduct = _cartesian.new()
      acrrproduct.xscale = obj.xscale
      acrrproduct.yscale = obj.yscale
      acrrproduct.areaextent = obj.areaextent
      acrrproduct.projection = obj.projection
      acrrproduct.product = obj.product
      acrrproduct.source = obj.source
      acrrproduct.time = etime
      acrrproduct.date = edate
      acrrproduct.startdate = startdt.strftime("%Y%m%d")
      acrrproduct.starttime = startdt.strftime("%H%M00")
      acrrproduct.enddate = enddt.strftime("%Y%m%d")
      acrrproduct.endtime = enddt.strftime("%H%M00")
      
    if obj.xscale != acrrproduct.xscale or obj.yscale != acrrproduct.yscale or \
       obj.projection.definition != acrrproduct.projection.definition:
      raise AttributeError("Scale or projdef inconsistancy for used area")
    par = obj.getParameter(quantity)
    if par == None:
      logger.warn("Could not find parameter (%s) for %s %s"%(acrr_quantity, obj.date, obj.time))
    else:
      pdfield = par.getQualityFieldByHowTask(distancefield) 
      if pdfield != None:
        acrr.sum(par, zr_a, zr_b)
      if qualityfield == None:
        qualityfield = pdfield 

  # accept, N, hours
  if not acrr.isInitialized():
    logger.info("No files can be found for acrr accumulation")
    return None
  
  acrrparam = acrr.accumulate(accept, N, interval)
  acrrproduct.addParameter(acrrparam)
  #acrrproduct.addQualityField(qualityfield) # TODO REMOVE ME????

  if applygra:
    acrrproduct, acrr_quantity = apply_gra_correction(None, acrrproduct, areaid, acrrparam.quantity, distancefield, zr_a, zr_b)

  rio = _raveio.new()
  rio.object = acrrproduct
  rio.filename = filename
  logger.info("Generated acrr product %s"%filename)
  rio.save()

def run_acrr_composite_generation(vpath, cpath, opath, areaid, interval, filesPerHour, 
                                  acceptableLoss, startdate, enddate, pattern, applygra=False,
                                  distancefield="se.smhi.composite.distance.radar", 
                                  detectors="ropo,beamb,overshooting,distance", ignore_malfunc=True, 
                                  product="pcappi", selectionMethod="NEAREST_RADAR", height=500.0, range=200000.0, elangle=0.0,
                                  ctfilter=True, qitotal_field=None, prodpar=None, zr_a=200.0, zr_b=1.5):
  filePeriod = 60/filesPerHour
  N = interval * filesPerHour + 1
  currdt = startdate
  startaccdate = currdt
  files = []
  while currdt <= enddate:
    volumepath = "%s/%s"%(vpath,currdt.strftime("%m/%d/%H/%M"))
    filename = generate_comp(volumepath, pattern, currdt.strftime("%Y%m%d"), currdt.strftime("%H%M00"), 
                             areaid, cpath, detectors, ignore_malfunc, "DBZH", product,
                             selectionMethod, height, range, elangle, ctfilter, qitotal_field, prodpar,
                             zr_a, zr_b)
    files.append(filename)
    
    if len(files) == N:
      logger.info("Generating ACRR product for %s"%currdt.strftime("%Y%m%d%H%M"))
      generate_acrr(files, opath, areaid, currdt.strftime("%H%M00"), currdt.strftime("%Y%m%d"), startaccdate, currdt, zr_a, zr_b, "DBZH", 
                    acceptableLoss, distancefield, interval, N, applygra)
      files=[filename]
      startaccdate = currdt

        
    currdt = currdt + datetime.timedelta(seconds=filePeriod*60)

def generate_gra(files, opath, areaid, etime, edate, startdt, enddt, zr_a=200.0, zr_b=1.5, 
                 quantity="DBZH", accept=0.0, distancefield="se.smhi.composite.distance.radar", 
                 interval=12, N=13, adjustmentfile=None, nrd_synop_path=None):
  
  acrrproduct = None
  acrr = _acrr.new()
  acrr.nodata = -1.0
  acrr.undetect = 0.0
  acrr.quality_field_name = distancefield

  for fname in files:
    obj = None
    try:
      rio = _raveio.open(fname)
      obj = rio.object
    except Exception as e:
      continue

    if _cartesianvolume.isCartesianVolume(obj):
      obj = obj.getImage(0)

    if not _cartesian.isCartesian(obj):
      raise AttributeError("Must call plugin with cartesian products")

    if acrrproduct == None:
      acrrproduct = _cartesian.new()
      acrrproduct.xscale = obj.xscale
      acrrproduct.yscale = obj.yscale
      acrrproduct.areaextent = obj.areaextent
      acrrproduct.projection = obj.projection
      acrrproduct.product = obj.product
      acrrproduct.source = obj.source
      acrrproduct.time = etime
      acrrproduct.date = edate

    if obj.xscale != acrrproduct.xscale or obj.yscale != acrrproduct.yscale or \
       obj.projection.definition != acrrproduct.projection.definition:
      raise AttributeError("Scale or projdef inconsistancy for used area")

    par = obj.getParameter(quantity)
    if par == None:
      logger.warn("Could not find parameter (%s) for %s %s"%(quantity, obj.date, obj.time))
    else:
      if par.getQualityFieldByHowTask(distancefield) != None:
        acrr.sum(par, zr_a, zr_b)

  # accept, N, hours
  acrrparam = acrr.accumulate(accept, N, interval)
  acrrproduct.addParameter(acrrparam)
    
  db = rave_dom_db.create_db_from_conf()
  if nrd_synop_path != None:
    matcher = nrd_obsmatcher(nrd_synop_path)
  else:
    matcher = obsmatcher.obsmatcher(db)
  
  points = matcher.match(acrrproduct, acc_period=interval, quantity="ACRR", how_task=distancefield, offset_hours=interval)
  if len(points) == 0:
    logger.warn("Could not find any matching observations")
    return None
  
  logger.info("Matched %d points between acrr product and observation db"%len(points))
  db.merge(points)

  d = acrrproduct.date
  t = acrrproduct.time

  tlimit = datetime.datetime(int(d[:4]), int(d[4:6]), int(d[6:8]), int(t[0:2]), int(t[2:4]), int(t[4:6]))
  tlimit = tlimit - datetime.timedelta(hours=interval*MERGETERMS)
  dlimit = datetime.datetime(int(d[:4]), int(d[4:6]), int(d[6:8]), int(t[0:2]), int(t[2:4]), int(t[4:6]))
  dlimit = dlimit - datetime.timedelta(hours=12*MERGETERMS)
  
  db.delete_grapoints(dlimit) # We don't want any points older than 12 hour * MERGETERMS back in time
  
  points = db.get_grapoints(tlimit); # Get all gra points newer than interval*MERGETERMS hours back in time
  logger.info("Using %d number of points for calculating the gra coefficients"%len(points))
  
  if adjustmentfile != None:
    significant, npoints, loss, r, sig, corr_coeff, a, b, c, m, dev = gra.generate(points, edate, etime, adjustmentfile)
  else:
    significant, npoints, loss, r, sig, corr_coeff, a, b, c, m, dev = gra.generate(points, edate, etime)

  # Also store the coefficients in the database so that we can search for them when applying the coefficients
  NOD = odim_source.NODfromSource(acrrproduct)
  if not NOD:
    NOD = ""
  grac = gra_coefficient(NOD, acrrproduct.date, acrrproduct.time, significant, npoints, loss, r, sig, corr_coeff, a, b, c, float(m), float(dev))
  db.merge(grac)

def create_gra_coefficients(vpath, cpath, opath, areaid, interval, filesPerHour,
                            acceptableLoss, startdate, enddate, pattern,
                            distancefield="se.smhi.composite.distance.radar",
                            detectors="ropo,beamb,overshooting,distance", ignore_malfunc=True, 
                            product="pcappi", selectionMethod="NEAREST_RADAR", height=500.0, range=200000.0, elangle=0.0,
                            ctfilter=True, qitotal_field=None, prodpar=None, zr_a=200.0, zr_b=1.5, adjustmentfile=None, nrd_synop_path=None):
  filePeriod = 60/filesPerHour
  N = interval * filesPerHour + 1
  currdt = startdate
  startaccdate = currdt
  files = []
  
  db = rave_dom_db.create_db_from_conf()
  db.purge_grapoints()
  while currdt <= enddate:
    volumepath = "%s/%s"%(vpath,currdt.strftime("%m/%d/%H/%M"))
    filename = generate_comp(volumepath, pattern, currdt.strftime("%Y%m%d"), currdt.strftime("%H%M00"), 
                             areaid, cpath, detectors, ignore_malfunc, "DBZH", product,
                             selectionMethod, height, range, elangle, ctfilter, qitotal_field, prodpar,
                             zr_a, zr_b)
    files.append(filename)
    
    if len(files) == N:
      logger.info("Generating gra coefficient for %s"%currdt.strftime("%Y%m%d%H%M"))
      
      generate_gra(files, opath, areaid, currdt.strftime("%H%M00"), currdt.strftime("%Y%m%d"), startaccdate, currdt, zr_a, zr_b, "DBZH", 
                   acceptableLoss, distancefield, interval, N, adjustmentfile, nrd_synop_path)
      files=[filename]
      startaccdate = currdt
      
    currdt = currdt + datetime.timedelta(seconds=filePeriod*60)

if __name__=="__main__":
  from optparse import OptionParser
  START_PERIOD="201410010000"
  END_PERIOD="201411010000"
  AREA="allbltgmaps_4000"
  INTERVAL=12 #HOURS
  FILES_PER_HOUR=4
  ZR_a=200.0
  ZR_b=1.5
  FT_UTC=6
  DISTANCE_FIELD="se.smhi.composite.distance.radar"
  ACCEPTABLE_LOSS=0
  N = INTERVAL * FILES_PER_HOUR + 1
  VOLUME_PATH="/storage/baltrad/data/2014_all"
  COMPOSITE_PATH="/storage/baltrad/acrr/composites"
  ACRR_PATH="/storage/baltrad/acrr"
  startdt = datetime.datetime.strptime(START_PERIOD, "%Y%m%d%H%M")
  enddt = datetime.datetime.strptime(END_PERIOD, "%Y%m%d%H%M")
  currdt = startdt
  minuteinterval = 60 / FILES_PER_HOUR
 
  usage = "usage: %prog --input=<input path> --startdt=<start datetime> --enddt=<end datetime> --pattern=<pattern> [--area=<area>] [--hours=hours] [--create-gra-coefficients]"
  usage += "\nGenerates acrr composites directly from polar scans and volumes."
  usage += "\nIf specifying --create-gra-coefficients, then a run will be performed that produces gra coefficients that can be used in the gra adjustment."
  
  parser = OptionParser(usage=usage)

  parser.add_option("-i", "--input", dest="inputpath", default=VOLUME_PATH,
                    help="Locating where file searching should be started. Expected subdirectory format is <yyyy>/<MM>/<dd>/<HH>/<mm>")

  parser.add_option("-c", "--coutput", dest="compositepath", default=COMPOSITE_PATH,
                    help="Locating where composites should be placed. Expected subdirectory format is <yyyy>/<MM>/<dd>/<HH>/<mm>")

  parser.add_option("-o", "--output", dest="outputpath", default=ACRR_PATH,
                    help="Locating where acrr products should be placed. Expected subdirectory format is <yyyy>/<MM>/<dd>/<HH>/<mm>")

  parser.add_option("-S", "--startdt", dest="startdt", default=START_PERIOD,
                    help="Start date/time. Format is <yyyy><MM><dd><HH><mm>.")

  parser.add_option("-E", "--enddt", dest="enddt", default=END_PERIOD,
                    help="End date/time. Format is <yyyy><MM><dd><HH><mm>.")

  parser.add_option("-p", "--pattern", dest="pattern", default="*pvol*.h5",
                    help="File pattern to search for, e.g. --pattern='se*pvol*.h5'.")

  parser.add_option("-a", "--area", dest="area", default="allbltgmaps_4000",
                    help="Name of Cartesian area to which to generate the composite.")

  parser.add_option("-H", "--hours", dest="hours", type="int", default=12,
                    help="Number of hours that should be included in each accumulation period, either 1,2,3,4,6,8,12 or 24. Default is 12")

  parser.add_option("--create-gra-coefficients", dest="create_gra_coefficients", action="store_true", default=False,
                    help="If gra coefficients should be created.")
  
  parser.add_option("--adjustmentfile", dest="adjustmentfile", default=None,
                    help="The adjustmentfile that gra coefficients should be written to. Default: None")

  parser.add_option("--apply-gra-correction", dest="applygra", action="store_true", default=False,
                    help="If gra coefficients should be applied to the acrr product.")

  parser.add_option("--nrd-synop-path", dest="nrd_synop_path", default=None,
                    help="If you want to use nordrad2 synop files for gra coefficient generation")

  (options, args) = parser.parse_args()
  
  print("Running with:")
  print("Inputpath: %s"%options.inputpath)
  print("Composite path: %s"%options.compositepath)
  print("Acrr path: %s"%options.outputpath)
  print("Datetime: %s -> %s"%(options.startdt, options.enddt))
  print("Pattern: %s"%options.pattern)
  print("Area: %s"%options.area)
  print("Hours: %d"%options.hours)
  print("Create gra coefficients: %s"%str(options.create_gra_coefficients)) 
  print("Adjustmentfile: %s"%options.adjustmentfile)
  print("Apply gra correction: %s"%str(options.applygra))
  print("N2 synop file path: %s"%str(options.nrd_synop_path))
  
  if options.create_gra_coefficients:
    create_gra_coefficients(options.inputpath, options.compositepath, options.outputpath, options.area, options.hours, 
                            FILES_PER_HOUR, ACCEPTABLE_LOSS,
                            datetime.datetime.strptime(options.startdt, "%Y%m%d%H%M"), 
                            datetime.datetime.strptime(options.enddt, "%Y%m%d%H%M"), 
                            options.pattern,zr_a=ZR_a,zr_b=ZR_b,adjustmentfile=options.adjustmentfile, nrd_synop_path=options.nrd_synop_path)    
  else:
    run_acrr_composite_generation(options.inputpath, options.compositepath, options.outputpath, options.area, options.hours, 
                                  FILES_PER_HOUR, ACCEPTABLE_LOSS,
                                  datetime.datetime.strptime(options.startdt, "%Y%m%d%H%M"), 
                                  datetime.datetime.strptime(options.enddt, "%Y%m%d%H%M"), 
                                  options.pattern, options.applygra, zr_a=ZR_a, zr_b=ZR_b)

    
  