#
# Utility to insert an hgvs variant into a transcript sequence.
# Generates a record corresponding to the modified transcript sequence,
# along with annotations for use in conversion to an hgvsp tag.
# Used in hgvsc to hgvsp conversion.
#
import math
import recordtype

from Bio.Seq import Seq

import hgvs.edit
from hgvs.location import CDS_START, CDS_END

DBG = False

class AltTranscriptData(recordtype.recordtype('AltTranscriptData', [
        'transcript_sequence', 'aa_sequence', 'cds_start', 'cds_stop', 'protein_accession',
        ('is_frameshift', False), ('variant_start_aa', None), ('frameshift_start', None), ('is_substitution', False),
        ('is_ambiguous', False)])):

    @classmethod
    def create_for_variant_inserter(cls, seq, cds_start, cds_stop, is_frameshift, variant_start_aa, accession,
                                    is_substitution=False, is_ambiguous=False):
        """Create a variant sequence using inputs from VariantInserter
        :param seq: DNA sequence wiith variant incorporated
        :type str or list
        :param cds_start: coding sequence start (1-based)
        :type int
        :param cds_stop: coding sequence stop (1-based)
        :type int
        :param is_frameshift: is this variant a frameshift
        :type bool
        :param variant_start_aa: AA start index (1-based) for this variant
        :type int
        :param accession: protein accession, e.g. NP_999999.2
        :type str
        :return variant sequence data
        :type recordtype
        """
        if isinstance(seq, basestring):
            seq = list(seq)
        seq_cds = seq[cds_start - 1:]
        if len(seq_cds) % 3 != 0:   # padding so biopython won't complain during the conversion
            seq_cds.extend(['N']*((3-len(seq_cds) % 3) % 3))
        seq_cds = ''.join(seq_cds)
        seq_aa = str(Seq(seq_cds).translate())
        stop_pos = seq_aa.find("*")
        if stop_pos != -1:
            seq_aa = seq_aa[:stop_pos + 1]

        alt_data = AltTranscriptData(''.join(seq), seq_aa, cds_start, cds_stop, accession,
                                     is_frameshift=is_frameshift, variant_start_aa=variant_start_aa,
                                     is_substitution=is_substitution, is_ambiguous=is_ambiguous)

        return alt_data


