###############################################################
# GDL, at University of Leeds and at IISc Bangalore, 2016-2020
# T cells in vitro. Time is measured in hours, distance in microns.
# Each cell has (x,y), radius and neighbourhood, local IL2 level.
# Cell death and division are included.
# Configuration at each hour stored as a .png file.
################################################################
from __future__ import print_function  # use python3 notation print()
from random import random
from matplotlib import use
use('Agg')
from datetime import date
import pylab as pl

# parameter values. 
# suggestion: first, run code with n0=1000 as a test run (1-2 minutes).
#             then, increase n0
n0 = 20000  # initial number of cells
mu = 0.008   # death rate
alph = 1.0  # signalling rate
gamma = 0.2  # TCR activation rate 0 < gamma < 1
tmax = 48.0  # Time is measured in hours
dt = 0.01  

class Cell(object):
    ''' store number'''
    number = 0
    def __init__(self):
        type(self).number += 1
        self.ct = ''
    def delete(self):
        type(self).number -= 1

class T(Cell):
   ''' T cell class. Attributes are:
   act activation status
   il2p IL-2 producing status
   cy cycling
   ge generation
   mt mitosis time
   b cell radius'''
   def __init__(self):
       Cell.__init__(self)
       self.cy = False
       self.ge = 0
       self.ct = 'T-cell'
   def __str__(self):
       return "%s cy=%s la=%f"%(self.ct,self.cy)
   def entercycle(self):
       '''cell.mt = time remaining until division'''
       self.cy = True
       self.mt = 24.0 if self.ge==0 else 12.0
   def activate(self):
       self.act = True
       self.cd69 *= 20.0
       self.r[-1] += r1*random()
   def advance(self):
       ''' reset attributes just before division '''
       self.ge += 1
       self.cy,self.act,self.il2p = False,False,False
       self.r[-1] = r0-0.5+random()
   def celldivision(self):
       ''' daughter cells close to mother'''
       self.advance()
       newcell = CD4()
       type(self).number += 1
       newcell.r[0],newcell.r[1] = self.r[0]+3*r1*random(),self.r[1]+3*r1*random()
       newcell.r[-1] = r0-0.5+random()
       newcell.vic,newcell.ge,newcell.IL2 = self.vic,self.ge,self.IL2
       return newcell
   def cellsurvive(self,p):
       if p > random():
           type(self).number -= 1
           cdfours.deadcelllist.append(self)
           return False
       return True
   def cellstep(self,dt):
       ''' Cycling cells count down; they may divide (return True).
           Activated cells may enter cycle or become IL2 producing.
           Noncycling cells can get an IL2 hit '''
       if self.cy:
           self.mt -= dt
           if self.mt < 0:
               return True
       urv = random()
       if urv < alph*self.IL2*dt :
           if urv < gamma*alph*self.IL2*dt: # TCR signal
               if self.act:
                   self.entercycle()
               else:
                   self.activate()
           else:   # IL-2 signal
               self.cd25 = pl.sqrt(cd25max*self.cd25*random())
               self.cd62l = pl.sqrt(cd62lmin*self.cd62l*random())
               if self.il2p:
                   self.il2p = False
               elif self.act:
                   self.il2p = True                   
       return False

class CD4(T):
    ''' CD4 T cell class '''
    def __init__(self):
        T.__init__(self)
        self.ct = 'CD4'
        self.cd69 = pl.exp(pl.normal(1.5,1.0))
        self.r = [0,0,0]        # x,y and radius     
        # Initially, ten percent of CD4 cells are CD25+CD62L+
        if random() < 0.10:
            self.cd25 = pl.exp(pl.normal(4.0,2.0))
            self.cd62l = pl.exp(pl.normal(1.5,0.5))
        else:
            self.cd25 = pl.exp(pl.normal(2.0,1.0))
            self.cd62l = pl.exp(pl.normal(5,1.0))
        # Initially, one percent of CD4 cells are activated
        # and half of those are IL2-producing
        self.il2p,self.act = False,False
        urv = random()
        if urv < 0.01:
            self.act = True
            self.cd69 *= 20
            if urv < 0.005:
                self.il2p = True

celltype={'CD4':CD4}

