#!/usr/bin/env python
"""
Tasks here prepare resources that are required for some
pipeline tasks, and are generally available from the
resources directory as specified by $GAX_RESOURCES or
in genairics config file 'resourcedir' option
"""
from datetime import datetime, timedelta
import luigi, os, sys, tempfile, pathlib, glob
from luigi.util import inherits, requires
from plumbum import local, colors
import logging

from genairics import config, gscripts

resourcedir = config.resourcedir
logresources = logging.Logger('genairics.resources')

#Set cksum program according to platform
if sys.platform == 'linux':
    cksum = local['sum']
elif sys.platform == 'darwin':
    cksum = local['cksum']['-o',1]
else: raise Exception('{} not supported. Run from docker container.'.format(sys.platform))

def requestFiles(urlDir,fileregex,outdir):
    """
    Download a set of files that match fileregex from urlDir
    checksums are expected to be generated by the "sum" command,
    as is the case for the ftp.ensembl.org repository
    """
    import requests, re
    link = re.compile(r'href="(.+?)"')
    if fileregex: fileregex = re.compile(fileregex)
    if not urlDir.endswith('/'): urlDir+='/'
    # Download checksums if present
    checksums = requests.get(urlDir+'CHECKSUMS')
    if checksums:
        from io import StringIO
        import pandas as pd
        checksums = pd.read_table(StringIO(checksums.text),sep='\s+',
                                  names=['checksum', 'octets', 'name'],index_col='name')
    else: checksums = None
    # Retrieve index
    r = requests.get(urlDir)
    for l in link.finditer(r.text):
        l = l.groups()[0]
        if not fileregex or fileregex.match(l):
            file_request = requests.get(urlDir+l, stream=True)
            with open(os.path.join(outdir,l),'wb') as outfile:
                for chunk in file_request.iter_content(chunk_size=512):
                    outfile.write(chunk)
            if checksums is not None:
                csum = int(cksum(os.path.join(outdir,l)).split()[0])
                if csum != checksums.checksum[l]:
                    logresources.warning('%s checksum did not match url location %s',csum,urlDir)

# Tasks
class InstallDependencies(luigi.Task):
    """
    Installs the genairics dependency programs
    """
    def run(self):
        local[gscripts % 'genairics_dependencies.sh']
    
class RetrieveGenome(luigi.Task):
    """
    Prepares the genome
    """
    genome = luigi.Parameter(default='homo_sapiens',description="genome species name")
    release = luigi.IntParameter(default=91,
                                 description=
"""ensembl release number of genome
for homo_sapiens, use the following :
 91 => GRCh38 aka hg38;
 75 => GRCh37 aka hg19
"""
    )

    def output(self):
        return luigi.LocalTarget(
            os.path.join(resourcedir,'ensembl',self.genome,'release-{}'.format(self.release))
        )

    def run(self):
        #Make temp dir for data
        local['mkdir']('-p',self.output().path+'_retrieving/dna')
        local['mkdir']('-p',self.output().path+'_retrieving/annotation')
        requestFiles(
            'http://ftp.ensembl.org/pub/release-{release}/fasta/{species}/dna/'.format(
                species=self.genome, release=self.release),
            r'.+.dna.chromosome.+.fa.gz',
            self.output().path+'_retrieving/dna'
        )
        requestFiles(
            'http://ftp.ensembl.org/pub/release-{release}/gtf/{species}/'.format(
                species=self.genome, release=self.release),
            r'.+.{release}.gtf.gz'.format(release=self.release),
            self.output().path+'_retrieving/annotation'
        )
        #Unzip all files
        local['gunzip'](*glob.glob(self.output().path+'_retrieving/*/*.gz'))
        #Rename temp dir to final/expected output dir
        os.rename(self.output().path+'_retrieving',self.output().path)

@requires(RetrieveGenome)
class RetrieveBlacklist(luigi.Task):
    """
    Info on blacklists: https://sites.google.com/site/anshulkundaje/projects/blacklists

    Available release:
    91 => GRCh38
    75 => GRCh37 (latest release of hg19)
    """
    availableBlacklists = {
        ('homo_sapiens', 91): "http://mitra.stanford.edu/kundaje/akundaje/release/blacklists/hg38-human/hg38.blacklist.bed.gz",
        ('homo_sapiens', 75):
        "http://mitra.stanford.edu/kundaje/akundaje/release/blacklists/hg19-human/wgEncodeHg19ConsensusSignalArtifactRegions.bed.gz"
    }
    
    def output(self):
        return luigi.LocalTarget(os.path.join(self.input().path,'blacklist.bed.gz'))

    def run(self):
        if (self.genome, self.release) not in self.availableBlacklists:
            raise NotImplementedError("blacklist for %s release %s is not available yet within genairics" % (self.genome,self.release))
        stdout = local['wget'](self.availableBlacklists[(self.genome, self.release)], '-O', self.output().path)
        logresources.info(stdout)

