
import collections
import datetime
import random
import re
import StringIO
import sys

import bio
import features
import vcf

DELETE_BASE = '-'

class ProbabilisticFasta(object):
  '''
    generate fasta with probabilities attached from mapped fragments
  '''
  
  def __init__( self, log=None ):
    self.genome = {} # counts over position { 'A': [ 1, 0, 2, ... ], 'G': [ 2, 2, 0, ... ] } roughly means [AGG][GG][AA]...
    self.insertions = {} # { 0: { 'AAA': 1, 'BB': 2 } } means at position 0, insertion of AAA with 1 confidence; BB with 2 confidence
    self.inserted = {}
    self.deleted = set()
    self.prior = {} # overall counts
    self.total = 0 # total # of nucleotides
    self.length = 0 # length of genome
    self.log = log

  def add( self, fragment, start, confidence=1.0, debug=None ):
    '''
      process a directly mapped fragment of dna
      @fragment: the piece of dna
      @start: where in the reference genome it starts
      @confidence: how confident in this mapping
    '''
    for i in xrange(0, len(fragment) ):
      value = fragment[i]
      position = start + i
      if value not in self.genome:
        self.genome[value] = []
        self.prior[value] = 0
      if len(self.genome[value]) < position + 1:
        self.genome[value].extend( [0] * ( position - len(self.genome[value]) + 1 ) )
        if len(self.genome[value]) > self.length:
          self.length = len(self.genome[value])
      #self.log( "adding confidence %f to %s at %i; start %i" % ( confidence, value, position, start ) )
      #self.log( "capacity %i" % ( len(self.genome[value] ) ) )
      if position >= 0: # mappers can clip off the end
        self.genome[value][position] += confidence
        self.prior[value] += 1
        self.total += 1
        #if position == 39415: # debugging
        #  self.log( 'start %i: added %f to %s, total %f, line %s' % ( start, confidence, value, self.genome[value][position], debug ) )

  def insert( self, fragment, start, confidence=1.0 ):
    '''
      @fragment: the piece of dna to insert
      @start: where the insertion starts on the reference
      @confidence: how confident in this insertion
    '''
    if start not in self.insertions:
      self.insertions[start] = {}
    if fragment not in self.insertions[start]:
      self.insertions[start] = { fragment: 0 }
    self.insertions[start][fragment] += confidence

  def delete( self, start, count, confidence=1.0 ):
    fragment = DELETE_BASE * count
    self.add( fragment, start, confidence )

  def count( self, position ):
    '''
      counts of nucleotides seen at a position
      e.g. { 'A': 10, 'G': 4 }
    '''
    result = collections.defaultdict(int)
    for value in self.genome:
      if len(self.genome[value]) > position:
        result[value] = self.genome[value][position]
    return result

  def confidence( self, position ):
    '''
      returns probabilities of nucleotides
      TODO priors
    '''
    counts = self.count( position ) # e.g. { 'A': 4, 'C': 1 }
    denominator = sum( counts[k] for k in counts ) + len(self.prior) # laplacian smoothing
    probabilities = {}
    for count in counts:
      probabilities[count] = ( counts[count] + 1.0 ) / denominator