#!/usr/bin/python3.6
'''
Copyright (C) 2016 - 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 merging scans and volumes

## @file
## @author Anders Henja, SMHI
## @date 2016-10-21

import os,getopt,sys,errno
import re, string, shutil
import fnmatch
import polar_merger
import _raveio

class merge_files(object):
  def __init__(self, indir, outdir, sources=None, ftype="SCAN", copy_ignored=False):
    self.indir=indir
    self.outdir=outdir
    self.sources=sources
    self.ftype=ftype.lower()
    if self.ftype not in ["scan","pvol"]:
      raise IOError("Not a valid type")
    self.pattern = re.compile("[a-z0-9A-Z]+_[A-Za-z0-9]+_(([0-9\.]+)_)?[0-9]{8}T[0-9]{4}Z_0x[0-9]+.h5")
    self.bnamepattern = re.compile("([a-z0-9A-Z]+_[A-Za-z0-9]+_(([0-9\.]+)_)?[0-9]{8}T[0-9]{4}Z)_0x[0-9]+.h5")
    self.copy_ignored=copy_ignored

  ##
  # Creates a directory with parents if necessary. If directory already exists nothing will
  # happen.
  # @param dname the name of the directory including parents
  def make_dir(self, dname):
    try:
      os.makedirs(dname)
    except OSError as e:
      if e.errno == errno.EEXIST and os.path.isdir(dname):
        pass
      else:
        raise

  ##
  # Reads all sources from the specified directory. Assumes the first 5 characters in the filenames
  # is the NOD.
  # @param root the directory to list
  # @return a list of NOD:s
  def read_sources(self, root):
    sources=[]
    files = [f for f in os.listdir(root) if os.path.isfile(os.path.join(root, f))]
    for f in files:
      if self.pattern.match(f):
        nod=f[:5]
        if nod not in sources:
          sources.append(nod)
    return sources    

  ##
  # Traverses through indir and all subdirectories and try to merge all files matching
  # the base filename pattern.
  #
  def find_and_merge(self):
    for root, d, files in os.walk(indir):
      base=root[len(self.indir):]
      if len(base)>0:
        if base[0]=='/':
          base=base[1:]

      sources = self.sources
      if sources is None:
        sources = self.read_sources(root)

      for src in sources:
        pvols = []
        scans = {}
        for item in fnmatch.filter(files, "%s_%s_*.h5"%(src,self.ftype)):
          m = self.pattern.match(item)
          if m:
            if m.group(2) is None:
              pvols.append("%s/%s"%(root,item))
            else:
              if not scans.has_key(m.group(2)):
                scans[m.group(2)]=[]
              scans[m.group(2)].append("%s/%s"%(root,item))
        if len(pvols)>1:
          self._merge_polar_volumes(pvols, "%s/%s"%(self.outdir, base))

        if len(scans.keys())>0:
          for k in scans.keys():
            if len(scans[k]) > 1:
              self._merge_polar_scans(scans[k], "%s/%s"%(self.outdir, base))

  def copy_files(self, files, outdir):
    for f in files:
      try:
        self.make_dir(outdir)
        shutil.copy(f, "%s/"%outdir)
      except Exception as e:
        print("Failed to copy %s to %s/"%(f,outdir))
        import traceback
        traceback.print_exc()
      
    

  ##
  # Merges polar volumes and places the result in outdir
  # @param pvols a list of filenames pointing at polar voluemes
  # @param outdir the directory where result should be placed
  def _merge_polar_volumes(self, pvols, outdir):
    import _rave
    _rave.setDebugLevel(_rave.Debug_RAVE_SPEWDEBUG)
    bname = os.path.basename(pvols[0])
    bname = "%s.h5"%self.bnamepattern.match(bname).group(1)

    self.make_dir(outdir)

    rio = _raveio.new()
    try:
      pm = polar_merger.polar_merger()
      rio.object = pm.merge_files(pvols)
      rio.save("%s/%s"%(outdir, bname))
    except Exception as e:
      print("Failed to merge %s"%str(pvols))
      if self.copy_ignored:
        self.copy_files(pvols, outdir)


  ##
  # Merges polar scans and places the result in outdir
  # @param scans a list of filenames pointing at polar scans
  # @param outdir the directory where result should be placed
  def _merge_polar_scans(self, scans, outdir):
    import _rave
    _rave.setDebugLevel(_rave.Debug_RAVE_SPEWDEBUG)
    bname = os.path.basename(scans[0])
    bname = "%s.h5"%self.bnamepattern.match(bname).group(1)

    self.make_dir(outdir)

    rio = _raveio.new()
    try:
      pm = polar_merger.polar_merger()
      rio.object = pm.merge_files(scans)
      rio.save("%s/%s"%(outdir, bname))
    except Exception as e:
      print("Failed to merge %s"%str(scans))
      if self.copy_ignored:
        self.copy_files(scans, outdir)

if __name__=="__main__":
  #for root, d, files in os.walk("archive"):
  #  for item in fnmatch.filter(files, "*.h5"):
  #    print "%s/%s"%(root,item)
  optlist = []
  args = []
  sources=None
  ftype="SCAN"
  indir="./archive"
  outdir="./out"
  copy_ignored=False
  try:
    optlist, args = getopt.getopt(sys.argv[1:], '', 
                                  ['indir=', 'outdir=', 'sources=','type=','copy-ignored'])
  except getopt.GetoptError as e:
    print(e.__str__())
    sys.exit(127)

  for o, a in optlist:
    if o == "--indir":
      indir=a
    elif o == "--outdir":
      outdir=a
    elif o == "--sources":
      sources = string.split(a,',')
    elif o == "--type":
      if a.lower() in ["pvol", "scan"]:
        ftype = a
      else:
        raise TypeError("type must be pvol or scan")
    elif o == "--copy-ignored":
      copy_ignored=True

  if not os.path.exists(outdir):
    raise IOError("Outdir %s does not exist"%outdir)
 
  mf = merge_files(indir, outdir, sources, ftype, copy_ignored)
  mf.find_and_merge()

