#!/usr/bin/env python 
"""
This script generates simulated dRNA-Seq reads (in .gff format) from known gene annotations. It is adapted from the gensimreads.py script that is part of
the RNASeqReadSimulator (https://github.com/davidliwei/RNASeqReadSimulator) originally developed by Wei Li (li.david.wei AT gmail.com).

USAGE 

  gensimreads_dRNAseq.py {OPTIONS} <GFF-File|->

PARAMETER

  GFF-File\tThe gene annotation file (in GFF format). Use '-' for STDIN input

OPTIONS

  -e/--expression [expression level file] \tSpecify the weight of each transcript. Each line in the file should have at least (NFIELD+1)  fields, with field 0 the annotation id, and field NFIELD the weight of this annoation. If this file is not provided, uniform weight is applied. 

  -n/--nreads readcnt \tSpecify the number of reads to be generated. Default 100000.

  -l/--readlen [read length] \tSpecify the read length. Default 32.

  -o/--output [output .bed file] \tSpecify the output file. Default STDOUT 

  -f/--field [NFIELD] \tThe field of each line as weight input. Default 7 (beginning from field 0) to compatible to genexplvprofile.py.

  --stranded \tThe reads are strand specific.
  
  --primary_enriched [enrichment_factor 0-100%]\tSimulate primary enrichement (differential RNA-seq), e.g. 50% of the reads are primary (5'-PPP)

NOTE 

  	1. The GFF file is required to sort according to the chromosome name and position. In Unix systems, use "sort -k 1,1 -k 4,4n in.gff > out.gff" to get a sorted version (out.BED) of the bed file (in.BED).  

  	2. Cannot handle reads spanning multiple exons. 

HISTORY
	03/03/2014
	  Support generating primary enriched libraries
	  Reads will be poisson distributed
	  Removed possibility for pairedend data
	  Removed possibility to provide positional bias file

AUTHOR
	Bjoern Voss
"""

from __future__ import print_function
import sys;
import subprocess;
import pydoc;
import os;
import random;
import bisect;
import math;
import numpy as np;
from getSegs import *;

import pdb;

# read length
readlen=32;
# number of reads to sample
readcnt=100000;

nfield=7;

if len(sys.argv)<2:
  print(pydoc.render_doc(sys.modules[__name__]));
  sys.exit();

allids={};
allidl=[];
allexp=[];

posweight=[];

#onbedfile=sys.argv[-1]+'.reads.bed';
onbedfile="-";

stranded=False;
primary_enriched=False;
enrichment_fac=0;

for i in range(len(sys.argv)):
  if i<len(sys.argv)-1:
    if sys.argv[i]=='-e' or sys.argv[i]=='--expression':
      # parse the annoatation file, and sum up the weights
      nline=0;
      totalweight=0;
      print('Reading annotation file...',file=sys.stderr);
      for lines in open(sys.argv[i+1]):
        nline=nline+1;
        if lines[0]=='#':
          continue;
        fields=lines.strip().split();
        if len(fields)<nfield+1:
          print('Error: the annotation file should include at least '+str(nfield+1)+' fields.',file=sys.stderr);
          sys.exit();
        allids[fields[0]]=0;
        totalweight+=float(fields[nfield]);
        allexp.append(totalweight);
        allidl.append(fields[0]);
      print('Read %d lines of the annotation' % nline,file=sys.stderr);
      #print('Total weight: %f' % sum(totalweight));
    if sys.argv[i]=='-n' or sys.argv[i]=='--nreads':
      readcnt=int(sys.argv[i+1]);
      print('Read count:',readcnt,file=sys.stderr);
    if sys.argv[i]=='-l' or sys.argv[i]=='--readlen':
      readlen=int(sys.argv[i+1]);
      print('Read length:',readlen,file=sys.stderr);
    if sys.argv[i]=='-o' or sys.argv[i]=='--output':
      onbedfile=sys.argv[i+1];
      print('Output bed file:',onbedfile,file=sys.stderr);
    if sys.argv[i]=='-f' or sys.argv[i]=='--field':
      nfield=int(sys.argv[i+1]);
      print('Field:',nfield,file=sys.stderr);
      print('Generate paired-end reads with mean and std '+str(pemean)+','+str(pestd),file=sys.stderr);
    if sys.argv[i]=='-h' or sys.argv[i]=='--help':
      print(pydoc.render_doc(sys.modules[__name__]));
      sys.exit();
    if sys.argv[i]=='--stranded':
      stranded=True;
    if sys.argv[i]=='--primary_enriched':
      primary_enriched=True;
      enrichment_fac=int(sys.argv[i+1]);

