#!/usr/bin/env python
#
# -----------------------------------------------------------------------------
# Copyright (c) 2016   Daniel Standage <daniel.standage@gmail.com>
# Copyright (c) 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 The Arabidopsis Information Resource.

GenomeDB implementation for data residing in TAIR. Currently tested only on
TAIR6, but presumably trivial to extend to other versions using Att6.yml as a
template.
"""

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


class TairDB(genhub.genomedb.GenomeDB):

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

        assert self.config['source'] == 'tair'
        assert 'version' in self.config
        assert str(self.config['version']) == '6'
        assert 'species' in self.config
        assert self.config['species'] == 'Arabidopsis thaliana'

        self.specbase = 'ftp://ftp.arabidopsis.org/home/tair'
        self.format_gdna = self.format_fasta
        self.format_prot = self.format_fasta

    def __repr__(self):
        return 'TAIR' + self.version

    @property
    def version(self):
        return str(self.config['version'])

    @property
    def gdnafilename(self):
        return self.config['gdna_filename']

    @property
    def gff3filename(self):
        return self.config['gff3_filename']

    @property
    def protfilename(self):
        return self.config['prot_filename']

    @property
    def gdnaurl(self):
        return self.config['gdna_url']

    @property
    def gff3url(self):
        return self.config['gff3_url']

    @property
    def proturl(self):
        return self.config['prot_url']

    def format_fasta(self, instream, outstream, logstream=sys.stderr):
        for line in instream:
            # No processing required.
            # 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:  # pragma: no cover
            excludefile = self.filter_file()
            cmds.append('grep -vf %s' % excludefile.name)
        cmds.append('sed "s/Index=/index=/"')
        cmds.append('tidygff3')
        cmds.append('genhub-format-gff3.py --source tair -')
        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:  # pragma: no cover
            os.unlink(excludefile.name)

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

    def protein_mapping(self, instream):
        locusid2name = dict()
        gene2loci = 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('Parent=([^;\n]+);Name=([^;\n]+)', attrs)
                assert idmatch, \
                    'Unable to parse mRNA and gene IDs: %s' % attrs.rstrip()
                geneid = idmatch.group(1)
                mrnaid = idmatch.group(2)
                ilocusid = gene2loci[geneid]
                ilocusname = locusid2name[ilocusid]
                yield mrnaid, ilocusname


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


def test_gdna_format():
    """TAIR6 gDNA formatting"""
    conf = genhub.test_registry.genome('Att6')
    tair_db = TairDB('Att6', conf, workdir='testdata/demo-workdir')
    tair_db.preprocess_gdna(logstream=None, verify=False)
    outfile = 'testdata/demo-workdir/Att6/Att6.gdna.fa'
    testoutfile = 'testdata/fasta/tair6-gdna-out.fa'
    assert filecmp.cmp(testoutfile, outfile), 'TAIR6 gDNA formatting failed'

    assert repr(tair_db) == 'TAIR6'
    ftpbase = ('ftp://ftp.arabidopsis.org/home/tair/Sequences/'
               'whole_chromosomes/OLD/chr%d.fas')
    urls = [ftpbase % x for x in (1, 2, 3, 4, 5)]
    assert tair_db.gdnaurl == urls, tair_db.gdnaurl
    url = ('ftp://ftp.arabidopsis.org/home/tair/Genes/TAIR6_genome_release/'
           'TAIR6_GFF3_genes.gff')
    assert tair_db.gff3url == url, tair_db.gff3url
    url = ('ftp://ftp.arabidopsis.org/home/tair/Genes/TAIR6_genome_release/'
           'TAIR6_pep_20060907')
    assert tair_db.proturl == url, tair_db.proturl


def test_annot_format():
    """TAIR6 annotation formatting"""
    conf = genhub.test_registry.genome('Att6')
    db = TairDB('Att6', conf, workdir='testdata/demo-workdir')
    db.preprocess_gff3(logstream=None, verify=False)
    outfile = 'testdata/demo-workdir/Att6/Att6.gff3'
    testfile = 'testdata/gff3/tair6-format.gff3'
    assert filecmp.cmp(outfile, testfile), 'TAIR6 annotation formatting failed'


def test_protids():
    """TAIR6: extract protein IDs from GFF3"""
    conf = genhub.test_registry.genome('Att6')
    db = TairDB('Att6', conf)
    protids = ['AT5G01010.1', 'AT5G01015.1', 'AT5G01020.1', 'AT5G01030.1',
               'AT5G01030.2', 'AT5G01040.1', 'AT5G01050.1', 'AT5G01060.1',
               'AT5G01070.1', 'AT5G01075.1']
    infile = 'testdata/gff3/tair6-format.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():
    """TAIR6: extract protein-->iLocus mapping from GFF3"""
    conf = genhub.test_registry.genome('Att6')
    db = TairDB('Xtro', conf)
    mapping = {'AT2G01800.1': 'Att6ILC-09797',
               'AT2G01810.1': 'Att6ILC-09799',
               'AT2G01820.1': 'Att6ILC-09801',
               'AT2G01830.1': 'Att6ILC-09803',
               'AT2G01830.2': 'Att6ILC-09803',
               'AT2G01830.3': 'Att6ILC-09803'}
    infile = 'testdata/gff3/tair6-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)
