#!/usr/bin/python

# Read an xmfa alignment and write some summary statistics

#from calc_similarity_profile_xmfa import *
from glint_misc import write_grfile
from Bio import SeqIO
import sys
import numpy
import re

LONG_GAP_THRESHOLD = 2
xmfa_block_header_re = re.compile("> ([0-9]*):([0-9]*)-([0-9]*) ([+-]) (.+)")

def parse_xmfa_header(instr):
  # Parse an xmfa header and return an ordered list of sequence names
  # For now we just associate them with seq1/seq2 in order

  seqnames = []

  loc = instr.tell()
  line = instr.readline()
  while (line[0] == "#"):
    if line[1:9] == "Sequence" and line[10:14] == "File":
      seqnames.append((line.split())[1])

    loc=instr.tell()
    line = instr.readline()

  instr.seek(loc)
  return(seqnames)

def parse_xmfa_block(instr):
  # Parse all entries coming up to the next '=' delimiter
  # return a dictionary containing the sequences indexed by name

  seq_dict = {}
  currloc = instr.tell()
  thisline = instr.readline()
  if (thisline == ""):
    return None

  while (thisline[0] != "="):
    #print "Working on block with head %s" % thisline
    header_match = xmfa_block_header_re.match(thisline)
    if header_match is None:
      raise("Failed to parse xmfa header line")

    matches = header_match.groups()
    seq_ind = int(matches[0])-1
    seq_start = int(matches[1])
    seq_end = int(matches[2])
    seq_ori = matches[3]
    seq_name = matches[4]
    this_seq = ""

    thisline = instr.readline()
    while (thisline[0] != ">" and thisline[0] != "="):
      this_seq += thisline.rstrip()
      thisline = instr.readline()

    seq_dict[seq_name] = { 'seq' : this_seq, 'start' : seq_start, 'end' : seq_end, 'ori' : seq_ori }

  return(seq_dict)



