#!/usr/bin/env python3
""" THIS MODULE CONTAINS ALL THE SHARED WRAPPER FUNCTIONS """
################################################################################
#                              CGE FUNCTION MODULE                             #
################################################################################
# This script is part of the CGE Pipeline structure
import sys, os, gzip, shutil, glob, re, json
from subprocess import Popen
from zipfile import ZipFile
from contextlib import closing

############# CLASSES #############
class Debug():
   """ Debug object, keeps track of all debug related matters, and provides easy
   access to logging debug text

   USAGE
      import os, sys
      debug = Debug()
      debug.setup(debug=True, logfile=None, stdout=None, stderr=None)
      debug.print_out('hello')
      debug.log_no_newline('Oh')
      debug.log_no_newline(' my')
      debug.log(' god!')
      # The graceful exit is only meant to be used with the integrated CGE platform
      debug.graceful_exit('error message')
   IMPORTING THE DEBUG OBJECT
      from cge.utility import debug
   """
   def __init__(self):
      """  """
      self.debug = False
      self.logfile = sys.stderr
      self.stdout = sys.stdout
      self.stderr = sys.stderr
      self.caught_error = None
   def setup(self, debug=None, logfile=None, stdout=None, stderr=None):
      """  """
      if debug   is not None: self.debug   = debug
      if logfile is not None: self.logfile = logfile
      if stdout  is not None: self.stdout  = stdout
      if stderr  is not None: self.stderr  = stderr
   def print_out(self, *lst):
      """ Print list of strings to the predefined stdout. """
      self.print2file(self.stdout, True, True, *lst)
   def print_err(self, *lst):
      """ Print list of strings to the predefined stdout. """
      self.print2file(self.stderr, False, True, *lst)
   def print2file(self, logfile, print2screen, addLineFeed, *lst):
      """ This function prints to the screen and logs to a file, all the strings
      given.
      # print2screen eg. True, *lst is a commaseparated list of strings
      """
      if addLineFeed:
         linefeed = '\n'
      else: linefeed = ''
      if print2screen: print(linefeed.join(str(string) for string in lst))
      try: file_instance = isinstance(logfile, file)
      except NameError as e:
         from io import IOBase
         try: file_instance = isinstance(logfile, IOBase)
         except: raise e
      if file_instance:
         logfile.write(linefeed.join(str(string) for string in lst) + linefeed)
      elif isinstance(logfile, str) and os.path.exists(logfile):
         with open_(logfile, 'a') as f:
            f.write(linefeed.join(str(string) for string in lst) + linefeed)
      elif not print2screen: # Print to screen if there is no outputfile
         print(linefeed.join(str(string) for string in lst))
   def log(self, *lst):
      """ Print list of strings to the predefined logfile if debug is set. and
      sets the caught_error message if an error is found
      """
      self.print2file(self.logfile, self.debug, True, *lst)
      if 'Error' in '\n'.join([str(x) for x in lst]):
         self.caught_error = '\n'.join([str(x) for x in lst])
   def log_no_newline(self, msg):
      """ print the message to the predefined log file without newline """
      self.print2file(self.logfile, False, False, msg)
   def graceful_exit(self, msg):
      """ This function Tries to update the MSQL database before exiting. """
      # Print stored errors to stderr
      if self.caught_error:
         self.print2file(self.stderr, False, False, self.caught_error)
      # Kill process with error message
      self.log(msg)
      sys.exit(1)

class adv_dict(dict):
   """ This class expands on the dictionary class by adding the gettree class
   method.
   """
   def get_tree(self, list_of_keys):
      """ gettree will extract the value from a nested tree
      
      INPUT
         list_of_keys: a list of keys ie. ['key1', 'key2']
      USAGE
      >>> # Access the value for key2 within the nested dictionary
      >>> adv_dict({'key1': {'key2': 'value'}}).gettree(['key1', 'key2'])
      'value'
      """
      cur_obj = self
      for key in list_of_keys:
         cur_obj = cur_obj.get(key)
         if not cur_obj: break
      return cur_obj
   def invert(self):
      ''' Return inverse mapping of dictionary with sorted values.
      USAGE
         >>> # Switch the keys and values
         >>> adv_dict({
         ...     'A': [1, 2, 3],
         ...     'B': [4, 2],
         ...     'C': [1, 4],
         ... }).invert()
         {1: ['A', 'C'], 2: ['A', 'B'], 3: ['A'], 4: ['B', 'C']}
      '''
      inv_map = {}
      for k, v in self.items():
         if sys.version_info < (3, 0):
            acceptable_v_instance = isinstance(v, (str, int, float, long))
         else:
            acceptable_v_instance = isinstance(v, (str, int, float))
         if acceptable_v_instance: v = [v]
         elif not isinstance(v, list):
            raise Exception('Error: Non supported value format! Values may only'
                            ' be numerical, strings, or lists of numbers and '
                            'strings.')
         for val in v:
            inv_map[val] = inv_map.get(val, [])
            inv_map[val].append(k)
            inv_map[val].sort()
      return inv_map