class CellPopulation(object):
    def __init__(self,ncells,ctype):
        self.ncells = ncells
        self.celllist = [celltype[ctype]() for i in pl.arange(ncells)] 
        self.deadcelllist = []

    def placecellsinwell(self):
        '''place cells with no overlap'''
        print('\n###placing '+str(n0)+' cells with sd='+str(sdxy),'###')
        maxn,i = 0,0
        for cell in self.celllist:
            b = r0-0.5+random()
            if cell.act:
                b += r1*random()
            i += 1
            n,placed = 0,False
            while not placed:
                n += 1
                x,y = pl.normal(0,sdxy),pl.normal(0,sdxy)
                if min([mydist(othercell.r,[x,y]) for othercell in self.celllist]) > 4*b*b:
                    placed=True
            # print some running information on screen
            for thisn in (10,100,200,400,1000):
                if n>=thisn and maxn<thisn:
                    print(i,'cells placed',n,'attempts')
            maxn = max(n,maxn)
            if i in (100,1000,n0/2,10000,3*n0/4,9*n0/10,n0):
                print(i,'cells placed',n,'attempts, max',maxn)
            cell.r = [x,y,b]  # position x,positiony,radius

    def testics(self):
        ''' deactivate all cells except for one at centre of well'''
        for cell in self.celllist:
            cell.il2p,cell.act = False,False
        distlist = [mydist([0,0],cell.r) for cell in self.celllist]
        ii = distlist.index(min(distlist))
        self.celllist[ii].il2p = True
        print('central cell has',len(self.celllist[ii].vic),'neighbours')

    def oneatcentre(self):
        ''' put, by hand, one IL-2 producing cell at centre of well'''
        distlist = [mydist([0,0],cell.r) for cell in self.celllist]
        ii = distlist.index(min(distlist))
        self.celllist[ii].il2p,self.celllist[ii].act = True,True
        print('central cell has',len(self.celllist[ii].vic),'neighbours')

    def plotcellsinwell(self,t):
        '''make diagram and two zoom insets'''
        spatialfig = pl.figure(figsize=(8,8))
        well = pl.axes()
        for cell in self.deadcelllist:
            well.add_patch(pl.Circle((cell.r[0],cell.r[1]),radius=cell.r[2],
                                     alpha=0.4,facecolor='0.5',edgecolor='k'))
        for cell in self.celllist:
            myc = 'r' if cell.il2p else 'b'                
            circle = pl.Circle((cell.r[0],cell.r[1]),radius=cell.r[2],
                               alpha=0.6,facecolor=myc,edgecolor=myc)
            well.add_patch(circle)
        well.axis('scaled',aspect='equal')
        [well.spines[s].set_visible(False) for s in ("top","right","bottom")]
        well.get_yaxis().tick_left()    
        well.plot([3*sdxy,3*sdxy+250],[-3.5*sdxy,-3.5*sdxy],'k',lw=4)
        well.text(3*sdxy+10,-3.75*sdxy,'250$\mu$m')
        well.text(-3.5*sdxy,3.5*sdxy,str(int(t+0.01))+'h')
        pl.xticks([])
        pl.yticks([-500,0,500],['-500$\mu$m','0','500$\mu$m'],fontsize=12)
        pl.xlim([-4*sdxy,4*sdxy])
        pl.ylim([-4*sdxy,4*sdxy])
        # centre
        cwell = pl.axes([0.77,0.42,0.16,0.16])
        for cell in self.deadcelllist:
            circle = pl.Circle((cell.r[0],cell.r[1]),radius=cell.r[2],
                               alpha=0.4,facecolor='0.5',edgecolor='k')
            cwell.add_patch(circle)
        for cell in self.celllist:
            if abs(cell.r[0])<50 and abs(cell.r[1])<50:
                myc = 'r' if cell.il2p else 'b'                
                circle = pl.Circle((cell.r[0],cell.r[1]),radius=cell.r[2],facecolor=myc,edgecolor=myc)
                cwell.add_patch(circle)
        cwell.axis('scaled')
        pl.xticks([])
        pl.xlim([-50,50])
        pl.ylim([-50,50])
        cwell.get_yaxis().tick_right()    
        pl.yticks([-50,0,50],['-50$\mu$m','0','50$\mu$m'],fontsize=10)
        # nearer edge
        ewell = pl.axes([0.77,0.70,0.16,0.16])
        for cell in self.deadcelllist:
            circle = pl.Circle((cell.r[0],cell.r[1]),radius=cell.r[2],
                               alpha=0.4,facecolor='0.5',edgecolor='k')
            ewell.add_patch(circle)
        for cell in self.celllist:
            if abs(cell.r[0]-350)<50 and abs(cell.r[1]-350)<50:
                myc = 'r' if cell.il2p else 'b'                
                circle = pl.Circle((cell.r[0],cell.r[1]),radius=cell.r[2],facecolor=myc,edgecolor=myc)
                ewell.add_patch(circle)
        ewell.axis('scaled')
        pl.xticks([])
        pl.xlim([300,400])
        pl.ylim([300,400])
        ewell.get_yaxis().tick_right()    
        pl.yticks([300,400],['300$\mu$m','400$\mu$m'],fontsize=10)
        mydate = date.today().strftime("%d") + date.today().strftime("%B")[:3]
        pl.savefig('cellsinwell'+str(int(n0/1000))+mydate+'h'+str(int(t+0.01)).zfill(2)+'.png')
        pl.close(spatialfig)

    def neighbourhood(self):
        ''' cell.vic is a list of all cells closer than ril2 to it'''
        rs = ril2*ril2
        for cell in self.celllist:
            cell.vic = [c for c in self.celllist if mydist(c.r,cell.r) < rs]
        nlist = [len(cell.vic) for cell in self.celllist]
        print('number of neighbours',min(nlist),'..',max(nlist),'average',pl.mean(nlist))

    def localil2(self):
        ''' find the local IL-2 concentration for each cell'''
        for cell in self.celllist:
            n3 = len([c for c in cell.vic if c.il2p])
            n6 = len([c for c in cell.vic if c.cd25 > cd25thr])
            cell.IL2 = n3/(n6+1.0)

    def step(self,dt):
        ''' in a time interval dt, 
        some cells die, some cells enter cycle and others complete it'''
        self.celllist = [cell for cell in self.celllist if cell.cellsurvive(mu*dt)]
        self.celllist += [cell.celldivision() for cell in self.celllist if cell.cellstep(dt)]

    def printstats(self,t):
        '''extract numbers of cells by type, mean and max IL2 attribute'''
        n1 = len(self.celllist)
        n2 = len([cell for cell in self.celllist if cell.act])
        n3 = len([cell for cell in self.celllist if cell.il2p])
        n4 = len([cell for cell in self.celllist if cell.cy])
        n5 = len([cell for cell in self.celllist if cell.ge != 0])
        n6 = len([cell for cell in self.celllist if cell.cd25 > cd25thr])
        n7 = len(self.deadcelllist)
        tlist.append(t)
        ntotcells.append(n1)
        nactcells.append(n2)
        nil2cells.append(n3)
        il2list = [cell.IL2 for cell in self.celllist]
        print("%3s %6i %4i %4i %4i %4i %5i %5i %.4f %.3f"%(str(int(t+dt+0.01))+'h',
               n1,n2,n3,n4,n5,n6,n7,pl.mean(il2list),max(il2list)))

    def plottimecourse(self):
        '''shows unactivated, activated and IL-2 producing cell numbers'''
        fig = pl.figure(figsize=(8,8))
        params = {'font.family': 'serif','legend.fontsize': 'small','legend.numpoints':1}
        pl.rcParams.update(params)
        myax = fig.add_axes([0.10,0.55,0.50,0.40])  
        pl.plot(tlist,ntotcells,label='all')
        pl.plot(tlist,nactcells,label='activated')
        pl.plot(tlist,nil2cells,label='IL-2 producing')
        pl.xlabel('$t$ / hours')
        pl.xticks([12,24,36,48])
        pl.ylabel('number of cells')
        pl.ylim([0,1.1*n0])
        pl.legend(ncol=3)
        [myax.spines[s].set_visible(False) for s in ("top","right")]
        myax.get_xaxis().tick_bottom()    
        myax.get_yaxis().tick_left()    
        pl.savefig('TimeCourse'+mydate+str(int(ril2/r0)).zfill(2)+'h'+str(n0)+'.png')