def update_matches_from_block(seqblock, strain_1_maps, strain_2_maps, strain_1_flags, strain_2_flags, strain_1_revmap, strain_2_revmap, strain_1_gapids, strain_2_gapids, strain_1_currgap, strain_2_currgap):
  # given an xmfa block, update some mapping information on the basis of matches/
  #  /mismatches within it

  block_1 = matches[seqnames[0]]
  block_2 = matches[seqnames[1]]
  start_1 = block_1['start']
  end_1 = block_1['end']
  seq_1 = block_1['seq']
  ori_1 = block_1['ori']
  step_1 = 1
  start_2 = block_2['start']
  end_2 = block_2['end']
  seq_2 = block_2['seq']
  ori_2 = block_2['ori']
  step_2 = 1

  # variables for keeping track of gaps in one sequence or the other
  # these should both be tuples giving the real positions in seq1 and 2 at the start
  this_gapstart_1 = None; # start position of current gap in seq 1
  this_gapstart_2  = None ;# start of current gap in gap 2

  # properly reverse things if we have a reversed alignment
  # each offset is our current position in the corresponding genome

  offset_1 = start_1 - 1
  offset_2 = start_2 - 1

  if (ori_1 == '-'):
    offset_1 = end_1 - 1
    step_1 = -1
  if (ori_2 == '-'):
    offset_2 = end_2 - 1
    step_2 = -1


  seqlen = len(seq_1)
  if (seqlen != len(seq_2)):
    raise("Sequences from block did not have equal length")

  for i in range(seqlen):
    base_1 = seq_1[i]
    base_2 = seq_2[i]

    if ((base_1 != '-') and  (base_2 != '-')):

      # match these bases
      strain_1_maps[offset_1].append(offset_2)
      strain_2_maps[offset_2].append(offset_1)
      strain_1_flags[offset_1] += 1
      strain_2_flags[offset_2] += 1
      strain_1_gapids[offset_1] = 0
      strain_2_gapids[offset_2] = 0
      strain_1_revmaps[offset_1] = offset_2
      strain_2_revmaps[offset_2] = offset_1
      if strain_1_flags[offset_1] > 1 or strain_2_flags[offset_2] > 1:
        print offset_1
        print offset_2
        print strain_1_flags[offset_1]
        print strain_2_flags[offset_2]
        raise("Found a multiple assignment of a genomic fragment!")

      # if there is currently a gap open, close it

      if ((this_gapstart_1 is not None) and (this_gapstart_2 is not None)):
        # handle the case where we had gaps in each sequence next to each other
        start1_1,start1_2 = this_gapstart_1
        start2_1, start2_2 = this_gapstart_2
        gaplen_1 = numpy.abs(offset_2 - start1_2)-1 ;# length of the gap in seq 1
        gaplen_2 = numpy.abs(offset_1 - start2_1)-1 ;# length of the gap in seq 2
        gapstart_1 = min(start1_2, start2_2, offset_2)
        gapstart_2 = min(start1_1, start2_1, offset_1)

        leftedge_in_1 = min(start1_1, start2_1, offset_1)
        rightedge_in_1 = max(start1_1, start2_1, offset_1)
        leftedge_in_2 = min(start1_2, start2_2, offset_2)
        rightedge_in_2 = max(start1_2, start2_2, offset_2)

        assign_gaps_to_endpoints(strain_1_maps,strain_2_revmaps,strain_2_flags,gapstart_1,gaplen_1,leftedge_in_1,rightedge_in_1, ori_1,ori_2)
        assign_gaps_to_endpoints(strain_2_maps,strain_1_revmaps,strain_1_flags,gapstart_2,gaplen_2,leftedge_in_2,rightedge_in_2, ori_2,ori_1)

        # also assign gap ids as appropriate
        if (gaplen_2 > LONG_GAP_THRESHOLD):
          strain_1_gapids[(gapstart_2 +1):gapstart_2 + gaplen_2+1] = strain_1_currgap
          strain_1_currgap += 1
        else:
          strain_1_gapids[(gapstart_2 +1):gapstart_2 + gaplen_2+1] = 1


        if (gaplen_1 > LONG_GAP_THRESHOLD):
          strain_2_gapids[(gapstart_1 +1):gapstart_1 + gaplen_1+1] = strain_2_currgap
          strain_2_currgap += 1
        else:
          strain_2_gapids[(gapstart_1 +1):gapstart_1 + gaplen_1+1] = 1

        
        this_gapstart_1 = None
        this_gapstart_2 = None


      if (this_gapstart_1 is not None):
        # assign the intervening range from sequence 2 to the endpoints in sequence 1
        start_1, start_2 = this_gapstart_1
        gaplen = numpy.abs(offset_2 - start_2)-1
        gapstart = min(start_2, offset_2)

        assign_gaps_to_endpoints(strain_1_maps,strain_2_revmaps,strain_2_flags, gapstart,gaplen,min(start_1,offset_1), max(start_1,offset_1),ori_1,ori_2)

        if gaplen > LONG_GAP_THRESHOLD:
          strain_2_gapids[(gapstart+1):gapstart+gaplen+1] = strain_2_currgap
          strain_2_currgap += 1
        else:
          strain_2_gapids[(gapstart+1):gapstart+gaplen+1] = 1

        this_gapstart_1 = None

      if (this_gapstart_2 is not None):
        # assign the intervening range from sequence 1 to the endpoints in sequence 2
        start_1, start_2 = this_gapstart_2
        gaplen = numpy.abs(offset_1 - start_1)-1
        gapstart = min(start_1, offset_1)

        assign_gaps_to_endpoints(strain_2_maps,strain_1_revmaps,strain_1_flags,gapstart,gaplen,min(start_2,offset_2),max(start_2,offset_2), ori_1,ori_2)

        if gaplen > LONG_GAP_THRESHOLD:
          strain_1_gapids[(gapstart+1):(gapstart+gaplen+1)] = strain_1_currgap
          strain_1_currgap += 1
        else:
          strain_1_gapids[(gapstart+1):(gapstart+gaplen+1)] = 1

        
        this_gapstart_2 = None

    if ( (base_1 == '-') and (base_2 != '-')):
      # open a gap in sequence 1
      if this_gapstart_1 is None:
        this_gapstart_1 = (offset_1-step_1, offset_2-step_2)

    if ( (base_1 != '-') and (base_2 == '-')):
      # open a gap in sequence 2
      if this_gapstart_2 is None:
        this_gapstart_2 = (offset_1-step_1, offset_2-step_2)

    # update offsets
    if (base_1 != '-'):
      offset_1 += step_1

    if (base_2 != '-'):
      offset_2 += step_2


  # return the updated maps
  return( (strain_1_maps, strain_2_maps, strain_1_flags, strain_2_flags, strain_1_revmaps, strain_2_revmaps, strain_1_gapids, strain_2_gapids, strain_1_currgap, strain_2_currgap))


