#!/usr/bin/env python
#
# -----------------------------------------------------------------------------
# Copyright (c) 2015-2016   Daniel Standage <daniel.standage@gmail.com>
# Copyright (c) 2015-2016   Indiana University
#
# This file is part of genhub (http://github.com/standage/genhub) and is
# licensed under the BSD 3-clause license: see LICENSE.txt.
# -----------------------------------------------------------------------------

"""
Retrieve and format data from RefSeq.

GenomeDB implementation for data residing in NCBI's RefSeq database.
"""

from __future__ import print_function
import filecmp
import gzip
import os
import re
import subprocess
import sys
import genhub


class RefSeqDB(genhub.genomedb.GenomeDB):

    def __init__(self, label, conf, workdir='.'):
        super(RefSeqDB, self).__init__(label, conf, workdir)

        assert self.config['source'] == 'refseq'
        assert 'branch' in self.config
        assert 'species' in self.config
        assert 'accession' in self.config
        assert 'build' in self.config

        species = self.config['species'].replace(' ', '_')
        self.acc = self.config['accession'] + '_' + self.config['build']

        base = 'ftp://ftp.ncbi.nlm.nih.gov/genomes/refseq'
        url_parts = [base, self.config['branch'], species,
                     'all_assembly_versions', self.acc]
        self.specbase = '/'.join(url_parts + [self.acc])
        self.format_gdna = self.format_fasta
        self.format_prot = self.format_fasta

    def __repr__(self):
        return 'RefSeq'

    @property
    def gdnafilename(self):
        return '%s_genomic.fna.gz' % self.acc

    @property
    def gff3filename(self):
        return '%s_genomic.gff.gz' % self.acc

    @property
    def protfilename(self):
        return '%s_protein.faa.gz' % self.acc

    @property
    def gdnaurl(self):
        return '%s_genomic.fna.gz' % self.specbase

    @property
    def gff3url(self):
        return '%s_genomic.gff.gz' % self.specbase

    @property
    def proturl(self):
        return '%s_protein.faa.gz' % self.specbase

    def format_fasta(self, instream, outstream, logstream=sys.stderr):
        for line in instream:
            # No processing required since new RefSeq migration (2015-11-30).
            # If any is ever needed again, do it here.
            print(line, end='', file=outstream)

    def format_gff3(self, logstream=sys.stderr, debug=False):
        cmds = list()
        cmds.append('gunzip -c %s' % self.gff3path)
        if 'annotfilter' in self.config:
            excludefile = self.filter_file()
            cmds.append('grep -vf %s' % excludefile.name)
        cmds.append('tidygff3')
        cmds.append('genhub-format-gff3.py --source refseq -')
        if 'fixseqreg' in self.config and self.config['fixseqreg'] is True:
            cmds.append('seq-reg.py - %s' % self.gdnafile)  # pragma: no cover
        cmds.append('gt gff3 -sort -tidy -o %s -force' % self.gff3file)

        commands = ' | '.join(cmds)
        if debug:  # pragma: no cover
            print('DEBUG: running command: %s' % commands, file=logstream)
        proc = subprocess.Popen(commands, shell=True, universal_newlines=True,
                                stderr=subprocess.PIPE)
        stdout, stderr = proc.communicate()
        for line in stderr.split('\n'):  # pragma: no cover
            if 'has not been previously introduced' not in line and \
               'does not begin with "##gff-version"' not in line and \
               'more than one pseudogene attribute' not in line and \
               line != '':
                print(line, file=logstream)
        assert proc.returncode == 0, \
            'annot cleanup command failed: %s' % commands
        if 'annotfilter' in self.config:
            os.unlink(excludefile.name)

    def gff3_protids(self, instream):
        protids = dict()
        for line in instream:
            if '\tCDS\t' not in line:
                continue
            idmatch = re.search('protein_id=([^;\n]+)', line)
            assert idmatch, 'cannot parse protein_id: ' + line
            protid = idmatch.group(1)
            if protid not in protids:
                protids[protid] = True
                yield protid

    def protein_mapping(self, instream):
        locusid2name = dict()
        gene2loci = dict()
        mrna2gene = dict()
        proteins = dict()
        protdups = dict()
        for line in instream:
            fields = line.split('\t')
            if len(fields) != 9:
                continue
            feattype = fields[2]
            attrs = fields[8]

            if feattype == 'locus':
                idmatch = re.search('ID=([^;\n]+);.*Name=([^;\n]+)', attrs)
                if idmatch:
                    locusid = idmatch.group(1)
                    locusname = idmatch.group(2)
                    locusid2name[locusid] = locusname
            elif feattype == 'gene':
                idmatch = re.search('ID=([^;\n]+);Parent=([^;\n]+)', attrs)
                if idmatch:
                    geneid = idmatch.group(1)
                    ilocusid = idmatch.group(2)
                    gene2loci[geneid] = ilocusid
                else:  # pragma: no cover
                    print('Unable to parse gene and iLocus IDs: %s' % attrs,
                          file=sys.stderr)
            elif feattype == 'mRNA':
                idmatch = re.search('ID=([^;\n]+);Parent=([^;\n]+)', attrs)
                assert idmatch, \
                    'Unable to parse mRNA and gene IDs: %s' % attrs.rstrip()
                mrnaid = idmatch.group(1)
                geneid = idmatch.group(2)
                mrna2gene[mrnaid] = geneid
            elif feattype == 'CDS':
                if 'exception=rearrangement required for product' in attrs:
                    continue
                idmatch = re.search('Parent=([^;\n]+).*protein_id=([^;\n]+)',
                                    attrs)
                assert idmatch, \
                    'Unable to parse protein and mRNA IDs: %s' % attrs
                mrnaid = idmatch.group(1)
                proteinid = idmatch.group(2)
                if proteinid not in proteins:
                    geneid = mrna2gene[mrnaid]
                    proteins[proteinid] = mrnaid
                    ilocusid = gene2loci[geneid]
                    ilocusname = locusid2name[ilocusid]
                    yield proteinid, ilocusname