class Reg:
   """
   NAME:        Reg - Extended Regular Expression Handler
   AUTHOR:      Martin Thomsen
   DESCRIPTION:
      This class enables a simplistic usage of regular expression to get
      contained groups in a match statement. But it also allows to do some of
      the normal re call, such as findall and sub.
   DEPENDENCIES:
      re (regular expression module)
   USAGE:
      >>> import re
      >>> RegEx = Reg(pattern, flag)
      >>> if RegEx.match(string):
      >>>    RegEx.getgroup(index)
   EXAMPLE:
      >>> RegEx = Reg('[^a]*(a)[^b]*(b)[^Y]*(Y)(a)?', 'I')
      >>> if RegEx.match('aBcdefgHIJKLmnOpqrstuvwxyz'):
      ...    print(RegEx.getgroup(0), # index=0 -> full match
      ...          RegEx.getgroup(1),
      ...          RegEx.getgroup(2),
      ...          RegEx.getgroup(3),
      ...          RegEx.getgroup(4))
      ...
      ('aBcdefgHIJKLmnOpqrstuvwxy', 'a', 'B', 'y', None)
      # NIFTY SUBSTITUTION LOOP
      >>> string = 'There are {%=count%} {%=animal%} on the {%=location%}!'
      >>> # Dictionary containing place-holders and values
      ... # (make sure all placeholders in the string is included!)
      ... d = { 'count': 5, 'animal': 'cows', 'location': 'Battle Field' }
      >>> # RE Object which matches place-holders in the string
      ... tmpPH = Reg('\{\%\=(\w+)\%\}', 'I')
      >>> # substitute all placeholders
      ... while tmpPH.match(string): string = tmpPH.sub(str(d[tmpPH.getgroup(1)]), string, 1)
      ...
      >>> print(string)
      There are 5 cows on the Battle Field!
      """
   def __init__(self, pattern, *flags):
      sd = {'T': 1, 'I': 2, 'L':4,'M':8,'S':16,'U':32,'X':64}
      try:
         flag = sum([sd[f] if f in sd else int(f) for f in set(flags)]) if flags else 0
      except:
         for f in flags:
            if not isinstance(f, int) and not f in sd:
               raise Exception("Error: Unrecognised flag argument '%s' for Reg call."%f)
         flag=0
      if flag: self.re = re.compile(pattern, flag)
      else: self.re = re.compile(pattern)
      self.matches = None
   def sub(self, replace, string, count=0):
      """ returns new string where the matching cases (limited by the count) in
      the string is replaced. """
      return self.re.sub(replace, string, count)
   def find_all(self, s):
      """ Finds all matches in the string and returns them in a tuple. """
      return self.re.findall(s)
   def match(self, s):
      """ Matches the string to the stored regular expression, and stores all
      groups in mathches. Returns False on negative match. """
      self.matches = self.re.search(s)
      return self.matches
   def get_group(self, x):
      """ Returns requested subgroup. """
      return self.matches.group(x)
   def get_groups(self):
      """ Returns all subgroups. """
      return self.matches.groups()

class REGroup():
   """ Regular Expression group object

   This class simplyfies the use of groups for the Sort2Groups function.
   """
   def __init__(self, pattern, flags=''):
      self.re = Reg(pattern, flags)
      self.list = []
   def match(self, s):
      """ Matching the pattern to the input string, returns True/False and
          saves the matched string in the internal list
      """
      if self.re.match(s):
         self.list.append(s)
         return True
      else: return False

