'''
The MIT License (MIT)

Copyright (c) <2014> <Mathias Lesche>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.

contact: mat.lesche(at)gmail.com
'''

''' python modules '''
import logging
import pysam

from collections import defaultdict

from datetime import date

from functools import partial

from hashlib import md5 as md5

from numpy import mean

from os.path import basename
from os.path import isfile
from os.path import getsize

from random import randint

from re import compile

from subprocess import Popen as Popen
from subprocess import PIPE as PIPE

from sys import _getframe as getframe  # @UnusedImport


''' own modules '''
from database.database import Database

from io_module import get_fileobject


'''
Class used as information set-up box for the files. It tries to extract
the library id, track id and read from the filename. If library doesn't work,
fallback will be to use the everything before the first _ in the name of the
file or the first ".".
If a connection to the database exists sample id and platform is extracted
by using the lib id.
Read group is set to random id (date, random id (100000-999999))
'''

class Fileclass(object):
  def __init__(self, filename, dbinst = None, ext = ''):
    self.__completepath = filename
    self.__filename = basename(filename)

    self.__logger = logging.getLogger('dsp.support.file')

    self.__libreg_read = compile('R[0-9]+(-R[0-9]+)*')
    self.__libreg_track = compile('Track(-[0-9]+)+')
    self.__libreg_lib = compile('L[0-9]+(-L[0-9]+)*')

    self.__libid = []
#     used if it is not possible to identify a lib id
    self.__alternative = False
    self.__sample = []
    self.__platform = []
    self.__track = []
    self.__lane = []
    self.__read = []
    self.__readcount = 0
    self.__size = 0
    self.__bases = 0
    self.__md5sum = ''
    self.__filetype = ''
    
    self.build_information()

    if self.__filename.endswith(('gz', 'zip', '7z', 'rar', 'tar.gz')):
      self.__zipped = True
    else:
      self.__zipped = False

    self.__dbinst = dbinst

    self.__rg = ''
    self.__ext = ext
    
#     if sample and platform information is available
    self.__use = False

    self.__module = ''


  def show_log(self, level, message):
    if level == 'debug':
      self.__logger.debug(message)
    elif level == 'info':
      self.__logger.info(message)
    elif level == 'warning':
      self.__logger.warning(message)
    elif level == 'error':
      self.__logger.error(message)
    elif level == 'critical':
      self.__logger.critical(message)

  def identify_lib(self):
    temp = self.__libreg_lib.match(self.__filename) 
    try:
      self.__libid = temp.group(0).split('-')
    except AttributeError:
      self.__alternative = True
      if '_' in self.__filename:
        self.__libid = [self.__filename.split('_')[0]]
      else:
        self.__libid = [self.__filename.split('.', 1)[0]]
    
  def identify_track(self):
    temp = self.__libreg_track.search(self.__filename) 
    try:
      self.__track = temp.group(0).split('-')[1:]
    except AttributeError:
      self.__track = []

  def identify_read(self):
    temp = self.__libreg_read.search(self.__filename) 
    try:
      self.__read = temp.group(0).split('-')
    except AttributeError:
      self.__read = []

  def build_information(self):
    self.identify_lib()
    self.identify_track()
    self.identify_read()
    self.calc_size()
    self.identify_filetype()
    if self.__alternative == True:
      self.show_log('warning', "File: {0} - Alternative Lib identification".format(self.__filename))
    elif len(self.__read) != 1:
      self.show_log('warning', "File: {0} - Read direction identification not possible".format(self.__filename))  
#     else:
#       self.show_log('debug', "File: {0} - with LB:{1} Track:{2} R:{3}".format(self.__filename, self.__libid, self.__track, self.__read))

  def show_information(self):
    msg = "File Information: '{0}' with LB:{1} SM:{2} PL:{3} R:{4} Tracks: {5} Reads: {6}".format(self.__filename, self.__libid, self.__sample, self.__platform, self.__read, len(self.__track), self.__readcount)
    self.show_log('info', msg)

  def set_information(self):
    if isinstance(self.__dbinst, Database):
      for lib in self.__libid:
        information = self.__dbinst.query_Library_with_Libid(lib[1:])
        if len(information) != 0:
          platform = self.__dbinst.query_Platforms_with_Platformid(information[0]['PLATFORMS_ID'])
          if len(platform) != 0:
            self.__platform.append(platform[0]['NAME'])
          sample = information[0]['NAME']
#           sample = self.__dbinst.query_Library_with_Libid_get_Samplename(lib[1:])
          if len(sample) != 0:
            self.__sample.append(sample)
#             self.__sample.append(sample[0]['NAME'])
      self.__platform = list(set(self.__platform))
      self.__sample = list(set(self.__sample))
      self.__libid = list(set(self.__libid))
      
