#!/usr/bin/python
# script to transform cluster tables generated by the "Clustering" toot (and aggregate + and -) in collection of SmRNAwidows Instances and perform operation on it
# version 1 - 19-11-2013
# Usage GetClustersInSmRNAwindows.py  <bowtie index> <bowtie input> <cluster table> <output>

import sys, subprocess
from collections import defaultdict # required for some SmRNAwindow attributes (readDic)
from numpy import mean, std # required for some SmRNAwindow methods
from smRtools import *

fasta_dic = get_fasta (sys.argv[1])
F = open (sys.argv[3], "r")
objDic = {}
for line in F:
  fields=line.split()
  gene = fields[0]+":"+fields[1]+"-"+fields[2]
  sequence = fasta_dic[fields[0]][int(fields[1])-1:int(fields[2])]
  objDic[gene] = SmRNAwindow(gene, sequence, fields[1]-1)
F.close()

#to be continued !! 19-11-2013
F = open (sys.argv[2], "r") # F is the bowtie output taken as input

for line in F:
  fields = line.split()
  polarity = fields[1]
  gene = fields[2]
  offset = int(fields[3])
  size = len (fields[4])
  try:
    objDic[gene].addread (polarity, offset, size)
  except KeyError:
    objDic[gene] = SmRNAwindow(gene, fasta_dic[gene])
    objDic[gene].addread (polarity, offset, size)
F.close()

OUT = open (sys.argv[2], "w")

print >> OUT, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % ("Transposon Id", "pi total", "pi strand+ ratio", "pi density", "zscore pi", "prob pi", "si total", "si strand+ ratio", "si density", "zscore si", "prob si", "pi/si ratio")
for transposon in objDic:
  pitotal = objDic[transposon].readcount(23,29)
  try:
    piratio = objDic[transposon].forwardreadcount(23,29) / float(pitotal)
  except ZeroDivisionError:
    piratio = "NA"
  pidensity = pitotal / float(objDic[transposon].size)                       
  pipi_z = objDic[transposon].z_signature(23,29,23,29, range(1,26) )[10]
  pipi_prob = objDic[transposon].hannon_signature(23,29,23,29, range(1,26) )[10] * 100

  sitotal = objDic[transposon].readcount(20,21)
  try:
    siratio = objDic[transposon].forwardreadcount(20,21) / float(sitotal)
  except ZeroDivisionError:
    siratio = "NA"
  sidensity = sitotal / float(objDic[transposon].size)
  sipi_z = objDic[transposon].z_signature(23,29,20,21, range(1,26) )[10]
  pisi_prob = objDic[transposon].hannon_signature(23,29,20,21, range(1,26) )[10] * 100
  try:
    pisi_ratio = float(pitotal)/(pitotal+sitotal)
  except ZeroDivisionError:
    pisi_ratio = "NA"
  print >> OUT, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (transposon, pitotal, piratio, pidensity, pipi_z, pipi_prob, sitotal, siratio, sidensity, sipi_z, pisi_prob, pisi_ratio )

OUT.close()