# -----------------------------------------------------------------------------
# Unit tests
# -----------------------------------------------------------------------------


def test_genome_download():
    """RefSeq chromosomes/scaffolds download"""
    config = genhub.test_registry.genome('Ador')
    testurl = ('ftp://ftp.ncbi.nlm.nih.gov/genomes/refseq/invertebrate/'
               'Apis_dorsata/all_assembly_versions/'
               'GCF_000469605.1_Apis_dorsata_1.3/'
               'GCF_000469605.1_Apis_dorsata_1.3_genomic.fna.gz')
    testpath = './Ador/GCF_000469605.1_Apis_dorsata_1.3_genomic.fna.gz'
    ador_db = RefSeqDB('Ador', config)
    assert '%r' % ador_db == 'RefSeq'
    assert ador_db.gdnaurl == testurl, \
        'scaffold URL mismatch\n%s\n%s' % (ador_db.gdnaurl, testurl)
    assert ador_db.gdnapath == testpath, \
        'scaffold path mismatch\n%s\n%s' % (ador_db.gdnapath, testpath)
    assert ador_db.compress_gdna is False

    config = genhub.test_registry.genome('Amel')
    testurl = ('ftp://ftp.ncbi.nlm.nih.gov/genomes/refseq/invertebrate/'
               'Apis_mellifera/all_assembly_versions/'
               'GCF_000002195.4_Amel_4.5/'
               'GCF_000002195.4_Amel_4.5_genomic.fna.gz')
    testpath = './Amel/GCF_000002195.4_Amel_4.5_genomic.fna.gz'
    amel_db = RefSeqDB('Amel', config)
    assert amel_db.gdnaurl == testurl, \
        'chromosome URL mismatch\n%s\n%s' % (amel_db.gdnaurl, testurl)
    assert amel_db.gdnapath == testpath, \
        'chromosome path mismatch\n%s\n%s' % (amel_db.gdnapath, testpath)
    assert ador_db.compress_gdna is False


def test_annot_download():
    """RefSeq annotation download"""
    config = genhub.test_registry.genome('Ador')
    testurl = ('ftp://ftp.ncbi.nlm.nih.gov/genomes/refseq/invertebrate/'
               'Apis_dorsata/all_assembly_versions/'
               'GCF_000469605.1_Apis_dorsata_1.3/'
               'GCF_000469605.1_Apis_dorsata_1.3_genomic.gff.gz')
    testpath = './Ador/GCF_000469605.1_Apis_dorsata_1.3_genomic.gff.gz'
    ador_db = RefSeqDB('Ador', config)
    assert ador_db.gff3url == testurl, \
        'annotation URL mismatch\n%s\n%s' % (ador_db.gff3url, testurl)
    assert ador_db.gff3path == testpath, \
        'annotation path mismatch\n%s\n%s' % (ador_db.gff3path, testpath)
    assert ador_db.compress_gff3 is False


