#!/usr/bin/env python
# -*- coding: utf-8

import os
import sys

import anvio
import anvio.tables as t
import anvio.utils as utils
import anvio.terminal as terminal
import anvio.summarizer as summarizer
import anvio.filesnpaths as filesnpaths

from anvio.errors import ConfigError, FilesNPathsError


__author__ = "A. Murat Eren"
__copyright__ = "Copyright 2016, The anvio Project"
__credits__ = ["Tom O. Delmont"]
__license__ = "GPL 3.0"
__version__ = anvio.__version__
__maintainer__ = "A. Murat Eren"
__email__ = "a.murat.eren@gmail.com"


run = terminal.Run()
progress = terminal.Progress()
pp = terminal.pretty_print


def main(args):
    filesnpaths.check_output_directory(args.output_dir)
    profile = utils.get_TAB_delimited_file_as_dictionary(args.profile)
    run.info('Num entries', pp(len(profile)))

    for field in [f for f in t.variable_nts_table_structure[1:] if f != 'split_name']:
        if field not in list(profile.values())[0]:
            raise ConfigError("The input file does not look like it is generated by anvi-gen-variability-profile.\
                                It is missing at least one field that should have appeared in this file (%s)" % field)

    nt_positions = set(e['unique_pos_identifier'] for e in list(profile.values()))
    run.info('Num positions', pp(len(nt_positions)))

    samples = sorted(list(set([e['sample_id'] for e in list(profile.values())])))
    run.info('Num samples', pp(len(samples)))

    data_dict = {}
    additional_data_dict = {}
    for entry in list(profile.values()):
        unique_pos_identifier = entry['unique_pos_identifier']
        sample = entry['sample_id']
        departure_from_consensus = float(entry['departure_from_consensus'])

        if unique_pos_identifier not in data_dict:
            data_dict[unique_pos_identifier] = dict(list(zip(samples, [0.0] * len(samples))))
            additional_data_dict[unique_pos_identifier] = {'Competing NTs': None, 'Position in codon': None}

        if departure_from_consensus > 0.01:
            additional_data_dict[unique_pos_identifier]['Competing NTs'] = entry['competing_nts']

        if departure_from_consensus > args.min_departure_from_consensus and departure_from_consensus < args.max_departure_from_consensus:
            data_dict[unique_pos_identifier][sample] = departure_from_consensus

        additional_data_dict[unique_pos_identifier]['Position in codon'] = {0: None, 1: '1st', 2: '2nd', 3: '3rd'}[int(entry['base_pos_in_codon'])]

    temp_view_data_file_path = filesnpaths.get_temp_file_path()
    temp_additional_data_file_path = filesnpaths.get_temp_file_path()

    utils.store_dict_as_TAB_delimited_file(data_dict, temp_view_data_file_path, headers = ['contig'] + samples)
    utils.store_dict_as_TAB_delimited_file(additional_data_dict, temp_additional_data_file_path, headers = ['contig'] + ['Competing NTs', 'Position in codon'])

    g = summarizer.AdHocRunGenerator(temp_view_data_file_path)
    g.linkage = 'ward'
    g.additional_view_data_file_path = temp_additional_data_file_path
    g.output_directory = args.output_dir
    g.generate()

    os.remove(temp_view_data_file_path)
    os.remove(temp_additional_data_file_path)

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='Take the output of anvi-gen-variability-profile, prepare an output for interactive interface')

    parser.add_argument('profile', help = 'The output file generated by anvi-gen-variability-profile')
    parser.add_argument('--min-departure-from-consensus', type = float, default = 0.00, metavar = "FLOAT",
                        help = "Minimum departure from consensus at a given variable nucleotide position. The default\
                        is %(default).2f.")
    parser.add_argument('--max-departure-from-consensus', type = float, default = 0.99, metavar = "FLOAT",
                        help = "Maximum departure from consensus at a given variable nucleotide position. The default\
                        is %(default).2f.")
    parser.add_argument(*anvio.A('output-dir'), **anvio.K('output-dir', {'required': True}))

    args = parser.parse_args()

    try:
        main(args)
    except ConfigError as e:
        print(e)
        sys.exit(-1)
    except FilesNPathsError as e:
        print(e)
        sys.exit(-1)
