###############################################################################
#                                                                             #
#    This program is free software: you can redistribute it and/or modify     #
#    it under the terms of the GNU General Public License as published by     #
#    the Free Software Foundation, either version 3 of the License, or        #
#    (at your option) any later version.                                      #
#                                                                             #
#    This program is distributed in the hope that it will be useful,          #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of           #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the            #
#    GNU General Public License for more details.                             #
#                                                                             #
#    You should have received a copy of the GNU General Public License        #
#    along with this program. If not, see <http://www.gnu.org/licenses/>.     #
#                                                                             #
###############################################################################

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import logging
from collections import defaultdict

from checkm.defaultValues import DefaultValues
from checkm.markerSets import MarkerSetParser, BinMarkerSets

from unitem.defaults import *


class Markers():
    """Read marker information."""

    def __init__(self):
        """Initialization."""
        
        self.logger = logging.getLogger('timestamp')
 
        parser = MarkerSetParser()
        bin_marker_sets = parser.parseTaxonomicMarkerSetFile(CHECKM_BAC_MS)
        self.bac_ms = bin_marker_sets.mostSpecificMarkerSet()
        
        bin_marker_sets = parser.parseTaxonomicMarkerSetFile(CHECKM_AR_MS)
        self.ar_ms = bin_marker_sets.mostSpecificMarkerSet()
        
        self.bac_markers_on_contigs = None
        self.ar_markers_on_contigs = None
        
    def read_table(self, marker_gene_table):
        """Read table indicating marker genes for each bin."""
        
        markers = defaultdict(lambda: defaultdict(list))
        if not os.path.exists(marker_gene_table):
            # did not identify any markers
            return markers

        with open(marker_gene_table) as f:
            f.readline()
            
            for line in f:
                line_split = line.strip().split('\t')
                
                bin_id = line_split[0]
                marker_id = line_split[1]
                gene_id = line_split[2]
                if '&&' in gene_id:
                    # indicates a marker gene is identified 
                    # in adjacent genes and is likely an
                    # assembly or gene calling error
                    gene_ids = gene_id.split('&&')
                    gene_id = gene_ids[0]
                scaffold_id = gene_id[0:gene_id.rfind('_')]
                
                markers[bin_id][marker_id].append(scaffold_id)
                
        return markers
        
    def evaluate_bac(self, gene_table, individual_markers=False):
        """Evaluate completeness and contamination of genome using bacterial marker sets."""
        
        comp, cont = self.bac_ms.genomeCheck(gene_table, individual_markers)
        return comp, cont
        
    def evaluate_ar(self, gene_table, individual_markers=False):
        """Evaluate completeness and contamination of genome using archaeal marker sets."""
        
        comp, cont = self.ar_ms.genomeCheck(gene_table, individual_markers)
        return comp, cont
        
    def evaluate(self, bac_gene_table, ar_gene_table):
        """Evaluate completeness and contamination of genome using best domain-level marker sets."""
        
        bac_comp, bac_cont = self.evaluate_bac(bac_gene_table, True)
        ar_comp, ar_cont = self.evaluate_ar(ar_gene_table, True)
        
        if bac_comp + bac_cont > ar_comp + ar_cont:
            # select domain set with the larget number of identified markers
            # including those present multiple times
            bac_comp, bac_cont = self.evaluate_bac(bac_gene_table, False)
            return 'Bacteria', bac_comp, bac_cont

        ar_comp, ar_cont = self.evaluate_ar(ar_gene_table, False)
        return 'Archaea', ar_comp, ar_cont
        
    def bin_quality(self, bin):
        """Estimate quality of bin."""
        
        # create gene tables for bin
        bac_gene_table = defaultdict(list)
        ar_gene_table = defaultdict(list)
        for cid in bin:
            for marker_id in self.bac_markers_on_contigs[cid]:
                bac_gene_table[marker_id].append(cid)
                
            for marker_id in self.ar_markers_on_contigs[cid]:
                ar_gene_table[marker_id].append(cid)
                
        domain, comp, cont = self.evaluate(bac_gene_table, 
                                                ar_gene_table)
                                                
        return domain, comp, cont
        
    def marker_gene_tables(self, profile_dir, binning_methods=None):
        """Get marker genes for bins across all binning methods."""

        markers = Markers()
        binning_methods_dir = os.path.join(profile_dir, BINNING_METHOD_DIR)
        gene_tables = {}
        for bm in os.listdir(binning_methods_dir):
            if binning_methods and bm not in binning_methods:
                continue # not is set of desired binning methods
                
            bac_mg_table = os.path.join(binning_methods_dir, bm, CHECKM_BAC_DIR, MARKER_GENE_TABLE)
            bac_gene_tables = markers.read_table(bac_mg_table)

            ar_mg_table = os.path.join(binning_methods_dir, bm, CHECKM_AR_DIR, MARKER_GENE_TABLE)
            ar_gene_tables = markers.read_table(ar_mg_table)
            
            # fill in any bins with no marker genes
            bin_dir = os.path.join(binning_methods_dir, bm, CHECKM_BAC_DIR, 'bins')
            for bin_id in os.listdir(bin_dir):
                if not os.path.isdir(os.path.join(bin_dir, bin_id)):
                    continue
            
                if bin_id not in bac_gene_tables:
                    bac_gene_tables[bin_id] = {}

                if bin_id not in ar_gene_tables:
                    ar_gene_tables[bin_id] = {}

            gene_tables[bm] = (bac_gene_tables, ar_gene_tables)
            
        self._markers_on_contigs(gene_tables)
            
        return gene_tables
        
    def _markers_on_contigs(self, gene_tables):
        """Get markers on each contig."""
        
        self.bac_markers_on_contigs = defaultdict(list)
        self.ar_markers_on_contigs = defaultdict(list)
        processed_contigs = set()
        for binning_method, (bac_gene_tables, ar_gene_tables) in gene_tables.items():
            scaffolds_in_binning_method = set()
            for bin_id in bac_gene_tables:
                for marker_id, scaffold_ids in bac_gene_tables[bin_id].items():
                    for scaffold_id in scaffold_ids:
                        if scaffold_id in processed_contigs:
                            continue
                        self.bac_markers_on_contigs[scaffold_id].append(marker_id)
                        
                    scaffolds_in_binning_method.update(scaffold_ids)
                            
            for bin_id in ar_gene_tables:
                for marker_id, scaffold_ids in ar_gene_tables[bin_id].items():
                    for scaffold_id in scaffold_ids:
                        if scaffold_id in processed_contigs:
                            continue
                        self.ar_markers_on_contigs[scaffold_id].append(marker_id)
                        
                    scaffolds_in_binning_method.update(scaffold_ids)
                        
            processed_contigs.update(scaffolds_in_binning_method)
            
    def create_gene_table(self, contig_ids):
        """Create gene tables for a set of contigs."""
        
        bac_table = defaultdict(list)
        ar_table = defaultdict(list)
        for cid in contig_ids:
            for mid in self.bac_markers_on_contigs[cid]:
                bac_table[mid].append(cid)
                
            for mid in self.ar_markers_on_contigs[cid]:
                ar_table[mid].append(cid)
                
        return bac_table, ar_table
        