#       only works of db inst exists
      fqdbinfo = self.__dbinst.query_FQstorage_with_filename_directory(self.__filename)
      if len(fqdbinfo) == 0:
        readcount = 0
        for i in self.__track:
          fqdbinfo = self.__dbinst.query_FQstorage_with_trackid(i)
          if len(fqdbinfo) != 0: readcount += fqdbinfo[0]['READCOUNT']
        self.__readcount = readcount
      else:
        self.__readcount = fqdbinfo[0]['READCOUNT']

    if len(self.__sample) != 0 and len(self.__platform) != 0:
      self.__use = True

    self.create_rg()
    self.show_information()

  @staticmethod
  def get_jobidentifier(inputfiles):
    libid = '-'.join(sorted(set(['-'.join(sorted(i.libid)) for i in inputfiles if i.libid != []])))
    sample = '-'.join(sorted(set(['-'.join(sorted(i.sample)) for i in inputfiles if i.sample != []])))
    read = '-'.join(sorted(set(['-'.join(sorted(i.read)) for i in inputfiles if i.read != []])))
    track = '-'.join(sorted(set(['-'.join(sorted(i.track)) for i in inputfiles if i.track != []])))
    alternative = '-'.join(sorted(set(['-'.join(sorted(i.alternative)) for i in inputfiles if i.alternative != []])))
    jobid = '_'.join([i for i in (libid, alternative, sample, track, read) if i != ''])
    if jobid == '':
      jobid = [i.filename.split('.', 1) for i in inputfiles]
    return jobid
  
  '''
  files are copied into a dictionary
  @param inputfiles: list of file objects
  @param sample: boolean
  '''
  @staticmethod
  def sort_files(inputfiles, sample = False):
    if sample:
      libdict = defaultdict(partial(defaultdict, partial( defaultdict, partial( defaultdict, list))))
      for fileinst in inputfiles:
        libdict['-'.join(sorted(set(fileinst.sample)))]['-'.join(sorted(set(fileinst.libid)))]['-'.join(sorted(set(fileinst.read)))]['-'.join(fileinst.track)].append(fileinst)
    else:
      libdict = defaultdict( partial( defaultdict, partial( defaultdict, list)))
      for fileinst in inputfiles:
        libdict['-'.join(sorted(set(fileinst.libid)))]['-'.join(sorted(set(fileinst.read)))]['-'.join(fileinst.track)].append(fileinst)
    return libdict

  '''
  method calculates the md5sum
  '''
  def calculate_md5sum(self):
    self.__md5sum = md5(open(self.__completepath, 'rb').read()).hexdigest()
    self.show_log('debug', "{0} - md5sum: {1}".format(self.__filename, self.__md5sum))
  
  '''
  method calculates the lines in a fastq file by calling shell commands
  '''
  def calculate_reads(self):
    if self.__completepath.endswith(('fastq.gz', 'fq.gz')):
      call = Popen(['zcat', self.__completepath], stdout=PIPE)
      countcall = Popen(['wc', '-l'], stdin=call.stdout, stdout=PIPE)
      self.__readcount = int(countcall.communicate()[0])/4
      call.stdout.close()
    elif self.__completepath.endswith(('fasta.gz', 'fa.gz')):
      call = Popen(['zcat', self.__completepath], stdout=PIPE)
      countcall = Popen(['wc', '-l'], stdin=call.stdout, stdout=PIPE)
      self.__readcount = int(countcall.communicate()[0])/2
      call.stdout.close()
    elif self.__completepath.endswith(('fastq', 'fq')):
      call = Popen(['cat', self.__completepath], stdout=PIPE)
      countcall = Popen(['wc', '-l'], stdin=call.stdout, stdout=PIPE)
      self.__readcount = int(countcall.communicate()[0])/4
      call.stdout.close()
    elif self.__completepath.endswith(('fasta', 'fa')):
      call = Popen(['cat', self.__completepath], stdout=PIPE)
      countcall = Popen(['wc', '-l'], stdin=call.stdout, stdout=PIPE)
      self.__readcount = int(countcall.communicate()[0])/2
      call.stdout.close()
    elif self.__completepath.endswith(('bam')):
      fileinput = pysam.AlignmentFile(self.__completepath, 'rb', check_header = False, check_sq = False)
      count = 0
      for entry in fileinput: #@UnusedVariable
        count += 1
      self.__readcount = count
    elif self.__completepath.endswith(('xml')):
      self.__readcount = 0
    else:
      self.__readcount = 0
    self.show_log('debug', "{0} - reads: {1}".format(self.__filename, self.__readcount))

  def identify_filetype(self):
    if self.__completepath.endswith(('fastq.gz', 'fq.gz', 'fastq', 'fq')):
      self.__filetype = 'fastq'
    elif self.__completepath.endswith(('fasta.gz', 'fa.gz', 'fasta', 'fa')):
      self.__filetype = 'fasta'
    elif self.__completepath.endswith(('bam')):
      self.__filetype = 'bam'
    elif self.__completepath.endswith(('xml')):
      self.__filetype = 'xml'
    else:
      self.__filetype = 'other'

  def calc_size(self):
    if isfile(self.__completepath): self.__size = getsize(self.__completepath) / 1000000.

  def calc_bases(self):
    counter, temp, seqlen = 0, [], []
    fopen = get_fileobject(self.__completepath, 'r')
    
    if self.__completepath.endswith(('fastq.gz', 'fq.gz', 'fastq', 'fq')):
      for line in fopen:
        temp.append(line.rstrip('\n'))
        temp.append(fopen.next().rstrip('\n'))
        temp.append(fopen.next().rstrip('\n'))
        temp.append(fopen.next().rstrip('\n'))
        if len(temp) == 4:
          counter += 1
          seqlen.append(len(temp[1]))
          temp = []
          if counter == 1000:
            break
    elif self.__completepath.endswith(('fasta.gz', 'fa.gz', 'fasta', 'fa')):
      for line in fopen:
        temp.append(line.rstrip('\n'))
        temp.append(fopen.next().rstrip('\n'))
        if len(temp) == 2:
          counter += 1
          seqlen.append(len(temp[1]))
          temp = []
          if counter == 1000:
            break
    elif self.__completepath.endswith(('bam')):
      fileinput = pysam.AlignmentFile(self.__completepath, 'rb', check_header = False, check_sq = False)
      for entry in fileinput:
        seqlen.append(len(entry.query_sequence))
    
    if len(seqlen) != 0 and self.__filetype != 'bam':
      self.__bases = (self.__readcount * mean(seqlen)) / 1000000.
    elif len(seqlen) != 0 and self.__filetype == 'bam':
      self.__bases = sum(seqlen) / 1000000.

  def get_alternative(self):
    return self.__alternative

  def get_completepath(self):
    return self.__completepath

  def get_ext(self):
    return self.__ext

  def get_filename(self):
    return self.__filename

  def get_libid(self):
    return self.__libid

  def get_module(self):
    return self.__module

  def get_platform(self):
    return self.__platform

  def get_dbinst(self):
    return self.__dbinst

  def get_lane(self):
    return self.__lane

  def get_read(self):
    return self.__read

  def get_readcount(self):
    return self.__readcount

  def get_rg(self):
    return self.__rg

  def get_sample(self):
    return self.__sample

  def get_track(self):
    return self.__track

  def get_use(self):
    return self.__use

  def get_zipped(self):
    return self.__zipped

  def get_md5sum(self):
    return self.__md5sum

  def get_bases(self,):
    return self.__bases
  
  def get_size(self,):
    return self.__size

  def get_filetype(self):
    return self.__filetype

  def set_alternative(self, alternative):
    self.__alternative = alternative

  def set_ext(self, ext):
    self.__ext = ext

  def set_filename(self, filename):
    self.__completepath = filename
    self.__filename = basename(filename)

  def set_libid(self, libid):
    self.__libid = libid

  def set_module(self, value):
    self.__module = value

  def set_platform(self, platform):
    self.__platform = platform

  def set_md5sum(self, md5sum):
    self.__md5sum = md5sum

  def create_rg(self):
