
from __future__ import division, print_function

import json

from csv import reader as csv_reader
from operator import itemgetter
from os.path import abspath, basename, dirname, exists, join, splitext
from warnings import warn

from BioExt import AmbigList, translate_ambiguous


__all__ = []


_scoredata = {
    'PI': {'name': 'PI.dat', 'offset': 0, 'gene': 'PR'},
    'NRTI': {'name': 'NRTI.dat', 'offset' : 99, 'gene': 'RT'},
    'NNRTI': {'name': 'NNRTI.dat', 'offset' : 99, 'gene': 'RT'}
}

COMPLETE_LIST, BINARY  = range(2)


class DrugCoord(object):
    DELETION, INSERTION, VALUE = range(3)

    def __init__(self, coord, klass, amino, value):
        if klass not in (DrugCoord.DELETION, DrugCoord.INSERTION, DrugCoord.VALUE):
            raise ValueError('DrugCoord klass must be one of DrugCoord.{DELETION, INSERTION, or VALUE}')
        if not isinstance(coord, int):
            raise ValueError('DrugCoord coord must be an index into a sequence')

        self._coord = coord
        self._klass = klass
        self._amino = amino.upper()
        self._value = value

    def __call__(self, seq):
        if not isinstance(seq, AmbigList):
            raise ValueError('seq must be an AmbigList, generated by translate_ambiguous in BioExt')

        c = self._coord
        if c >= len(seq):
            warn("sequence too short to evaluate at position '%d'" % self._coord)
            return c, 0, 0

        ndegen = len(seq[c])
        r = 0
        if self._klass == DrugCoord.DELETION:
            if seq[c] == '-':
                r = self._value
        elif self._klass == DrugCoord.INSERTION:
            warn('INSERTION UNEVALUATED!!! We need to think about how to communicate insertion events')
        else:
            if self._amino in seq[c]:
                r = self._value
        return c, r, ndegen


class DrugData(object):

    def __init__(self, name, gene, drug_coords):
        self._name = name
        self._gene = gene
        self._drug_coords = drug_coords

    def __call__(self, seq):
        # sort by the coordinate
        tmp = sorted((d(seq) for d in self._drug_coords), key=itemgetter(0))
        # group by the coordinate
        r = []
        prev = -1
        for coord, s, n in tmp:
            # if we're the previous coord,
            # append to that list
            if len(r) > 0 and coord == prev:
                r[-1].append((coord, s, n))
            # otherwise, make a new list
            else:
                r.append([(coord, s, n)])
            prev = coord
        return r


def read_scorefile(filepath, gene, offset):
    if not exists(filepath):
        raise ValueError("file '%s' doesn't exist!" % filepath)

    with open(filepath) as fh:
        reader = csv_reader(fh, delimiter='\t')

        header = reader.next()

        # some error correction
        if len(header) < 3:
            raise RuntimeError('expected at least 3 columns')
        if header[0].lower() != 'position':
            raise RuntimeError("expected first column label to be 'position'")
        if header[1].lower() != 'aa':
            raise RuntimeError("expected second column label to be 'aa'")

        drug_idxs = [(i + 2, d) for i, d in enumerate(header[2:])]
        pos_by_drug = dict((d, []) for d in header[2:])

        for row in reader:
            # coordinates are assumed 1-indexed
            coord = int(row[0]) + offset - 1
            amino = row[1].upper()
            klass = (DrugCoord.DELETION  if amino == '-' else
                     DrugCoord.INSERTION if amino == 'INSERT' else
                     DrugCoord.VALUE)
            for idx, drug in drug_idxs:
                val = int(row[idx])
                if val > 0:
                    pos_by_drug[drug].append(DrugCoord(coord, klass, amino, val))

        return dict((d, DrugData(d, gene, pos_by_drug[d])) for _, d in drug_idxs)


def load_scoredata():
    scorefiledir = join(dirname(abspath(__file__)), 'data', 'scores')
    # scorefiles = iglob(join(scorefiledir, '*.dat'))
    # return a dict of dicts, key = drug class, value = dict of DrugData's,
    # keyed on the drug name
    r = {}
    for klass, data in _scoredata.items():
        fp = join(scorefiledir, data['name'])
        gene = data['gene']
        off = data['offset']
        r[klass] = read_scorefile(fp, gene, off)
    return r


OPTIMISTIC, PESSIMISTIC = range(2)

def DRAM_calculator_p(msa, mode=None, noise_threshold=4):
    if mode is None:
        mode = OPTIMISTIC

    if mode not in (OPTIMISTIC, PESSIMISTIC):
        raise ValueError('mode must be one of either OPTIMISTIC or PESSIMISTIC')

    table = load_scoredata()

    resolver = max if mode == OPTIMISTIC else min

    r = []
    for seq in msa:
        ambiglist = translate_ambiguous(seq)
        seq_r = {}
        for klass, drugdatas in table.items():
            klass_r = {}
            for drug, drugdata in drugdatas.items():
                # grab a list of lists,
                # where each sublist is for a particular coord
                coord_groups = drugdata(ambiglist)
                # now grab either the min or max from these sublists,
                # depending on how optimistic or pessimistic we are being
                # and use a generator as we can calculate in a single pass
                coord_score_ndegens = (resolver(g, key=itemgetter(1)) for g in coord_groups)
                # grab some stats, in a single pass
                covered, missing, noise, score = 0, 0, 0, 0
                for _, s, n in coord_score_ndegens:
                    if n > noise_threshold:
                        noise += 1
                    elif s > 0:
                        covered += 1
                        score += s
                    else: # s == 0 => missing
                        missing += 1
                klass_r[drug] = (covered, missing, noise, score)
            seq_r[klass] = klass_r
        r.append(seq_r)
    return r