def assign_gap(startind, target_map, target_flags, target_revmap, target_gapids, origin_map, origin_flags, origin_revmap, origin_gapids, origin_currgap):
  # assign an insertion starting at startind in the origin sequence to 
  #  the corresponding endpoints in the target
  # return the index in the origin sequence immediately after the end of the gap

  # first find the extent of the gap
  end_ind = startind 

  # do a special check to see if we're at the very beginning
  if startind == 0:
    while (origin_flags[startind] < 1):
      startind -= 1

    startind += 1

  while origin_flags[end_ind] < 1:
    end_ind += 1


    # do a special check to make sure we don't run off the end
    if end_ind == len(origin_flags):
      # this has already been handled if the first base was unassigned
      # otherwise, we should add it now

      if (origin_flags[0] < 1):
        end_ind = -1
        
      else:
        end_ind = 0

      break

  gaplen = numpy.abs(end_ind - startind)

  if (end_ind == 0):
    gaplen = len(origin_flags) - startind

  if (end_ind < 0):
    return(len(origin_flags) + 1)

  # find the points in the target sequence mapping to these endpoints
  # both end_ind and (startind-1) should map

  start_map = None
  end_map = None

  start_target = startind-1
  end_target = end_ind

  origin_flags[startind:end_ind] += 1
  
  if gaplen > LONG_GAP_THRESHOLD:
    origin_gapids[startind:end_ind] = origin_currgap
    origin_currgap += 1
  else:
    origin_gapids[startind:end_ind] = 1

  start_map = origin_revmap[start_target]
  end_map = origin_revmap[end_target]


  if end_ind >= len(origin_flags):
    end_map = 0

  if (start_map == None or end_map == None):
    raise(ValueError("Failed to find an endpoint for gap"))


  target_map[start_map] += range(startind, startind+gaplen/2)
  target_map[end_map] += range(startind + gaplen/2, startind+gaplen)


  if (end_ind < startind):
    return(len(origin_flags))

  return(end_ind, origin_currgap)


def write_all_weights(outfile, seqlen, strain_maps, strain_gapids, other_maps):
  # write all quantities of interest to a table
  all_locs = numpy.array(range(seqlen))
  all_dir_weights  = calc_dir_weights(seqlen, strain_maps)
  all_inv_weights = calc_inv_weights(seqlen, other_maps)

  outstr = open(outfile, 'w')
  outstr.write('%s;%s;%s;%s\n' % ('loc', 'dir_weight', 'inv_weight', 'gap_id'))

  for i in range(len(all_locs)):
    outstr.write('%i;%i;%f;%i\n' % (all_locs[i]+1, all_dir_weights[i], all_inv_weights[i], strain_gapids[i]))

  outstr.close()
  
def write_debug_maps(outfile,mappings,gapids):
  # write the mappings and gap ids to a debug file
  ostr=open(outfile,'w')
  for i in range(len(mappings)):
    ostr.write('%i\t%s\t%s\n' % (i+1,mappings[i],gapids[i]))

  ostr.close()

