#!/usr/bin/env python2.4
# 
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# Written (W) 2007 Gunnar Raetsch
# Written (W) 2006-2007 Soeren Sonnenburg
# Copyright (C) 2007 Fraunhofer Institute FIRST and Max-Planck-Society
# 
import os
import os.path
import sys
import pickle
import bz2
import numpy
import optparse

import genomic
import model
import seqdict
import shogun.Structure

from content_sensors import content_sensors
from signal_detectors import signal_detectors
from plif import plif

msplicer_version='v 0.1'

class msplicer:
	def __init__(self):
		self.model = None
		self.plif = None
		self.signal = None
		self.content = None

	def load_model(self, filename):
		sys.stderr.write('loading model file\n')
		f=None
		picklefile=filename+'.pickle'
		if os.path.isfile(picklefile):
			self.model=pickle.load(file(picklefile))	
		else:
			if filename.endswith('.bz2'):
				f=bz2.BZ2File(filename);
			else:
				f=file(filename);

			self.model=model.parse_file(f)
			f.close()

			f=file(picklefile,'w')
			pickle.dump(self.model, f)
			f.close()

		self.plif=plif(self.model)
		self.signal=signal_detectors(self.model)
		self.content=content_sensors(self.model)

	def compute_seqmatrix(self, seq):
		positions=[(0,0,0)]
		positions.append((seq.start,0,1))
		positions.append((seq.end,0,4))
		positions.append((len(seq.seq)-1,0,5))

		positions.extend(zip(seq.preds['donor'].positions,
			seq.preds['donor'].scores,
			len(seq.preds['donor'].positions)*[2]))

		positions.extend(zip(seq.preds['acceptor'].positions,
			list(seq.preds['acceptor'].scores),
			len(seq.preds['acceptor'].positions)*[3]))

		positions.sort(cmp=lambda x,y : int(x[0]-y[0]))

		seqmatrix=-numpy.infty * numpy.ones((6,len(positions)))
		for i in xrange(len(positions)):
			seqmatrix[positions[i][2],i]=positions[i][1]

		positions= numpy.array([ x[0] for x in positions ], numpy.int32)
		return (seqmatrix, positions)

	def initialize_dynprog(self, seq):
		dyn=shogun.Structure.DynProg()
		n=len(self.model.p)
		dyn.set_N(n)
		dyn.set_p_vector(self.model.p)
		dyn.set_q_vector(self.model.q)
		dyn.set_a_trans_matrix(self.model.a_trans)

		#design scoring seqmatrix
		(seqmatrix, positions) = self.compute_seqmatrix(seq)

		dyn.init_mod_words_array(numpy.array([[1,1,1,1,1,1,1,1],[0,0,0,0,0,0,0,0]],numpy.int32).transpose())
		assert(dyn.check_svm_arrays())
		dyn.best_path_set_seq(seqmatrix)
		dyn.best_path_set_pos(positions)
		dyn.best_path_set_orf_info(-numpy.ones((n,2),numpy.int32))
		dyn.best_path_set_plif_list(self.plif.get_plif_array())

		plifidmat=numpy.array([
					 [-1,-1,-1,-1,-1,-1],
					 [-1,-1, 2,-1, 5,-1],
					 [-1,-1,-1, 6,-1,-1],
					 [-1,-1, 4,-1, 3,-1],
					 [-1,-1,-1,-1,-1,-1],
					 [-1,-1,-1,-1,-1,-1]],numpy.int32).transpose()
		dyn.best_path_set_plif_id_matrix(plifidmat)
		plifstatemat=numpy.array([[-1,-1, 1, 0,-1,-1]],numpy.int32).transpose()
		dyn.best_path_set_plif_state_signal_matrix(plifstatemat)
		s=[]; s+=seq.seq;
		dyn.best_path_set_single_genestr(numpy.array(s))
		dyn.best_path_set_dict_weights(self.content.get_dict_weights())
		return (dyn,positions)

	def write_gff(self, outfile, pred, name, score, skipheader):
		descr=list()
		for i in xrange(pred.shape[0]):
			d=dict()
			d['seqname']=name
			d['source']='msplicer'
			d['feature']='exon'
			d['start']=pred[i,0]
			d['end']=pred[i,1]
			d['score']=score
			d['strand']='+'
			d['frame']=0
			descr.append(d)

		genomic.write_gff(outfile, ('msplicer',msplicer_version), ('DNA', name), descr, skipheader)

	def predict_file(self, fname, (start,end)):
		skipheader=False
		fasta_dict = genomic.read_fasta(file(fname))
		sys.stderr.write('found fasta file with ' + `len(fasta_dict)` + ' sequence(s)\n')
		seqs= seqdict.seqdict(fasta_dict, (start,end))

		#get donor/acceptor signal predictions for all sequences
		self.signal.predict_acceptor_sites_from_seqdict(seqs)
		self.signal.predict_donor_sites_from_seqdict(seqs)

		for seq in seqs:
			#initialize dynamic programming, with content sensors
			#signal detectors, Plifs and HMM like model
			(dyn,positions)=self.initialize_dynprog(seq)

			#compute max likely path
			dyn.best_path_call(1, False)
			scores=dyn.best_path_get_scores()
			states=dyn.best_path_get_states()
			pos=dyn.best_path_get_positions()
			pred=positions[pos[0][0:numpy.where(pos[0]==-1)[0]][1:-1]]
			pred=pred.reshape((len(pred)/2,2))
			self.write_gff(outfile, pred, seq.name, scores, skipheader)
			skipheader=True