#     if len(self.__track) != 0:
#       track = 'Track-'+'-'.join(self.__track)
#     else:
#       track = ''
#     libid = '-'.join(self.__libid)
#     sample = '-'.join(self.__sample)
#     read = '-'.join(self.__read)
    
    rgstring = '{0}_{1}'.format(str(date.today()), randint(100000, 999999))
#     rgstring = '_'.join([i for i in (libid, sample, track) if i != ''])
    
    if self.__ext == '':
      self.__rg = '{0}'.format(rgstring)
    else:
      self.__rg = '{0}.{1}'.format(rgstring, self.__ext)

  def set_readcount(self, readcount):
    self.__readcount = readcount

  def set_bases(self, bases):
    self.__bases = bases
  
  def set_size(self, size):
    self.__size = size

  def set_dbinst(self, dbinst):
    self.__dbinst = dbinst

  def set_sample(self, sample):
    self.__sample = sample

  def set_track(self, track):
    self.__track = track
    
  def set_lane(self, lane):
    self.__lane = lane

  def set_use(self, use):
    self.__use = use

  def set_zipped(self, zipped):
    self.__zipped = zipped
  
  def set_rg(self, rgstring):
    self.__rg = rgstring

  def set_filetype(self, filetype):
    self.__filetype = filetype

  md5sum = property(get_md5sum, set_md5sum)
  filetype = property(get_filetype, set_filetype)
  alternative = property(get_alternative, set_alternative)
  bases = property(get_bases, set_bases)
  size = property(get_size, set_size)
  completepath = property(get_completepath)
  ext = property(get_ext, set_ext)
  filename = property(get_filename, set_filename)
  libid = property(get_libid, set_libid, None, None)
  platform = property(get_platform, set_platform, None, None)
  dbinst = property(get_dbinst, set_dbinst)
  readcount = property(get_readcount, set_readcount)
  rg = property(get_rg, set_rg, None, None)
  lane = property(get_lane, set_lane)
  read = property(get_read, None, None, None)
  sample = property(get_sample, set_sample, None, None)
  track = property(get_track, set_track, None, None)  
  use = property(get_use, set_use, None, None)
  zipped = property(get_zipped, set_zipped)
  module = property(get_module, set_module, None, None)