def calc_dir_weights(seqlen, strain_maps):
  # calculate direct weights
  # for each base, this is the number of bp in the other genome that
  # map to this one. Thus, this will be high at sites that flank insertions in the
  # other genome

  return([len(strain_maps[i]) for i in range(seqlen)])

def calc_inv_weights(seqlen1, strain_maps_2):
  # calculate the inverse weights
  # for each base in sequence 1, the inverse weight is 1/X, where X is 
  # the number of bases in seq1 that map to a given base in seq2

  weights = [0 for i in range(seqlen1)]

  for i in range(len(strain_maps_2)):
    this_maps = strain_maps_2[i]
    if len(this_maps) < 1:
      continue

    thisweight = 1.0/len(this_maps)

    for j in this_maps:
      weights[j] = thisweight

  return(weights)

def assign_gaps_to_endpoints(target_maps, other_revmaps, other_flags, gapstart, gaplen, leftedge,rightedge, ori_target, ori_other):
  """
  Helper function that assigns a gap appropriately

  target_maps -- the forward map for the strain containing the gap
  other_revmaps -- the reverse map for the other strain
  other_flags -- flags indicating the gap status of other_strain
  gapstart,gaplen -- start and length of the gap in other_strain
  leftedge,rightedge -- bases in target strain aligning to the ends of the gap
  ori_target, ori_other -- orientations of the two strains in this aligned block
  """

  firstblock = range(gapstart+1, gapstart+gaplen / 2+1)
  secondblock = range(gapstart+gaplen/2+1, gapstart+gaplen+1)

  # make sure we orient the aligned region appropriately
  if ori_target == ori_other:
    leftblock = firstblock
    rightblock = secondblock
  else:
    leftblock=secondblock
    rightblock=firstblock

  # now update the arrays appropriately
  target_maps[leftedge] += leftblock
  target_maps[rightedge] += rightblock

  other_revmaps[leftblock] = leftedge
  other_revmaps[rightblock] = rightedge

  other_flags[(gapstart+1):(gapstart+gaplen+1)]+=1