# DRAM_AMBIGS_PESSIMISTIC, DRAM_AMBIGS_OPTIMISTIC = range (2)
# DRAM_REPORT_COMPLETE, DRAM_REPORT_BY_DRUG, DRAM_REPORT_BY_MUTATION, DRAM_REPORT_BY_CLASS = range (4)
# 
# def DRAM_reporter (positional_resistance_data, resistance_table = None, ambiguity_resolution = None, report_mode = None, mut_score_limit = 30, drug_resistance_limit = 30):
# # positional_resistance_data is returned by DRAM_calculator
#     if resistance_table is None:
#         if '>resistance table' in positional_resistance_data:
#             resistance_table = positional_resistance_data['>resistance table']
#         else:
#             resistance_table = import_mdr_scores_from_stanford ()
#     if ambiguity_resolution is None:
#         ambiguity_resolution = DRAM_AMBIGS_PESSIMISTIC
#     if report_mode is None:
#         report_mode = DRAM_REPORT_COMPLETE
# 
#     dram_report = {}
# 
#     # generate the full report by default, then condense if necessary
#     # score by drug (with comments if sequenced completely or has noisy data): min and max
#     # resistance mutations by drug class (can be limited to major)
#     # resistance to classes e.g. NRTI : ['susceptible', 'resistant']
# 
# 
#     for sequence_id in positional_resistance_data:
#         if sequence_id == '>resistance table':
#             continue
#         sequence_info = positional_resistance_data[sequence_id]
#         sequence_report = {'by drug' : {}}
#         for drug_name in resistance_table['by drug']:
#             total_score = 0
#             covered_positions = 0
#             noisy_positions   = 0
#             missing_positions = 0
#             major_DRAM        = []
# 
#             list_of_positions_by_drug = resistance_table['by drug'][drug_name]
# 
#             for position in list_of_positions_by_drug:
#                 score = 0
#                 in_this_sequence = sequence_info[str(position)]
#                 resolution_scores = {}
# 
#                 if in_this_sequence is None:
#                     missing_positions += 1
#                     continue
#                 elif in_this_sequence[0] == 'Noisy':
#                     noisy_positions += 1
#                 else:
#                     covered_positions += 1
#                     for value in in_this_sequence:
#                         if value != 'WT':
#                             try:
#                                 score = resistance_table['mutations'][value][drug_name]
#                             except KeyError:
#                                 score = 0
#                         else:
#                             score = 0
#                         resolution_scores [value] = score
#                     mutation,score  = _DRAM_resolve_score (resolution_scores, ambiguity_resolution)
#                     if score >= mut_score_limit:
#                         major_DRAM.append(resistance_table["mutations"][mutation]["Original"])
# 
#                 total_score += score
# 
#             sequence_report ['by drug'][drug_name] = {'covered' : covered_positions, 'missing' : missing_positions, 'noisy': noisy_positions, 'score':  total_score ,'mutations' : major_DRAM if len (major_DRAM) else None}
# 
# 
#         sequence_report ['by class'] = {}
#         for drug_class, class_info in resistance_table['by class'].items():
#             report = {'mutations' : set(), 'resistant' : []}
#             for drug_id in class_info['names']:
#                 #print (drug_id, report)
#                 if sequence_report ['by drug'][drug_id]['score'] >= drug_resistance_limit:
#                     report['resistant'].append (drug_id)
#                 if sequence_report ['by drug'][drug_id]['mutations'] is not None:
#                     report['mutations'].update (sequence_report ['by drug'][drug_id]['mutations'])
#             report['mutations'] = list (report['mutations'])
#             sequence_report ['by class'][drug_class] = report
# 
#         dram_report [sequence_id] = sequence_report
#         #assert False
# 
#     if report_mode == DRAM_REPORT_BY_DRUG:
#     # in this mode the report is of the form
#     # seq id => list of drugs to which the sequence is resistant, grouped by class
#         drug_report = {}
#         for seq_id in dram_report:
#             seq_drug_report = {}
#             for drug_class, drug_value in dram_report[seq_id]['by class'].items():
#                 seq_drug_report [drug_class] = drug_value['resistant']
#             drug_report [seq_id] = seq_drug_report
#         return drug_report
#     elif report_mode == DRAM_REPORT_BY_MUTATION:
#     # in this mode the report is of the form
#     # seq id => list of major mutations
#         mut_report = {}
#         for seq_id in dram_report:
#             seq_mut_report = set()
#             for drug_class, drug_value in dram_report[seq_id]['by class'].items():
#                 seq_mut_report.update(drug_value['mutations'])
#             mut_report [seq_id] = list(seq_mut_report)
#         return mut_report
#     elif report_mode == DRAM_REPORT_BY_CLASS:
#     # in this mode the report is of the form
#     # seq id => list of drug classes to which there are major mutations
#         class_report = {}
#         for seq_id in dram_report:
#             class_list = set()
#             for drug_class, drug_value in dram_report[seq_id]['by class'].items():
#                 if len (drug_value['resistant']):
#                     class_list.add (drug_class)
#             class_report [seq_id] = list(class_list)
#         return class_report
# 
# 
#     return dram_report


if __name__ == "__main__":
    from sys import argv
    assert len (argv) > 1, "This program requires a single argument -- the .json file of sequences split by positions"

    mode = None

#     if len (argv) == 2:
#         with open (argv[1]) as positional_output:
#             DRAM_data = DRAM_calculator (json.load (positional_output))
#     else:
#         with open (argv[1]) as positional_output:
#             DRAM_data = json.load (positional_output)
#             mode = int (argv[2])
# 
#     print (json.dumps(DRAM_reporter (DRAM_data,report_mode = mode),sort_keys=True, indent=4))