def mydist(r1,r2):
    ''' returns square of distance '''
    return (r1[0]-r2[0])**2+(r1[1]-r2[1])**2

def makehist(mybins,mylist):
    myhist = []
    for i,x in enumerate(mybins[:-1]):
        y = mybins[i+1]
        myhist.append(len([t for t in mylist if x<t and t<y]))
    return myhist

# parameters (radii and thresholds)
r0,r1 = 4,3 # radius of a resting cell, max expansion of activated cell
ril2 = 7*r0 # IL-2 neighbourhood radius (distances in microns)
cd25max,cd25thr,cd62lmin = 1000,100,10

# a population of CD4 T cells, initially n0 cells
cdfours = CellPopulation(n0,'CD4') 
# place the cells and compile lists of neighbours of each cell
sdxy = 0.45*r0*pl.sqrt(n0)
cdfours.placecellsinwell()
cdfours.neighbourhood()
cdfours.oneatcentre()
cdfours.plotcellsinwell(0)

# run the experiment, make .png files
print('r0,r1,ril2=',str(r0),str(r1),str(ril2))
print('################################################')
print('n1..n7 = live, activated, IL-2p, cycling, divided, CD25+, dead cells')
print('\nStart in silico run.   n1..n7,          meanIL2, maxIL2')
t,nfig = 0,0
tlist,ntotcells,nactcells,nil2cells = [],[],[],[]
while t < tmax:
    cdfours.localil2()
    cdfours.step(dt)
    if (t+dt-int(t+dt+dt/10)) < dt/2:
        mydate = date.today().strftime("%d") + date.today().strftime("%B")[:3]
        # hourly figure showing cell positions, with IL-2 producing cells in red
        cdfours.plotcellsinwell(t+dt) 
        cdfours.printstats(t+dt)
    t += dt
cdfours.plottimecourse()