#!/usr/bin/python

# filter a table to contain only bases present in a given score file
# we record in each position in the table the average of the bp level data that it represents
# for the gapid, we instead take the one most represented by the bp covered by the probe

import sys
import table
import numpy

def merge_bp_data(target_inds,dirweights,invweights,gapids):
  # appropriately merge data from a given set of probes
  # we average the direct and inverse weights, and choose the most frequent
  #  gapid among the probes here
  # return a tuple containing the three merged values

  new_dirweight = numpy.mean(dirweights[target_inds])
  new_invweight = numpy.mean(invweights[target_inds])
  new_gapid = numpy.argmax(numpy.bincount(gapids[target_inds]))

  return (new_dirweight,new_invweight,new_gapid)

def read_grfile(filename, skiprows=0):
  """
  Read data from a .gr file into a numpy array

  In the process we transpose the array, so that it has a[0] the indices and a[1] the values
  """

  locs = numpy.loadtxt(filename, dtype='int', usecols=(0,), skiprows=skiprows)
  vals = numpy.loadtxt(filename, usecols=(1,), skiprows=skiprows)
  return (locs,vals)


intab=sys.argv[1]
scorefile=sys.argv[2]
outtab=sys.argv[3]

mytab = table.read_table(intab, sep=";", header=True, guesstypes=True)
locs, vals = read_grfile(scorefile, skiprows=1)

genome_locs =mytab.get_column('loc')
genome_dirweights = numpy.array(mytab.get_column('dir_weight'))
genome_invweights = numpy.array(mytab.get_column('inv_weight'))
genome_gapids = numpy.array(mytab.get_column('gap_id'))

new_locs = []
new_dirweights = []
new_invweights = []
new_gapids = []

# now, for each probe on the array, find the bp closer to it than anything else,
#  and merge their properties for output

#we check (and subsequently assume) that both the probes and genome are sorted
def is_sorted(lst): 
  return (sorted(lst) == lst)

if not is_sorted(list(genome_locs)):
  raise('Genomic locations in alignment summary file were not sorted!')

if not is_sorted(list(locs)):
  raise('Microarray probes were not sorted!')

# if everything looks ok, then find the set of genomic locations closest to each probe in turn

# do some special treatment for bases between the first and last probes
last_incl_inds = []
last_center = locs[-1]
first_center = locs[0]
genome_length = genome_locs[-1]

dist_last_first = first_center - (last_center - genome_length)
lastprobe_lastind = last_center + (dist_last_first/2)

# now go through and find the range of bp corresponding to each probe
thisprobe_start = last_center - genome_length + (dist_last_first/2)
thisprobe_ind = 0
thisprobe_loc = locs[0]

#print "starting to go through the array"
while (thisprobe_ind < (len(locs) - 1) ):
  nextprobe_loc = locs[thisprobe_ind + 1]
  interprobe_dist = nextprobe_loc - thisprobe_loc
  thisprobe_end = thisprobe_loc + (interprobe_dist/2)

  thisprobe_inds = range(thisprobe_start,thisprobe_end)
  #print "using probes %s for probe centered at %i" % (thisprobe_inds, thisprobe_loc)
  dirweight,invweight,gapid = merge_bp_data(thisprobe_inds, genome_dirweights,genome_invweights,genome_gapids)
  new_locs.append(thisprobe_loc)
  new_dirweights.append(dirweight)
  new_invweights.append(invweight)
  new_gapids.append(gapid)

  # now roll over to the next probe
  thisprobe_ind += 1
  thisprobe_loc = locs[thisprobe_ind]
  thisprobe_start = thisprobe_end

# handle the final probe appropriately
last_inds = range(thisprobe_start,lastprobe_lastind)
dirweight,invweight,gapid = merge_bp_data(last_inds, genome_dirweights,genome_invweights,genome_gapids)
#print "using probes %s for probe centered at %i" % (last_inds, thisprobe_loc)
new_locs.append(locs[-1])
new_dirweights.append(dirweight)
new_invweights.append(invweight)
new_gapids.append(gapid)


# now set up an appropriate table for output
newtab = table.Table()
newtab.add_column(new_locs,'loc')
newtab.add_column(new_dirweights,'dir_weight')
newtab.add_column(new_invweights,'inv_weight')
newtab.add_column(new_gapids,'gap_id')

table.write_table(outtab, newtab, sep=";")