def test_proteins_download():
    """RefSeq protein download"""
    config = genhub.test_registry.genome('Ador')
    testurl = ('ftp://ftp.ncbi.nlm.nih.gov/genomes/refseq/invertebrate/'
               'Apis_dorsata/all_assembly_versions/'
               'GCF_000469605.1_Apis_dorsata_1.3/'
               'GCF_000469605.1_Apis_dorsata_1.3_protein.faa.gz')
    testpath = ('/home/gandalf/HymHub/Ador/'
                'GCF_000469605.1_Apis_dorsata_1.3_protein.faa.gz')
    ador_db = RefSeqDB('Ador', config, workdir='/home/gandalf/HymHub')
    assert ador_db.proturl == testurl, \
        'protein URL mismatch\n%s\n%s' % (ador_db.proturl, testurl)
    assert ador_db.protpath == testpath, \
        'protein path mismatch\n%s\n%s' % (ador_db.protpath, testpath)
    assert ador_db.compress_prot is False


def test_gdna_format():
    """RefSeq gDNA formatting"""
    conf = genhub.test_registry.genome('Hsal')
    hsal_db = RefSeqDB('Hsal', conf, workdir='testdata/demo-workdir')
    hsal_db.preprocess_gdna(logstream=None, verify=False)
    outfile = 'testdata/demo-workdir/Hsal/Hsal.gdna.fa'
    testoutfile = 'testdata/fasta/hsal-first-7-out.fa'
    assert filecmp.cmp(testoutfile, outfile), 'Hsal gDNA formatting failed'

    conf = genhub.test_registry.genome('Tcas')
    tcas_db = RefSeqDB('Tcas', conf, workdir='testdata/demo-workdir')
    tcas_db.preprocess_gdna(logstream=None, verify=False)
    outfile = 'testdata/demo-workdir/Tcas/Tcas.gdna.fa'
    testoutfile = 'testdata/fasta/tcas-first-33-out.fa'
    assert filecmp.cmp(testoutfile, outfile), 'Tcas gDNA formatting failed'


def test_annot_format():
    """RefSeq annotation formatting"""
    conf = genhub.test_registry.genome('Aech')
    aech_db = RefSeqDB('Aech', conf, workdir='testdata/demo-workdir')
    aech_db.preprocess_gff3(logstream=None, verify=False)
    outfile = 'testdata/demo-workdir/Aech/Aech.gff3'
    testfile = 'testdata/gff3/ncbi-format-aech.gff3'
    assert filecmp.cmp(outfile, testfile), 'Aech annotation formatting failed'

    conf = genhub.test_registry.genome('Pbar')
    conf['annotfilter'] = 'NW_011933506.1'
    pbar_db = RefSeqDB('Pbar', conf, workdir='testdata/demo-workdir')
    pbar_db.preprocess_gff3(logstream=None, verify=False)
    outfile = 'testdata/demo-workdir/Pbar/Pbar.gff3'
    testfile = 'testdata/gff3/ncbi-format-pbar.gff3'
    assert filecmp.cmp(outfile, testfile), 'Pbar annotation formatting failed'

    conf = genhub.test_registry.genome('Ador')
    conf['annotfilter'] = ['NW_006264094.1', 'NW_006263516.1']
    ador_db = RefSeqDB('Ador', conf, workdir='testdata/demo-workdir')
    ador_db.preprocess_gff3(logstream=None, verify=False)
    outfile = 'testdata/demo-workdir/Ador/Ador.gff3'
    testfile = 'testdata/gff3/ncbi-format-ador.gff3'
    assert filecmp.cmp(outfile, testfile), 'Ador annotation formatting failed'


def test_prot_ncbi():
    """RefSeq protein formatting"""
    conf = genhub.test_registry.genome('Hsal')
    hsal_db = RefSeqDB('Hsal', conf, workdir='testdata/demo-workdir')
    hsal_db.preprocess_prot(logstream=None, verify=False)
    outfile = 'testdata/demo-workdir/Hsal/Hsal.all.prot.fa'
    testoutfile = 'testdata/fasta/hsal-13-prot-out.fa'
    assert filecmp.cmp(testoutfile, outfile), 'Hsal protein formatting failed'


def test_protids():
    """RefSeq: extract protein IDs from GFF3"""
    conf = genhub.test_registry.genome('Xtro')
    db = RefSeqDB('Xtro', conf)
    protids = ['XP_012809995.1', 'XP_012809996.1', 'XP_012809997.1',
               'XP_012809998.1']
    infile = 'testdata/gff3/xtro-3genes.gff3'
    testids = list()
    with open(infile, 'r') as instream:
        for protid in db.gff3_protids(instream):
            testids.append(protid)
    assert sorted(protids) == sorted(testids), \
        'protein ID mismatch: %r %r' % (protids, testids)