############# ITERATORS #############
def seqs_from_file(filename, exit_on_err=False, return_qual=False):
   """Extract sequences from a file
   
   Name:
      seqs_from_file
   Author(s):
      Martin C F Thomsen
   Date:
      18 Jul 2013
   Description:
      Iterator which extract sequence data from the input file
   Args:
      filename: string which contain a path to the input file
   Supported Formats:
      fasta, fastq
   
   USAGE:
   >>> import os, sys
   >>> # Create fasta test file
   >>> file_content = ('>head1 desc1\nthis_is_seq_1\n>head2 desc2\n'
                       'this_is_seq_2\n>head3 desc3\nthis_is_seq_3\n')
   >>> with open_('test.fsa', 'w') as f: f.write(file_content)
   >>> # Parse and print the fasta file
   >>> for seq, name, desc in SeqsFromFile('test.fsa'):
   ...    print ">%s %s\n%s"%(name, desc, seq)
   ...
   >head1 desc1
   this_is_seq_1
   >head2 desc2
   this_is_seq_2
   >head3 desc3
   this_is_seq_3
   """
   # VALIDATE INPUT
   if not isinstance(filename, str):
      msg = 'Filename has to be a string.'
      if exit_on_err:
         sys.stderr.write('Error: %s\n'%msg)
         sys.exit(1)
      else: raise IOError(msg)
   if not os.path.exists(filename):
      msg = 'File "%s" does not exist.'%filename
      if exit_on_err:
         sys.stderr.write('Error: %s\n'%msg)
         sys.exit(1)
      else: raise IOError(msg)
   
   # EXTRACT DATA
   with open_(filename,"rt") as f:
      query_seq_segments = []
      seq, name, desc, qual = '', '', '', ''
      add_segment = query_seq_segments.append
      for l in f:
         if len(l.strip()) == 0: continue
         #sys.stderr.write("%s\n"%line)
         fields=l.strip().split()
         if l.startswith(">"):
            # FASTA HEADER FOUND
            if query_seq_segments != []:
               # YIELD SEQUENCE AND RESET
               seq = ''.join(query_seq_segments)
               yield (seq, name, desc)
               seq, name, desc = '', '', ''
               del query_seq_segments[:]
            name = fields[0][1:]
            desc = ' '.join(fields[1:])
         
         elif l.startswith("@"):
            # FASTQ HEADER FOUND
            name = fields[0][1:]
            desc = ' '.join(fields[1:])
            try:
               # EXTRACT FASTQ SEQUENCE
               seq  = next(f).strip().split()[0]
               # SKIP SECOND HEADER LINE AND QUALITY SCORES
               l = next(f)
               qual = next(f).strip() # Qualities
            except:
               break
            else:
               # YIELD SEQUENCE AND RESET
               if return_qual:
                  yield (seq, qual, name, desc)
               else:
                  yield (seq, name, desc)
               seq, name, desc, qual = '', '', '', ''
         
         elif len(fields[0])>0:
            # EXTRACT FASTA SEQUENCE
            add_segment(fields[0])
      
      # CHECK FOR LAST FASTA SEQUENCE
      if query_seq_segments != []:
         # YIELD SEQUENCE
         seq = ''.join(query_seq_segments)
         yield (seq, name, desc)

############# FUNCTIONS #############
def open_(filename, mode=None, compresslevel=9):
   """Switch for both open() and gzip.open().
   
   Determines if the file is normal or gzipped by looking at the file
   extension.
   
   The filename argument is required; mode defaults to 'rb' for gzip and 'r'
   for normal and compresslevel defaults to 9 for gzip.
   
   >>> import gzip
   >>> from contextlib import closing
   >>> with open_(filename) as f:
   ...     f.read()
   """
   if filename[-3:] == '.gz':
      if mode is None: mode = 'rt'
      return closing(gzip.open(filename, mode, compresslevel))
   else:
      if mode is None: mode = 'r'
      return open(filename, mode)

def load_json(json_object):
   ''' Load json from file or file name '''
   content = None
   if isinstance(json_object, str) and os.path.exists(json_object):
      with open_(json_object) as f:
         try:
            content = json.load(f)
         except Exception as e:
            debug.log("Warning: Content of '%s' file is not json."%f.name)
   elif hasattr(json_object, 'read'):
      try:
         content = json.load(json_object)
      except Exception as e:
         debug.log("Warning: Content of '%s' file is not json."%json_object.name)
   else:
      debug.log("%s\nWarning: Object type invalid!"%json_object)
   return content

def sort2groups(array, gpat=['_R1','_R2']):
   """ Sort an array of strings to groups by patterns """
   groups = [REGroup(gp) for gp in gpat]
   unmatched = []
   for item in array:
      matched = False
      for m in groups:
         if m.match(item):
            matched = True
            break
      if not matched: unmatched.append(item)
   return [sorted(m.list) for m in groups], sorted(unmatched)

