#!/usr/bin/etc python
from numpy.random import choice
from itertools import count,chain
from copy import deepcopy
import sys

###############
# Globals
###############
possible_states = ["progenitor","differentiated"]

#################
# Classes
#################
class Cell:
	def __init__(self,pDiff,state,parent_name=-1):
		self.parent_name = parent_name
		self.name = nameGenerator.next()
		#assert(0>pDiff>1)
		self.pDiff = float(pDiff)
		self.weights = [1-self.pDiff,self.pDiff]
		#assert (sum(self.weights)==1)
		self.cycles = 0
		#assert(state in self.possible_states)
		self.state = state
	def divide(self):
		if(self.state=="progenitor"):
			newState=choice(possible_states,p=self.weights)
			daughterState = choice(possible_states,p=self.weights)
			self.state = newState
			self.cycles = 1
			return([self,Cell(parent_name=self.name,pDiff=self.pDiff,state=newState)])
		else:
			return([self])
	def __str__(self):
		return("Cell%d:%s:%d" % (self.name,self.state,self.cycles))
	def __repr__(self):
		return(self.__str__())


nStart = 10
pDiff = 0.45 # probability that a cell division will produce a daughter cell that is 'post-mitotic' (e.g. differentiated)
nCycles = 30
nameGenerator = count(start=0,step=1)
cells = {0:[Cell(pDiff=pDiff,state="progenitor") for i in xrange(nStart)]} #initialize progenitor pool


# Simulate Cell cycle for nCycles
for i in xrange(1,nCycles+1):
	sys.stderr.write("%d" % i)
	cells[i] = [deepcopy(x) for x in cells[i-1]]
	cells[i] = [x.divide() for x in cells[i]]
	#print(cells[i])
	cells[i] = list(chain.from_iterable(cells[i]))

#Report Summary stats by cycle
print("%s: %s\t%s\t%s" % ("round", "total_cells", "fracProgenitors","fracDifferentiated"))
for k in cells.keys():
	nCells = len(cells[k])
	fracProgenitors = (len([i for i in cells[k] if i.state == "progenitor"])/float(nCells))*100
	fracDiff = 100-fracProgenitors
	print("%s: %d\t%0.2f\t%0.2f" % (k,nCells,fracProgenitors,fracDiff))