if __name__ == "__main__":

  if len(sys.argv) != 7:
    print "Usage: calc_seq_mapping_from_xmfa.py mauve_xmfa_file fastafile1 fastafile2 seqname_1 seqname_2 output_prefix"
    sys.exit()

  mauve_outfile = sys.argv[1]
  seq1file = sys.argv[2]
  seq2file = sys.argv[3]
  seq1name = sys.argv[4]
  seq2name = sys.argv[5]
  outprefix=sys.argv[6]

  # load up the data
  seq1 = SeqIO.read(open(seq1file), "fasta")
  seq2 = SeqIO.read(open(seq2file), "fasta")
  len1 = len(seq1.seq)
  len2 = len(seq2.seq)

  # initialize arrays to store mapping information

  ## maps ofthe bases in each sequence corresponding to the other
  strain_1_maps = [ [] for i in range(len1)] 
  strain_2_maps = [ [] for i in range(len2)] 

  ## flags for whether each base has been assigned
  strain_1_flags = numpy.zeros(len1)
  strain_2_flags = numpy.zeros(len2)

  # island IDs for each sequence
  # 0 is for bases that align, and 1 is for gaps less than LONG_GAP_THRESHOLD
  # all others in each sequence are numbered sequentially
  strain_1_gapids = numpy.zeros(len1) - 1
  strain_2_gapids = numpy.zeros(len2) - 1
  strain_1_currgap = 2
  strain_2_currgap = 2

  # list telling which base in the other genome this one maps to
  strain_1_revmaps = numpy.array([-1 for i in range(len1)])
  strain_2_revmaps = numpy.array([-1 for i in range(len2)])


  # Now go through all of the aligned regions from the mauve alignment file and flag bases that match
  instr = open(mauve_outfile)

  seqnames = parse_xmfa_header(instr)

  # process each set of alignments
  matches = "hello"
  print "Processing alignments"
  while (matches is not None):


    matches = parse_xmfa_block(instr)
    if (matches is None):
      break
    if len(matches.keys()) <2:
      continue
    if len(matches.keys()) > 2:
      raise("Got a block with more than 2 entries... I don't know how to handle this")

    # parse the block and update arrays accordingly
    strain_1_maps, strain_2_maps, strain_1_flags, strain_2_flags, strain_1_revmaps, strain_2_revmaps, strain_1_gapids, strain_2_gapids, strain_1_currgap, strain_2_currgap = update_matches_from_block(matches, strain_1_maps, strain_2_maps, strain_1_flags, strain_2_flags, strain_1_revmaps, strain_2_revmaps, strain_1_gapids, strain_2_gapids, strain_1_currgap, strain_2_currgap)
    
  # now deal with any gaps remaining
  # as a simple heuristic, for each gap in one sequence, we assign that gap to the two
  # endpoints in the other sequence that best match it

  # first go through sequence 1
  print "looking for orphans in sequence 1"
  i=0
  while (i < len1):

    #if (i % 1000 == 0):
    #  print i 

    if (strain_1_flags[i] < 1):
      i,strain_1_currgap=assign_gap(i, strain_2_maps, strain_2_flags, strain_2_revmaps, strain_2_gapids,strain_1_maps, strain_1_flags, strain_1_revmaps, strain_1_gapids, strain_1_currgap)
    else:
      i += 1

  print "looking for orphans in sequence 2"
  i=0
  while (i < len2):
    #if (i % 1000 == 0):
    #  print i 
    if (strain_2_flags[i] < 1):
      i, strain_2_currgap=assign_gap(i, strain_1_maps, strain_1_flags, strain_1_revmaps, strain_1_gapids, strain_2_maps, strain_2_flags, strain_2_revmaps, strain_2_gapids, strain_2_currgap)
    else:
      i += 1


  # now make sure everything got assigned
  n_zero_1 = 0
  n_doub_1 = 0
  n_zero_2 = 0
  n_doub_2 = 0

  for flag in strain_1_flags:
    if flag < 0:
      n_zero_1 += 1
    if flag > 1:
      n_doub_1 += 1

  for flag in strain_2_flags:
    if flag < 0:
      n_zero_2 += 1
    if flag > 1:
      n_doub_2 += 1

  print "-------------------------------------------------"
  print "Summary of alignment:"
  print ""
  print "Abberations in assignment:"
  print " Unassigned from 1: %i" % n_zero_1
  print " Multiply assigned from 1: %i" % n_doub_1
  print " Unassigned from 2: %i" % n_zero_2
  print " Multiply assigned from 2: %i" % n_doub_2

  # write all of the resulting weights and other daya
  # direct weights are the number of bases in the other sequence that
  #  map to each base in this one

  # inverse weights are the number of bases that share this base's assignment 
  #  in the other sequence 

  ## write some debug information
  #write_debug_maps('%s_strain1_debug_map.txt' % outprefix,strain_1_maps,strain_1_gapids)
  #write_debug_maps('%s_strain2_debug_map.txt' % outprefix,strain_2_maps,strain_2_gapids)


  write_all_weights('%s_%s_vs_%s_allweights.txt' % (outprefix, seq1name, seq2name), len1, strain_1_maps, strain_1_gapids, strain_2_maps)
  write_all_weights('%s_%s_vs_%s_allweights.txt' % (outprefix, seq2name, seq1name), len2, strain_2_maps, strain_2_gapids, strain_1_maps)


#  # also write gr file
#  write_direct_weights("seq1_direct_weights.txt", len1, strain_1_maps)
#  write_direct_weights("seq2_direct_weights.txt", len2, strain_2_maps)
#
#
#  # and the inverse weights, the number of bases that share this base's assignment
#  #  in the other sequence
#  write_inverse_weights("seq1_inverse_weights.txt", len1, strain_2_maps)
#  write_inverse_weights("seq2_inverse_weights.txt", len2, strain_1_maps)

  print "calc_seq_mapping_from_xmfa.py finished successfully"

