#!/usr/bin/env python3

from Bio.Blast.Applications import NcbiblastnCommandline
import csv
import os
from proceseq_16s.utilities import remove_leading_string
import pydoc
import tempfile


class Analyzer():
    '''Analyze sequences in fasta file by blast

    Parameters
    ----------
    fasta_file : str
        Path to the fasta file containing sequences to be analyzed (default: None)
    dbase : str
        Path to the database used by blast (including name of the database)
        (default: None)
    taxonomy_folder : str
        Path to the folder containing taxonomy files *names.dmp* and *nodes.dmp*
        downloaded from the NCBI web (default: None)
    taxonomy_ranks : tuple
        Taxonomy ranks returned by *element_to_tsv* method
        (default: ('kingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species'))
    bitscore_steps : int
        Number of the highest bitscore results parsed from the blast analysis results
        (default: 3)
    nproc : int
        Number of CPU threads used by blast (default: 8)

    Example
    -------
    ::

        from proceseq_16s import blast

        blast_analyzer = blast.analyzer('./sequences.fa', database='./database',
                                        taxonomy_folder='./taxonomy/')

        with open('taxonomized_seqs.tsv', 'w') as output_file:

            for seq_id in blast_analyzer.taxonomy:

                taxonomy_for_tsv = blast_analyzer.element_to_tsv(seq_id)

                output_file.write('Sequence ID: {}\\n'
                                  'Taxonomy: {}\\n'.format(seq_id, taxonomy_for_tsv))

    Attributes
    ----------
    bitscore_steps : int
        Number of the highest bitscore results parsed from the blast analysis results
    dbase : str
        Path to the database used by blast (including name of the database)
    fasta_file : str
        Path to the fasta file containing sequences to be analyzed
    nproc : int
        Number of CPU threads used by blast
    taxonomy_folder : str
        Path to the folder containing taxonomy files *names.dmp* and *nodes.dmp*
        downloaded from the NCBI web
    taxonomy_ranks : tuple
        Taxonomy ranks returned by *element_to_tsv* method
    '''

    _taxonomy_names = None
    _taxonomy_nodes = None

    @classmethod
    def _parse_taxonomy_nodes(cls, path_to_dir=None):
        '''Read taxonomy nodes from nodes.dmp file downloaded from NCBI web

        Parses nodes.dmp file and saves parsed result to *cls._taxonomy_nodes*
        class parameter. Data are parsed as:

        cls._taxonomy_nodes = {taxid:(parent_id, rank)}

        Parameters
        ----------
        path_to_dir : str
            Path to the directory containing nodes.dmp file
        '''
        cls._taxonomy_nodes = {}
        with open(os.path.join(path_to_dir, 'nodes.dmp'), 'r') as f:
            cls._z_taxonomy_nodes = {}
            for line in f.readlines():
                taxid, parent_id, rank, *_ = line.split('|')
                cls._taxonomy_nodes[int(taxid.strip())] = \
                    (int(parent_id.strip()), rank.strip())

    @classmethod
    def _parse_taxonomy_names(cls, path_to_dir=None):
        '''Read taxonomy names from names.dmp file downloaded from NCBI web

        Parses names.dmp file and saves parsed result to cls._taxonomy_names
        class parameter. Data are parsed as:

        cls._taxonomy_names = {taxid:{class_name:name}}

        Parameters
        ----------
        path_to_dir : str
            Path to the directory containing names.dmp file
        '''
        names_file = os.path.join(path_to_dir, 'names.dmp')
        with open(names_file, 'r') as f:
            cls._taxonomy_names = {}
            for line in f.readlines():
                taxid, name, _, class_name, _ = line.split('|')
                try:
                    cls._taxonomy_names[int(taxid.strip())][class_name.strip()] = \
                        name.strip()
                except KeyError:
                    cls._taxonomy_names[int(taxid.strip())] = \
                        {class_name.strip(): name.strip()}

    def __init__(self, fasta_file=None, dbase=None, taxonomy_folder=None,
                 taxonomy_ranks=None, bitscore_steps=3, nproc=8):
        self.dbase = dbase
        self.fasta_file = fasta_file
        self.taxonomy_folder = taxonomy_folder
        self.bitscore_steps = bitscore_steps
        self.nproc = nproc
        if taxonomy_ranks is None:
            self.taxonomy_ranks = ('kingdom', 'phylum', 'class', 'order', 'family',
                                   'genus', 'species')
        else:
            self.taxonomy_ranks = taxonomy_ranks
        self._columns = ('qseqid staxid evalue mismatch gapopen pident length qstart '
                         'qend sstart send bitscore sseqid').split()
        self._taxonomy = None

    def _taxonomize(self):
        '''Analyze input fasta file by blast and parse results

        Parameters
        ----------
        None

        Returns
        -------
        dict
            Parsed alignment and taxonomy results. Results are returned as a
            dictionary mapping sequence identifiers as keys to taxonomy information.
            Taxonomy information is presented as a list of all blast hits. Every hit is
            saved as another dictionary mapping alignement and taxonomy parameters
            to results of blast analysis:
            {seq_id:[{'taxonomy': {rank: taxon}, 'length': length, ...},
                     {'taxonomy': {rank: taxon}, 'length': length, ...}
                    ]
            }
        '''
        if self.dbase is None:
            raise ValueError('No blast database provided!')
        if self.taxonomy_folder is None:
            raise ValueError('No folder containing taxonomy data provided!')

        with tempfile.NamedTemporaryFile(suffix='-blast.csv',
                                         mode='w',
                                         delete=False) as tmp_file:
            blast_command = NcbiblastnCommandline(
                query=self.fasta_file,
                task='megablast',
                num_threads=self.nproc,
                db=self.dbase,
                out=tmp_file.name,
                max_hsps=1,
                outfmt=('"6 {}"'.format(' '.join(self._columns))),
            )

            blast_command()

        self._taxonomy_file = tmp_file.name
        return self.parse_results()

    def _add_blast_taxonomy(self, taxon_id, taxonomy_folder=None):
        '''Create blast taxonomy information from dump files

        Parse names.dmp and nodes.dmp files from the local NCBI database
        and create dictionary containing taxonomic levels starting from
        provided blast taxonomy ID up to the kingdom.

        Parameters
        ----------
        taxon_id : int
            Blast taxonomy ID (staxids)
        taxonomy_folder : str
            Path to the folder containing blast taxonomic tree files
            (names.dmp and nodes.dmp). *self.taxonomy_folder* attribute
            is used, if this parameter is omitted.

        Returns
        -------
        dict
            Dictonary with taxonomic ranks as keys
            and taxon names as values
        '''
        if taxonomy_folder is None:
            taxonomy_folder = self.taxonomy_folder
        if self._taxonomy_nodes is None or self._taxonomy_names is None:
            self._parse_taxonomy_names(taxonomy_folder)
            self._parse_taxonomy_nodes(taxonomy_folder)

        taxonomy = {}

        while taxon_id != 1:
            rank = self._taxonomy_nodes[taxon_id][1]
            if rank == 'superkingdom':
                rank = 'kingdom'
            taxonomy[rank] = self._taxonomy_names[taxon_id]['scientific name']
            taxon_id = self._taxonomy_nodes[taxon_id][0]

        return taxonomy

    @property
    def taxonomy(self):
        '''Taxonomized sequences from input fasta file

        Parsed alignment and taxonomy results. Results are returned as a
        dictionary mapping sequence identifiers as keys to taxonomy information.
        Taxonomy information is presented as a list of all blast hits. Every hit is
        saved as another dictionary mapping alignement and taxonomy parameters
        to results of blast analysis:
        {seq_id:[{'taxonomy': {rank: taxon}, 'length': length, ...},
        {'taxonomy': {rank: taxon}, 'length': length, ...}]}

        :type: dict
        :getter: Returns taxonomized sequences
        '''
        self._taxonomize()
        return self._taxonomy

    def parse_results(self, taxonomy_file=None):
        '''Parse results of the blast analysis

        Parameters
        ----------
        taxonomy_file : str
            Path to a file with results of blast analysis (required)

        Returns
        -------
        dict
            Parsed alignment and taxonomy results. Results are returned as a
            dictionary mapping sequence identifiers as keys to taxonomy information.
            Taxonomy information is presented as a list of all blast hits. Every hit is
            saved as another dictionary mapping alignement and taxonomy parameters
            to results of blast analysis:
            {seq_id:[{'taxonomy': {rank: taxon}, 'length': length, ...},
                {'taxonomy': {rank: taxon}, 'length': length, ...}]}
        '''
        if taxonomy_file is None:
            try:
                taxonomy_file = self._taxonomy_file
            except AttributeError:
                raise AttributeError('Missing file with blast results!')

        last_sequence = None

        with open(taxonomy_file, 'r') as csv_file:
            csv_reader = csv.reader(csv_file, delimiter='\t')
            self._taxonomy = {}

            for row in csv_reader:
                if len(row) == 0:
                    continue
                sequence = row[0]
                results = dict(zip(self._columns[1:], row[1:]))
                if sequence != last_sequence:
                    last_bitscore = 10000
                    bitscore_rank = 0
                    self._taxonomy[sequence] = []
                    taxids = set()
                    last_sequence = sequence
                if float(results['bitscore']) < last_bitscore:
                    last_bitscore = float(results['bitscore'])
                    bitscore_rank += 1
                if (bitscore_rank > self.bitscore_steps or
                        results['staxid'] in taxids):
                    continue
                taxids.add(results['staxid'])
                taxonomy = self._add_blast_taxonomy(int(results['staxid']))
                taxonomy['species'] = remove_leading_string(
                    taxonomy['species'], taxonomy.get('genus', '') + ' ')
                results['taxonomy'] = taxonomy
                self._taxonomy[sequence].append(results)

        return self._taxonomy

    def element_to_tsv(self, seq_id):
        '''Format results from the last analyzed taxonomy for output to tsv file

        Blast alignment and taxonomization will be performed, if it hasn't been
        done manually.

        Parameters
        ----------
        seq_id : str
            Identifier (key in *self.taxonomy* attribute) of the sequence,
            whose results will be formated and returned (required)

        Returns
        -------
        str
            Result of the analysis formated for saving to the output tsv file
        '''
        if self._taxonomy is None:
            self._taxonomize()

        taxonomy_for_tsv = ''
        try:
            extracted_taxonomies = self._taxonomy[seq_id]
            for blast_result in extracted_taxonomies:
                final_taxonomy = [blast_result['taxonomy'].get(taxonomy_rank, 'NA')
                                  for taxonomy_rank in self.taxonomy_ranks]
                taxonomy_for_tsv += ('\tBlast-bits:{}\t{}\tmismatch:{}\tgaps:{}\te_val:{}'
                                     '\n').format(blast_result['bitscore'],
                                                  '\t'.join(final_taxonomy),
                                                  blast_result['mismatch'],
                                                  blast_result['gapopen'],
                                                  blast_result['evalue'])

        except KeyError:
            pass
            # print('Blast - sequence ID {} not found.'.format(seq_id))
            # taxonomy_for_tsv = ''

        return taxonomy_for_tsv


if __name__ == '__main__':
    exit(pydoc.render_doc(__name__))  # pragma: no cover
