#! encoding: utf-8
from collections import namedtuple
try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict

from .creator import Creator
from .utils import get_seq


class Dataset(object):
    """User's class for making datasets of several formats. It needs as input
    a list of SeqRecord-expanded objects with as much info as possible:

    Parameters:
        seq_records (list):     SeqRecordExpanded objects. The list should be
                                sorted by gene_code and then voucher code.
        format (str):           NEXUS, PHYLIP, TNT, MEGA, GenBankFASTA.
        partitioning (str):     Partitioning scheme: ``by gene`` (default),
                                ``by codon position`` (each) and ``1st-2nd, 3rd``.
        codon_positions (str):  Can be ``1st``, ``2nd``, ``3rd``, ``1st-2nd``,
                                ``ALL`` (default).
        aminoacids (boolean):   Returns the dataset as aminoacid sequences.
        degenerate (str):       Method to degenerate nucleotide sequences,
                                following Zwick et al. Can be ``S``, ``Z``,
                                ``SZ`` and ``normal``.
        outgroup (str):         voucher code to be used as outgroup for NEXUS
                                and TNT files.

    Attributes:
         _gene_codes_and_lengths (dict):   in the form ``gene_code: list``
                                           The list contains sequence lengths for its
                                           sequences. We assume the longest to be the
                                           real gene_code sequence length.

    Example:

        >>> dataset = Dataset(seq_records, format='NEXUS', codon_positions='1st',
        ...                   partitioning='by gene',
        ...                   )
        >>> print(dataset.dataset_str)
        '#NEXUS
        blah blah
        '
        >>> dataset = Dataset(seq_records, format='PHYLIP', codon_positions='ALL',
        ...                   partitioning='by gene',
        ...                   )
        >>> print(dataset.dataset_str)
        '100 10
        blah blah
        '
    """
    def __init__(self, seq_records, format=None, partitioning=None,
                 codon_positions=None, aminoacids=None, degenerate=None,
                 outgroup=None):
        self.warnings = []
        self.seq_records = self.sort_seq_records(seq_records)
        self.gene_codes = None
        self.number_taxa = None
        self.number_chars = None
        self.reading_frames = {}

        self.format = format
        self.partitioning = partitioning
        self.codon_positions = codon_positions
        self.aminoacids = aminoacids
        self.degenerate = degenerate
        self.outgroup = None

        self._validate_codon_positions(codon_positions)
        self._validate_partitioning(partitioning)
        self._validate_outgroup(outgroup)

        self.data = None
        self._gene_codes_and_lengths = OrderedDict()
        self._prepare_data()
        self.extra_dataset_str = None
        self.dataset_str = self._create_dataset()

    def sort_seq_records(self, seq_records):
        """Makes sure that SeqExpandedRecords are sorted by gene_code and then
        by voucher code.
        """
        unsorted_gene_codes = set([i.gene_code for i in seq_records])
        sorted_gene_codes = list(unsorted_gene_codes)
        sorted_gene_codes.sort(key=lambda x: x.lower())

        unsorted_voucher_codes = set([i.voucher_code for i in seq_records])
        sorted_voucher_codes = list(unsorted_voucher_codes)
        sorted_voucher_codes.sort(key=lambda x: x.lower())

        if sorted_gene_codes == list(unsorted_gene_codes) and \
                sorted_voucher_codes == list(unsorted_voucher_codes):
            return seq_records
        else:
            sorted_seq_records = []
            append = sorted_seq_records.append
            for gene_code in sorted_gene_codes:
                for voucher_code in sorted_voucher_codes:
                    for seq_record in seq_records:
                        if seq_record.gene_code == gene_code and \
                                seq_record.voucher_code == voucher_code:
                            append(seq_record)
            return sorted_seq_records

    def _validate_partitioning(self, partitioning):
        if partitioning is None:
            self.partitioning = 'by gene'
        elif partitioning not in ['by gene', 'by codon position', '1st-2nd, 3rd']:
            raise AttributeError("Partitioning parameter should be one of these: "
                                 "None, 'by gene', 'by codon position', '1st-2nd, 3rd")
        elif partitioning in ['by codon position', '1st-2nd, 3rd'] \
                and self.degenerate:
            raise ValueError("Cannot degenerate if partitions scheme is {0!r}".format(
                partitioning))
        elif partitioning in ['by codon position', '1st-2nd, 3rd'] and self.format == 'MEGA':
            raise ValueError("Cannot produce MEGA dataset with codon positions in different partitions")

    def _validate_codon_positions(self, codon_positions):
        if codon_positions is None:
            self.codon_positions = 'ALL'
        elif codon_positions not in ['1st', '2nd', '3rd', '1st-2nd', 'ALL']:
            raise AttributeError("Codon positions parameter should be one of these: "
                                 "None, '1st', '2nd', '3rd', '1st-2nd', 'ALL'")

    def _validate_outgroup(self, outgroup):
        if outgroup:
            good_outgroup = False
            for seq_record in self.seq_records:
                if seq_record.voucher_code == outgroup:
                    good_outgroup = True
                    break
            if good_outgroup:
                self.outgroup = outgroup
            else:
                raise ValueError("The given outgroup {0!r} cannot be found in the"
                                 "input sequence records.")
        else:
            self.outgroup = None

    def _prepare_data(self):
        """
        Creates named tuple with info needed to create a dataset.

        :return: named tuple
        """
        self._extract_genes()
        self._extract_total_number_of_chars()
        self._extract_number_of_taxa()
        self._extract_reading_frames()

        Data = namedtuple('Data', ['gene_codes', 'number_taxa', 'number_chars',
                                   'seq_records', 'gene_codes_and_lengths',
                                   'reading_frames'])
        self.data = Data(self.gene_codes, self.number_taxa, self.number_chars,
                         self.seq_records, self._gene_codes_and_lengths,
                         self.reading_frames)

    def _extract_genes(self):
        gene_codes = [i.gene_code for i in self.seq_records]
        unique_gene_codes = list(set(gene_codes))
        # this is better: unique_gene_codes.sort(key=str.lower)
        # but will not work in python2
        unique_gene_codes.sort(key=lambda x: x.lower())
        self.gene_codes = unique_gene_codes

    def _extract_total_number_of_chars(self):
        """
        sets `self.number_chars` to the number of characters as string.
        """
        self._get_gene_codes_and_seq_lengths()

        sum = 0
        for seq_length in self._gene_codes_and_lengths.values():
            sum += sorted(seq_length, reverse=True)[0]
        self.number_chars = str(sum)

    def _get_gene_codes_and_seq_lengths(self):
        for seq_record in self.seq_records:
            if seq_record.gene_code not in self._gene_codes_and_lengths:
                self._gene_codes_and_lengths[seq_record.gene_code] = []

            if self.aminoacids:
                seq = seq_record.translate()
            elif not self.aminoacids and self.degenerate is not None:
                seq = seq_record.degenerate(method=self.degenerate)
            else:
                sequence = get_seq(seq_record, self.codon_positions)
                seq = sequence.seq
                if sequence.warning:
                    self.warnings.append(sequence.warning)
            self._gene_codes_and_lengths[seq_record.gene_code].append(len(seq))

    def _extract_number_of_taxa(self):
        """
        sets `self.number_taxa` to the number of taxa as string
        """
        n_taxa = dict()
        for i in self.seq_records:
            if i.gene_code not in n_taxa:
                n_taxa[i.gene_code] = 0
            n_taxa[i.gene_code] += 1
        number_taxa = sorted([i for i in n_taxa.values()], reverse=True)[0]
        self.number_taxa = str(number_taxa)

    def _extract_reading_frames(self):
        for seq_record in self.seq_records:
            if seq_record.gene_code not in self.reading_frames:
                self.reading_frames[seq_record.gene_code] = seq_record.reading_frame

    def _create_dataset(self):
        creator = Creator(self.data, format=self.format,
                          codon_positions=self.codon_positions,
                          partitioning=self.partitioning,
                          aminoacids=self.aminoacids,
                          degenerate=self.degenerate,
                          outgroup=self.outgroup,
                          )
        self.warnings = creator.warnings
        self.extra_dataset_str = creator.extra_dataset_str
        dataset_str = creator.dataset_str

        return dataset_str