def sort_and_distribute(array, splits=2):
   """ Sort an array of strings to groups by alphabetically continuous
       distribution
   """
   if not isinstance(array, (list,tuple)): raise TypeError("array must be a list")
   if not isinstance(splits, int): raise TypeError("splits must be an integer")
   remaining = sorted(array)
   if sys.version_info < (3, 0):
      myrange = xrange(splits)
   else:
      myrange = range(splits)
   groups = [[] for i in myrange]
   while len(remaining) > 0:
      for i in myrange:
         if len(remaining) > 0: groups[i].append(remaining.pop(0))
   return groups

def mkpath(filepath, permissions=0o777):
   """ This function executes a mkdir command for filepath and with permissions
   (octal number with leading 0 or string only)
   # eg. mkpath("path/to/file", "0o775")
   """
   # Converting string of octal to integer, if string is given.
   if isinstance(permissions, str):
      permissions = sum([int(x)*8**i for i,x in enumerate(reversed(permissions))])
   # Creating directory
   if not os.path.exists(filepath):
      debug.log("Creating Directory %s (permissions: %s)"%(
         filepath, permissions))
      os.makedirs(filepath, permissions)
   else:
      debug.log("Warning: The directory "+ filepath +" already exists")
   return filepath

def create_zip_dir(zipfile_path, *file_list):
   """ This function creates a zipfile located in zipFilePath with the files in
   the file list
   # fileList can be both a comma separated list or an array
   """
   try:
      if isinstance(file_list, (list, tuple)): #unfolding list of list or tuple
         if len(file_list) == 1:
            if isinstance(file_list[0], (list, tuple)): file_list = file_list[0]
      #converting string to iterable list
      if isinstance(file_list, str): file_list = [file_list]
      if file_list:
         with ZipFile(zipfile_path, 'w') as zf:
            for cur_file in file_list:
               if '/' in cur_file:
                  os.chdir('/'.join(cur_file.split('/')[:-1]))
               elif '/' in zipfile_path:
                  os.chdir('/'.join(zipfile_path.split('/')[:-1]))
               zf.write(cur_file.split('/')[-1])
      else:
         debug.log('Error: No Files in list!',zipfile_path+' was not created!')
   except Exception as e:
      debug.log('Error: Could not create zip dir! argtype: '+
                 str(type(file_list)), "FileList: "+ str(file_list),
                 "Errormessage: "+ str(e))

def file_zipper(root_dir):
   """ This function will zip the files created in the runroot directory and
   subdirectories """
   # FINDING AND ZIPPING UNZIPPED FILES
   for root, dirs, files in os.walk(root_dir, topdown=False):
      if root != "":
         if root[-1] != '/': root += '/'
         for current_file in files:
            filepath = "%s/%s"%(root, current_file)
            try:
               file_size = os.path.getsize(filepath)
            except Exception as e:
               file_size = 0
               debug.log('Error: file_zipper failed to zip following file '+filepath, e)
            # Excluding small files, gzipped files and links
            if (         file_size > 50
                 and     current_file[-3:] != ".gz"
                 and not os.path.islink(filepath)
               ):
               if current_file[-4:] == ".zip":
                  # Unzip file
                  ec = Popen('unzip -qq "%s" -d %s > /dev/null 2>&1'%(filepath, root), shell=True).wait()
                  if ec > 0:
                     debug.log('Error: fileZipper failed to unzip following file %s'%filepath)
                     continue
                  else:
                     ec = Popen('rm -f "%s" > /dev/null 2>&1'%(filepath), shell=True).wait()
                     if ec > 0: debug.log('Error: fileZipper failed to delete the original zip file (%s)'%filepath)
                     filepath = filepath[:-4]
                  # Saving a gzipped version
                  with open_(filepath, 'rb') as f, open_(filepath+".gz", 'wb', 9) as gz:
                     gz.writelines(f)
                  # Deleting old (non-zipped) file
                  try: os.remove(filepath)
                  except OSError as e:
                     debug.log(("WARNING! The file %s could not be "
                                    "removed!\n%s")%(current_file, e))

def file_unzipper(directory):
   """ This function will unzip all files in the runroot directory and
   subdirectories
   """
   debug.log("Unzipping directory (%s)..."%directory)
   #FINDING AND UNZIPPING ZIPPED FILES
   for root, dirs, files in os.walk(directory, topdown=False):
      if root != "":
         orig_dir = os.getcwd()
         os.chdir(directory)
         Popen('gunzip -q -f *.gz > /dev/null 2>&1', shell=True).wait()
         Popen('unzip -qq -o "*.zip" > /dev/null 2>&1', shell=True).wait()
         Popen('rm -f *.zip > /dev/null 2>&1', shell=True).wait()
         os.chdir(orig_dir)