def print_version():
	sys.stderr.write('mSplicer '+msplicer_version+'\n')

def parse_options():
	parser = optparse.OptionParser(usage="usage: %prog [options] seq.fa")
	
	parser.add_option("-o", "--outfile", type="str", default='stdout',
			                  help="File to write the results to")
	parser.add_option("-v", "--version", default=False,
			                  help="Show some more information")
	parser.add_option("--start", type="int", default=499,
			                  help="coding start (zero based, relative to sequence start)")
	parser.add_option("--stop", type="int", default=-499,
			                  help="""coding stop (zero based, if positive relative to
							  sequence start, if negative relative to sequence end)""")
	parser.add_option("--model", type="str", default='WS160',
			                  help="mSplicer Model to use in predicting")

	(options, args) = parser.parse_args()
	if options.version:
		print_version()
		sys.exit(0)

	if len(args) != 1:
		parser.error("incorrect number of arguments")

	fafname=args[0]
	if not os.path.isfile(fafname):
		parser.error("fasta file does not exist")

	if options.model.endswith('gc'):
		gc=1
		model=options.model[:-2]
	else:
		gc=0
		model=options.model

	modelfname = 'data/msplicer_elegans%s_gc=%d_orf=0.dat.bz2' % (model, gc)

	if not os.path.isfile(modelfname):
		parser.error("model should be one of  WS120, WS120gc, WS150, WS160, WS160gc")

	if options.outfile == 'stdout':
		outfile=sys.stdout
	else:
		try:
			outfile=file(options.outfile,'w')
		except IOError:
			parser.error("could not open %s for writing" % options.outfile)


	if options.stop<0:
		options.stop-=5;

	if options.start<80:
		parser.error("--start value must be >=80")

	if options.stop > 0 and options.start >= options.stop - 80:
		parser.error("--stop value must be > start + 80")

	if options.stop < 0 and options.stop > -80:
		parser.error("--stop value must be <= - 80")

	return ((options.start,options.stop), fafname, modelfname, outfile)


if __name__ == '__main__':
	dyn=shogun.Structure.DynProg()
	(startstop, fafname, modelfname, outfile ) = parse_options()
	p=msplicer()
	p.load_model(modelfname);
	p.predict_file(fafname, startstop)
