#!/usr/bin/env python
"""
COG_phylum_table.py

Count single copy genes (scgs) for each phylum per cluster of contigs.

The output table will contain three columns for each phylum x.
The column "x total" will contain the number of COGS that are considered
scgs for this phylum and will thus be the same for each row.
The column "x >1" will contain the number of COGS from these phylum
specific scgs that are present in more than 1 copy for the cluster
corresponding to that row.
The column "x ==1" will contain the number of COGS from these phylum
specific scgs that are present in exactly 1 copy for the cluster 
corresponding to that row.
"""
import sys
import argparse
from cogler.io import read_blast_output, read_gff_file, read_clustering_file
import cogler
from collections import defaultdict
import pandas as pd
import logging

def get_records_from_file(queries, cog_file):
    # Read a simple tsv file with two columns, cddid and cogid respectively.
    with open(cog_file, 'r') as cf:
        records = dict([(row.split('\t')[0], {'Accession': row.split('\t')[1].strip()}) for row in cf.readlines()])
    return records

def main(args):
    RPSBLAST_SCOVS_THRESHOLD = args.scovs_threshold
    RPSBLAST_PIDENT_THRESHOLD = args.pident_threshold

    records, sseq_ids = read_blast_output(args.blastoutfile)

    # Retrieve the cog accession number from file
    cogrecords = get_records_from_file(sseq_ids, args.cdd_cog_file)

    features_per_contig = defaultdict(list)

    # For every hit in blast output, determine if the hit is valid,
    # and add it to a list for that contig
    for record_d in records:
        pident_above_threshold = record_d['pident'] >= RPSBLAST_PIDENT_THRESHOLD

        # A certain fraction of the cog should be covered to avoid the same cog 
        # to be counted twice in the case when a cog is split across two or more contigs.
        alignment_length_in_subject = abs(record_d['send'] - record_d['sstart']) + 1
        percent_seq_covered = (alignment_length_in_subject / record_d['slen']) * 100.0
        seq_cov_above_threshold =  percent_seq_covered >= RPSBLAST_SCOVS_THRESHOLD
        
        if pident_above_threshold and seq_cov_above_threshold:
            cog_accession = cogrecords[record_d['sseqid'].split('|')[2]]['Accession']
            contig = (args.separator).join(record_d['qseqid'].split(args.separator)[:-1])
            features_per_contig[contig].append(cog_accession)

    # Load clustering
    cluster_per_contig = read_clustering_file(args.cluster_file)

    cogs_per_contig = cogler.CogsPerContig(cluster_per_contig, features_per_contig)

    o = cogler.Output(cogs_per_contig)

    phyla_df = pd.read_table(args.marker_file, index_col=0)
    for phylum_name, row in phyla_df.iterrows():
        phylum_name = phylum_name.strip()
        # scgs will be the column names for each column in
        # COG0001-end with a nonzero value for this phylum
        scgs = list(row.loc['COG0001':][row > 0].index.values)
        if not scgs:
            logging.warning("Phylum {0} had no scgs, skipping".format(phylum_name))
            continue
        if row['Number_genera'] < args.min_genera:
            logging.warning("Phylum {0} had too few genera included, skipping".format(phylum_name))
            continue
        phylum = cogler.Phylum(phylum_name, scgs)
        o.add_phylum(phylum)
    o.result_matrix.to_csv(args.output_file, index_label='cluster')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(usage=__doc__)
    parser.add_argument('-b', '--blastoutfile', required=True,
            help=('REQUIRED: Output of rpsblast run against the COG database, '
               'assumed to be in tabular format whith '
               'columns: qseqid sseqid evalue pident score qstart qend sstart send length slen. '
               'The contigs ids are assumed to be recoverable by removing the last underscore '
               'and the characters following it from the qseqid column.' ))
    parser.add_argument('-g', '--gfffile',
           help=('GFF file generated by e.g. prodigal '
           'only needed if the contig names are not recoverable from the '
           'blast output file.'))
    parser.add_argument('-c', '--cluster_file', required=True,
            help=('REQUIRED: Clustering file from concoct execution.'))
    parser.add_argument('-m', '--marker_file', required=True,
            help=('REQUIRED: File containing a table of genes that will be used as marker genes'
               'per phylum. The table should have phylum as the first column, '
               'and have one column named Number_genera. The columns for the '
               'COG ids should be contained in one chunk, without any other '
               'columns inbetween, ranging from COG00001 to COG5665.'))
    parser.add_argument('-s', '--scovs-threshold', type=float, default=50.0,
           help='Threshold covered in percent, default=50.0')
    parser.add_argument('-p', '--pident-threshold', type=float, default=0.0,
           help='Threshold identity in percent, default=0.0')
    parser.add_argument('--cdd_cog_file', required=True,
            help = ('REQUIRED: Supply a cdd to cog mapping file in a tsv format.'))
    parser.add_argument('--min_genera', default=0, type=int,
            help=('REQUIRED: The minimum number of generas for a phylum needed to be '
                 'included in the analysis'))
    parser.add_argument('--separator', default="_",
           help=('Character that is used to separate the contig id from the '
                 'protein identifier. Everything before the last occurence ' 
                 'of this character will be used as the contig id. Default '
                 'value is "_"'))
    parser.add_argument('--output_file', required=True,
            help=("REQUIRED: The file name where the output table will be printed."))
    args = parser.parse_args()

    if args.gfffile:
        raise NotImplementedError

    main(args)