def move_file(src, dst):
   """ this function will simply move the file from the source path to the dest
   path given as input
   """
   # Sanity checkpoint
   src = re.sub('[^\w/\-\.\*]', '', src)
   dst = re.sub('[^\w/\-\.\*]', '', dst)
   if len(re.sub('[\W]', '', src)) < 5 or len(re.sub('[\W]', '', dst)) < 5:
      debug.log("Error: Moving file failed. Provided paths are invalid! src='%s' dst='%s'"%(src, dst))
   else:
      # Check destination
      check = False
      if dst[-1] == '/':
         if os.path.exists(dst):
            check = True # Valid Dir
         else:
            debug.log("Error: Moving file failed. Destination directory does not exist (%s)"%(dst)) #DEBUG
      elif os.path.exists(dst):
         if os.path.isdir(dst):
            check = True # Valid Dir
            dst += '/' # Add missing slash
         else:
            debug.log("Error: Moving file failed. %s exists!"%dst)
      elif os.path.exists(os.path.dirname(dst)):
         check = True # Valid file path
      else:
         debug.log("Error: Moving file failed. %s is an invalid distination!"%dst)
      if check:
         # Check source
         files = glob.glob(src)
         if len(files) != 0:
            debug.log("Moving File(s)...", "Move from %s"%src, "to %s"%dst)
            for file_ in files:
               # Check if file contains invalid symbols:
               invalid_chars = re.findall('[^\w/\-\.\*]', os.path.basename(file_))
               if invalid_chars:
                  debug.graceful_exit(("Error: File %s contains invalid "
                                      "characters %s!"
                                      )%(os.path.basename(file_), invalid_chars))
                  continue
               # Check file exists
               if os.path.isfile(file_):
                  debug.log("Moving file: %s"%file_)
                  shutil.move(file_, dst)
               else:
                  debug.log("Error: Moving file failed. %s is not a regular file!"%file_)
         else: debug.log("Error: Moving file failed. No files were found! (%s)"%src)

def copy_file(src, dst, ignore=None):
   """ this function will simply copy the file from the source path to the dest
   path given as input
   """
   # Sanity checkpoint
   src = re.sub('[^\w/\-\.\*]', '', src)
   dst = re.sub('[^\w/\-\.\*]', '', dst)
   if len(re.sub('[\W]', '', src)) < 5 or len(re.sub('[\W]', '', dst)) < 5:
      debug.log("Error: Copying file failed. Provided paths are invalid! src='%s' dst='%s'"%(src, dst))
   else:
      # Check destination
      check = False
      if dst[-1] == '/':
         if os.path.exists(dst):
            check = True # Valid Dir
         else:
            debug.log("Error: Copying file failed. Destination directory does not exist (%s)"%(dst)) #DEBUG
      elif os.path.exists(dst):
         if os.path.isdir(dst):
            check = True # Valid Dir
            dst += '/' # Add missing slash
         else:
            debug.log("Error: Copying file failed. %s exists!"%dst)
      elif os.path.exists(os.path.dirname(dst)):
         check = True # Valid file path
      else:
         debug.log("Error: Copying file failed. %s is an invalid distination!"%dst)
      if check:
         # Check source
         files = glob.glob(src)
         if ignore is not None: files = [fil for fil in files if not ignore in fil]
         if len(files) != 0:
            debug.log("Copying File(s)...", "Copy from %s"%src, "to %s"%dst) #DEBUG
            for file_ in files:
               # Check file exists
               if os.path.isfile(file_):
                  debug.log("Copying file: %s"%file_) #DEBUG
                  shutil.copy(file_, dst)
               else:
                  debug.log("Error: Copying file failed. %s is not a regular file!"%file_) #DEBUG
         else: debug.log("Error: Copying file failed. No files were found! (%s)"%src) #DEBUG

def copy_dir(src, dst):
   """ this function will simply copy the file from the source path to the dest
   path given as input
   """
   try:
      debug.log("copy dir from "+ src, "to "+ dst)
      shutil.copytree(src, dst)
   except Exception as e:
      debug.log("Error: happened while copying!\n%s\n"%e)

# Initiate Shared Global Objects
debug = Debug()