def test_protmap():
    """RefSeq: extract protein-->iLocus mapping from GFF3"""
    conf = genhub.test_registry.genome('Xtro')
    db = RefSeqDB('Xtro', conf)
    mapping = {'XP_012809997.1': 'XtroILC-43374',
               'XP_012809996.1': 'XtroILC-43374',
               'XP_012809995.1': 'XtroILC-43373',
               'XP_012809998.1': 'XtroILC-43374'}
    infile = 'testdata/gff3/xtro-3genes-loci.gff3'
    testmap = dict()
    with open(infile, 'r') as instream:
        for protid, locid in db.protein_mapping(instream):
            testmap[protid] = locid
    assert mapping == testmap, \
        'protein mapping mismatch: %r %r' % (mapping, testmap)


def test_cleanup():
    """RefSeq cleanup task"""
    config = genhub.test_registry.genome('Aech')
    db = RefSeqDB('Aech', config, workdir='testdata/demo-workdir')
    delfiles = ['testdata/demo-workdir/Aech/Aech.gff3']
    testfiles = db.cleanup(None, False, True)
    assert testfiles == delfiles, '%r %r' % (testfiles, delfiles)
    delfiles = ['testdata/demo-workdir/Aech/Aech.gff3',
                'testdata/demo-workdir/Aech/GCF_000204515.1_Aech_3.9_genomic.'
                'gff.gz']
    testfiles = db.cleanup(None, True, True)
    assert set(testfiles) == set(delfiles), '%r %r' % (testfiles, delfiles)

    config = genhub.test_registry.genome('Bdis')
    db = RefSeqDB('Bdis', config, workdir='testdata/demo-workdir')
    delfiles = ['testdata/demo-workdir/Bdis/Bdis.gdna.fa',
                'testdata/demo-workdir/Bdis/Bdis.gff3',
                'testdata/demo-workdir/Bdis/Bdis.ilocus.mrnas.gff3',
                'testdata/demo-workdir/Bdis/Bdis.ilocus.mrnas.txt',
                'testdata/demo-workdir/Bdis/Bdis.miloci.fa',
                'testdata/demo-workdir/Bdis/Bdis.miloci.gff3',
                'testdata/demo-workdir/Bdis/Bdis.mrnas.txt',
                'testdata/demo-workdir/Bdis/Bdis.simple-iloci.txt',
                'testdata/demo-workdir/Bdis/ilens.temp']
    testfiles = db.cleanup(None, False, True)
    assert set(testfiles) == set(delfiles), '%r %r' % (testfiles, delfiles)
    delfiles = ['testdata/demo-workdir/Bdis/Bdis.gdna.fa',
                'testdata/demo-workdir/Bdis/Bdis.gff3',
                'testdata/demo-workdir/Bdis/Bdis.ilocus.mrnas.gff3',
                'testdata/demo-workdir/Bdis/Bdis.ilocus.mrnas.txt',
                'testdata/demo-workdir/Bdis/Bdis.mrnas.txt',
                'testdata/demo-workdir/Bdis/Bdis.simple-iloci.txt',
                'testdata/demo-workdir/Bdis/ilens.temp']
    testfiles = db.cleanup(['.miloci.'], False, True)
    assert set(testfiles) == set(delfiles), '%r %r' % (testfiles, delfiles)
    delfiles = ['testdata/demo-workdir/Bdis/Bdis.gdna.fa',
                'testdata/demo-workdir/Bdis/Bdis.gff3',
                'testdata/demo-workdir/Bdis/Bdis.ilocus.mrnas.gff3',
                'testdata/demo-workdir/Bdis/Bdis.ilocus.mrnas.txt',
                'testdata/demo-workdir/Bdis/Bdis.mrnas.txt',
                'testdata/demo-workdir/Bdis/ilens.temp']
    testfiles = db.cleanup(['.miloci.', 'simple'], False, True)
    assert set(testfiles) == set(delfiles), '%r %r' % (testfiles, delfiles)

    config = genhub.test_registry.genome('Vcar')
    db = RefSeqDB('Vcar', config, workdir='testdata/demo-workdir')
    nodelfile = 'testdata/demo-workdir/Vcar/Vcar.protein2ilocus.txt'
    testfiles = db.cleanup(None, False, True)
    assert nodelfile not in testfiles, 'incorrectly deleted file'


def test_get_map():
    """RefSeq get prot map"""
    config = genhub.test_registry.genome('Vcar')
    db = RefSeqDB('Vcar', config, workdir='testdata/demo-workdir')
    mapping = dict()
    for protid, locid in db.get_prot_map():
        mapping[protid] = locid
    testmap = {'XP_002945607.1': 'VcarILC-00002',
               'XP_002945976.1': 'VcarILC-00003',
               'XP_002945608.1': 'VcarILC-00004',
               'XP_002945609.1': 'VcarILC-00005'}
    assert mapping == testmap, 'get prot map fail: %r %r' % (testmap, mapping)