gfffile=sys.argv[-1];

# if no annotation file is specified, use uniform distri.
print('Assigning weights...',file=sys.stderr);
if len(allexp)==0:
  totalweight=0;
  for lines in open(gfffile):
    bedfield=lines.strip().split();
    attribs = {};
    for x in bedfield[8].split(';'):
      templ = x.split('=');
      attribs[templ[0]]=templ[1];
    allids[attribs['ID']]=0;
    totalweight+=1;
    allexp.append(totalweight);
    allidl.append(attribs['ID']);

# sampling process
print('Sampling...',file=sys.stderr);
for j in range(readcnt):
  k=random.random()*totalweight;
  sel=bisect.bisect_right(allexp,k);
  allids[allidl[sel]]+=1;

print('Total assigned reads:',sum(allids.values()),file=sys.stderr);

if onbedfile!="-":
  onfid=open(onbedfile,'w');
else:
  onfid=sys.stdout;


nlines=0;

totalgenreads=0;
# read bed file
for lines in open(gfffile):
  # update line counter
  nlines=nlines+1;
  if nlines %10000==1:
    print('Processing '+str(nlines)+' lines...',file=sys.stderr);
  # parse lines
  bedfield=lines.strip().split();
  if lines.startswith('#'):
    continue;
  attribs = {};
  for x in bedfield[8].split(';'):
    templ = x.split('=');
    attribs[templ[0]]=templ[1];
    
  if bedfield[6]=='+':
    direction=1;
  elif bedfield[6]=='-':
    direction=-1;
  else:
    print('Error: incorrect field in field[5] %s:' %bedfield[6],file=sys.stderr);
  if attribs['ID'] not in allids:
    # the current id not found, continue
    continue;
  nreads=allids[attribs['ID']];
  if nreads<1:
    continue;
  # parse all segments
  fieldrange=(int(bedfield[3]),int(bedfield[4]));
  exonlen=[int(bedfield[4])-int(bedfield[3])+1];
  exonstart=[int(bedfield[3])];
  totallen=sum(exonlen);
  # here, we randomly choose one position
  selrange=totallen-readlen+1;
  if selrange<1:
    print('Ignore annotation',attribs['ID'],'of length',totallen,'Reads:',allids[attribs['ID']],file=sys.stderr);
    continue;
  totalgenreads+=nreads;
  cumlen=[];cumlen.extend(exonlen);
  for i in range(1,len(cumlen)):
    cumlen[i]=cumlen[i]+cumlen[i-1];
  #
  # Poisson sampling with numpy:
  #
  distrlist=[];
  if primary_enriched==True:
    tss_reads = (enrichment_fac*nreads)/100;
    if tss_reads >= nreads:
      continue;
    distr = np.random.poisson((nreads-tss_reads)/(selrange-1), selrange-1);
    for t in range(0,int(tss_reads)):
      distrlist.append(0);
  else:
    distr = np.random.poisson(nreads/selrange, selrange);
  
  for t in range(0,len(distr)-1):
    for z in range(1,distr[t]):
      distrlist.append(t);

  for tpos in distrlist:
    if direction==-1:
        tpos=selrange-1-tpos;
    pos=tpos2pos(tpos,cumlen,exonstart);
    if True:
      (startrange,lenrange,status)=getSegs(pos,readlen,1,exonstart,exonlen);
      if status!=0:
        print('Status:',status,', pos:', pos,'out of',len(cumlen),len(distr),nreads,'strand',direction,file=sys.stderr);
        continue;
      lineid="%s_e_%d_%s_%d" % (attribs['ID'],t,bedfield[0],pos);
      # random direction
      if stranded==False or direction==0:
        thisdir=random.choice([1,-1]);
      else:
        thisdir=direction;
      writeBedline(onfid,lineid,bedfield[0],thisdir,startrange,lenrange);
    else:
      print(bedfield[0],file=sys.stdout);

print('Total '+str(nlines)+' lines...',file=sys.stderr);
print('Total '+str(totalgenreads)+' reads...',file=sys.stderr);
if onbedfile!="-":
  onfid.close();