class AltSeqBuilder(object):

    EXON = "exon"
    INTRON = "intron"
    F_UTR = "five utr"
    T_UTR = "three utr"

    def __init__(self, var_c, transcript_data):
        """Constructor

        :param variant representation of hgvs variant
        :type variant
        :param transcript_data representation of transcript
        :type recordtype

        """
        self._var_c = var_c
        self._transcript_data = transcript_data
        if DBG:
            print transcript_data.transcript_sequence

    def build_altseq(self):
        """given a variant and a sequence, incorporate the variant and return the new sequence

        Data structure returned is analogous to the data structure used to return the variant sequence,
        but with an additional parameter denoting the start of a frameshift that should affect all bases
        downstream.

        :returns variant sequence data
        :type list of dictionaries
        """
        NOT_CDS = "not_cds_variant"

        type_map = {hgvs.edit.NARefAlt: self._incorporate_delins,
                    hgvs.edit.Dup: self._incorporate_dup,
                    hgvs.edit.Repeat: self._incorporate_repeat,
                    NOT_CDS: self._create_alt_eq_ref_noncds
                    }

        # TODO - loop over each allele rather than assume only 1 variant; return a list for now
        alt_data = []

        variant_location = self._get_variant_region()

        if variant_location == self.EXON:
            edit_type = type(self._var_c.posedit.edit)
        elif variant_location == self.INTRON:
            edit_type = NOT_CDS
        elif variant_location == self.T_UTR:
            edit_type = NOT_CDS
        elif variant_location == self.F_UTR:
            # TODO - handle case where variant introduces a Met (new start)
            edit_type = NOT_CDS
        else:   # should never get here
            raise ValueError("value_location = {}".format(variant_location))

        try:
            this_alt_data = type_map[edit_type]()
        except KeyError as e:
            raise NotImplementedError("c to p translation unsupported for {} type {}".format(self._var_c, edit_type))

        # get the start of the "terminal" frameshift (i.e. one never "cancelled out")
        this_alt_data = self._get_frameshift_start(this_alt_data)
        alt_data.append(this_alt_data)
        if DBG:
            print this_alt_data.transcript_sequence

        return alt_data

    def _get_variant_region(self):
        """Categorize variant by location in transcript (5'utr, exon, intron, 3'utr)
        :return "exon", "intron", "five_utr", "three_utr"
        """
        if self._var_c.posedit.pos.start.datum == CDS_END and self._var_c.posedit.pos.end.datum == CDS_END:
            result = self.T_UTR
        elif self._var_c.posedit.pos.start.base < 0 and self._var_c.posedit.pos.end.base < 0:
            result = self.F_UTR
        elif self._var_c.posedit.pos.start.offset != 0 or self._var_c.posedit.pos.end.offset != 0:
            # leave out anything intronic for now
            result = self.INTRON
        else:   # anything else that contains an exon
            result = self.EXON
        return result

    # def _is_intron_only(self):
    #     """Checks if variant is entirely intronic"""
    #
    #     # case 1: same base,same sign on offset -> start, end are anchored from the same base
    #     same_start = self._var_c.posedit.pos.start.base == self._var_c.posedit.pos.end.base
    #     pos_offset = self._var_c.posedit.pos.start.offset > 0 and self._var_c.posedit.pos.end.offset > 0
    #     neg_offset = self._var_c.posedit.pos.start.offset < 0 and self._var_c.posedit.pos.end.offset < 0
    #
    #     # case 2: start, end between 2 different exon bases but don't overlap either
    #     same_start_plus_1 = (self._var_c.posedit.pos.start.base + 1) == self._var_c.posedit.pos.end.base
    #     pm_offset = self._var_c.posedit.pos.start.offset > 0 and self._var_c.posedit.pos.end.offset < 0
    #
    #     # other types of introns would overlap at least 1 exon base
    #
    #     return (same_start and (pos_offset or neg_offset)) or (same_start_plus_1 and pm_offset)


    def _incorporate_delins(self):
        """Incorporate delins"""
        seq, cds_start, cds_stop, start, end = self._setup_incorporate()

        ref = self._var_c.posedit.edit.ref
        alt = self._var_c.posedit.edit.alt
        ref_length = end - start if ref is not None else 0  # can't just get from ref since ref isn't always known
        alt_length = len(self._var_c.posedit.edit.alt) if self._var_c.posedit.edit.alt is not None else 0
        net_base_change = alt_length - ref_length
        cds_stop += net_base_change

        # incorporate the variant into the sequence (depending on the type)
        is_substitution = False
        if ref is not None and alt is not None:     # delins or SNP
            seq[start:end] = list(alt)
            if len(ref) == 1 and len(alt) == 1:
                is_substitution = True
        elif ref is not None:                       # deletion
            del seq[start:end]
        else:                                       # insertion
            seq[start + 1:start + 1] = list(alt)    # insertion in list before python list index

        if DBG:
            print "net base change: {}".format(net_base_change)
        is_frameshift = net_base_change % 3 != 0
        # use max of mod 3 value and 1 (in event that indel starts in the 5'utr range)
        variant_start_aa = max(int(math.ceil((self._var_c.posedit.pos.start.base) / 3.0)), 1)

        alt_data = AltTranscriptData.create_for_variant_inserter(seq, cds_start, cds_stop,
                                                                 is_frameshift, variant_start_aa,
                                                                 self._transcript_data.protein_accession,
                                                                 is_substitution=is_substitution)
        return alt_data

    def _incorporate_dup(self):
        """Incorporate dup into sequence"""
        seq, cds_start, cds_stop, start, end = self._setup_incorporate()

        dup_seq = seq[start:end]
        seq[end:end] = dup_seq
        
        is_frameshift = len(dup_seq) % 3 != 0
        variant_start_aa = int(math.ceil((self._var_c.posedit.pos.end.base + 1) / 3.0))

        alt_data = AltTranscriptData.create_for_variant_inserter(seq, cds_start, cds_stop,
                                                                 is_frameshift, variant_start_aa,
                                                                 self._transcript_data.protein_accession)
        return alt_data

    def _incorporate_repeat(self):
        """Incorporate repeat int sequence"""
        # TODO - implement
        raise NotImplementedError("hgvs c to p conversion does not support {} type: repeats".format(self._var_c))

    def _setup_incorporate(self):
        """Helper to setup incorporate functions
        :return (transcript sequence, cds start [1-based], cds stop [1-based],
        cds start index in seq [inc, 0-based], cds end index in seq [excl, 0-based])
        :type (list, int, int, int, int)
        """
        seq = list(self._transcript_data.transcript_sequence)

        # get initial start/end points; will modify these based on the variant length
        cds_start = self._transcript_data.cds_start
        cds_stop = self._transcript_data.cds_stop

        start_end = []
        for pos in (self._var_c.posedit.pos.start, self._var_c.posedit.pos.end):
            # list is zero-based; seq pos is 1-based
            if pos.datum == CDS_START:
                if pos.base < 0:                  # 5' UTR
                    result = cds_start - 1
                else:                             # cds/intron
                    if pos.offset <= 0:
                        result = (cds_start - 1) + pos.base - 1
                    else:
                        result = (cds_start - 1) + pos.base
            elif pos.datum == CDS_END:            # 3' UTR
                result = cds_stop + pos.base - 1
            else:
                raise NotImplementedError("Unsupported/unexpected location")
            start_end.append(result)

        # unpack; increment end by 1 (0-based exclusive)
        (start, end) = start_end
        end += 1

        if DBG:
            print "len seq:{} cds_start:{} cds_stop:{} start:{} end:{}".format(len(seq), cds_start, cds_stop, start, end)
        return seq, cds_start, cds_stop, start, end

    def _create_alt_eq_ref_noncds(self):
        """Create an alt seq that matches the reference (for non-cds variants)"""
        alt_data = AltTranscriptData.create_for_variant_inserter(list(self._transcript_data.transcript_sequence),
                                                                 self._transcript_data.cds_start,
                                                                 self._transcript_data.cds_stop,
                                                                 False,
                                                                 None,
                                                                 self._transcript_data.protein_accession,
                                                                 is_ambiguous=True)
        return alt_data

    def _get_frameshift_start(self, variant_data):
        """Get starting position (AA ref index) of the last frameshift
        which affects the rest of the sequence, i.e. not offset by subsequent frameshifts
        :param variant_data: info on each variant
        :type RecordType
        :return variant data with additional field for AA index (1-based) of the frameshift start
        """
        # TODO - implement for 2+ variants

        if DBG:
            print "is_frameshift:{}".format(variant_data.is_frameshift)
            print "variant_start_aa:{}".format(variant_data.variant_start_aa)
        if variant_data.is_frameshift:
            variant_data.frameshift_start = variant_data.variant_start_aa
        return variant_data
## <LICENSE>
## Copyright 2014 HGVS Contributors (https://bitbucket.org/invitae/hgvs)
## 
## Licensed under the Apache License, Version 2.0 (the "License");
## you may not use this file except in compliance with the License.
## You may obtain a copy of the License at
## 
##     http://www.apache.org/licenses/LICENSE-2.0
## 
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.
## </LICENSE>