@requires(RetrieveGenome)
class Bowtie2Index(luigi.Task):
    """
    Index that can be used by bowtie2 aligner for ChIPseq, ATACseq, variant calling
    """
    def output(self):
        return (
            luigi.LocalTarget(
                os.path.join(resourcedir,'ensembl/{species}/release-{release}/genome_index'.format(
                    species=self.genome,release=self.release))
            ),
            luigi.LocalTarget(
                os.path.join(resourcedir,'ensembl/{species}/release-{release}/genome_index/build_completed'.format(
                    species=self.genome,release=self.release))
            )
        )
    
    def run(self):
        genomeDir = self.input().path
        os.mkdir(self.output()[0].path)
        stdout = local['bowtie2-build'](
            '--threads', config.threads,
            ','.join(glob.glob(os.path.join(genomeDir,'dna')+'/*.fa*')),
            os.path.join(self.output()[0].path, self.genome)
        )
        logresources.info(stdout)
        pathlib.Path(self.output()[1].path).touch()        

@inherits(RetrieveGenome)
class STARandRSEMindex(luigi.Task):
    """
    Index that can be used by both STAR aligner and RSEM counter for transcriptomics
    """
    def requires(self):
        return self.clone_parent()
    
    def output(self):
        return (
            luigi.LocalTarget(
                os.path.join(resourcedir,'ensembl/{species}/release-{release}/transcriptome_index'.format(
                    species=self.genome,release=self.release))
            ),
            luigi.LocalTarget(
                os.path.join(resourcedir,'ensembl/{species}/release-{release}/transcriptome_index/build_completed'.format(
                    species=self.genome,release=self.release))
            )
        )
    
    def run(self):
        genomeDir = self.input().path
        os.mkdir(self.output()[0].path)
        stdout = local[gscripts % 'buildRSEMindex.sh'](
            *glob.glob(os.path.join(genomeDir,'annotation')+'/*.gtf'),
            config.threads,
            ','.join(glob.glob(os.path.join(genomeDir,'dna')+'/*.fa*')),
            os.path.join(self.output()[0].path, self.genome)
        )
        logresources.info(stdout)
        pathlib.Path(self.output()[1].path).touch()

@inherits(RetrieveGenome)
class STARindex(luigi.Task):
    """
    Index that can be used by STAR aligner for genome mapping 
    (for genomic variant discovery, ChiPseq, ATACseq, )

    NOT YET TESTED, for the moment use STARandRSEMindex
    """
    sjdbOverhang = luigi.IntParameter(default=100, description='library prep read length. default of 100 should generally be ok')
    
    def requires(self):
        return self.clone_parent()
    
    def output(self):
        return (
            luigi.LocalTarget(
                os.path.join(resourcedir,'ensembl/{species}/release-{release}/genome_index_overhang{overhang}'.format(
                    species=self.genome,release=self.release,overhang=self.sjdbOverhang))
            ),
            luigi.LocalTarget(
                os.path.join(resourcedir,'ensembl/{species}/release-{release}/genome_index_overhang{overhang}/build_completed'.format(
                    species=self.genome,release=self.release,overhang=self.sjdbOverhang))
            )
        )
    
    def run(self):
        genomeDir = self.input().path
        os.mkdir(self.output()[0].path)
        stdout = local['STAR'](
            '--runThreadN', config.threads,
            '--runMode', 'genomeGenerate',
            '--genomeDir', os.path.join(self.output()[0].path, self.genome),
            '--genomeFastaFiles', ','.join(glob.glob(os.path.join(genomeDir,'dna')+'/*.fa*')),
            '--sjdbOverhang', self.sjdbOverhang
            
        )
        #TODO STAR Log.out will be produced should be moved to genome index dir
        logresources.info(stdout)
        pathlib.Path(self.output()[1].path).touch()
