import numpy
import functools
import math

def fval(t,W):
   if callable(W):
      return W(t)
   return W

def L(t,t0,W):
   ft0=fval(t,t0)
   fW=fval(t,W)
   x=(t-ft0)/fW
   if x<-18:
      return 0
   if x>18:
      return 1
   return 1/(1+numpy.exp(-x))
def dL(t,t0,W):
   return L(t,t0,W)*(1-L(t,t0,W))

def dLdW(t,t0,W):
   fW=fval(t,W)
   ft0=fval(t,t0)
   return dL(t,t0,W)*(t-ft0)/fW/fW

def dLdt0(t,t0,W):
   fW=fval(t,W)
   return dL(t,t0,W)/fW

def logisticPar(par):
   t0=par['t0']['value']
   W=par['W']['value']
   f=functools.partial(L,t0=t0,W=W)
   
   #D1[x] could be a function !!!
   Dt0=par['t0']['derivatives']
   dfdt0=functools.partial(dLdt0,t0=t0,W=W)
   dt0=generate(dfdt0,Dt0)
   #D2[x] could be a function !!!
   DW=par['W']['derivatives']
   dfdW=functools.partial(dLdW,t0=t0,W=W)
   dW=generate(dfdW,DW)
   #merge
   return Object(f,[dt0,dW])

def deltaW(t,t1,t2):
   ft1=fval(t,t1)
   ft2=fval(t,t2)
   return ft2-ft1

def deltaT(t,t1):
   ft1=fval(t,t1)
   return t-t1

def fq(t,t1,t2):
   return deltaT(t,t1)/deltaW(t,t1,t2)

def f1(t,C,t1,t2,L1,L2):
   fC=fval(t,C)
   return fC*fq(t,t1,t2)*L1(t)*(1-L2(t))

def f2(t,C,L1,L2):
   fC=fval(t,C)
   return fC*L1(t)*L2(t)

def lg(t,C,t1,t2,eps,L1,L2):
   fC=fval(t,C)
   return eps*fC+f1(t,C,t1,t2,L1,L2)+f2(t,C,L1,L2)

def dlgdt1(t,C,t1,t2,L1,L2):
   fC=fval(t,C)
   return -fC*L1(t)*(1-L2(t))*(1+fq(t,t1,t2))/deltaW(t,t1,t2)

def dlgdt2(t,C,t1,t2,L1,L2):
   return -f1(t,C,t1,t2,L1,L2)/deltaW(t,t1,t2)

def dlgdL1(t,C,t1,t2,L1,L2):
   fC=fval(t,C)
   fL2=L2(t)
   return fC*(fq(t,t1,t2)*(1-fL2)+fL2)

def dlgdL2(t,C,t1,t2,L1,L2):
   fC=fval(t,C)
   return fC*L1(t)*(1-fq(t,t1,t2))

def dlgdC(t,C,t1,t2,L1,L2):
   fL2=L2(t)
   return L1(t)*(fq(t,t1,t2)*(1-fL2)+fL2)

def lgFixedSlope(t,alpha,t1,t2,eps,L1,L2):
   fA=fval(t,alpha)
   ft1=fval(t,t1)
   dW=deltaW(t,t1,t2)
   return fA*(dW*eps+L1(t)*((t-ft1)*(1-L2(t))+dW*L2(t)))

def dlgFixedSlopedt1(t,alpha,L1,L2):
   fA=fval(t,alpha)
   return -fA*L1(t)

def dlgFixedSlopedt2(t,alpha,L1,L2):
   fA=fval(t,alpha)
   return -fA*L1(t)*L2(t)

def dlgFixedSlopedL1(t,alpha,t1,t2,L1,L2):
   return lgFixedSlope(t,alpha,t1,t2,L1,L2)/L1(t)

def dlgFixedSlopedL2(t,alpha,t1,t2,L1,L2):
   fA=fval(t,alpha)
   ft2=fval(t,t2)
   return fA*L1(t)*(ft2-t)

def dlgFixedSlopedalpha(t,alpha,t1,t2,L1,L2):
   fA=fval(t,alpha)
   return lgFixedSlope(t,alpha,t1,t2,L1,L2)/fA

def linearGrowth(par):
   #make sure we never return zero or negative
   C=par['C']['value']
   t1=par['t1']['value']
   t2=par['t2']['value']
   eps=1e-8

   pL1=logisticPar({'t0':par['t1'],'W':par['W1']})
   pL2=logisticPar({'t0':par['t2'],'W':par['W2']})
   L1=pL1['value']
   L2=pL2['value']
   f=functools.partial(lg,C=C,t1=t1,t2=t2,eps=eps,L1=L1,L2=L2)

   Dt1=par['t1']['derivatives']
   dfdt1=functools.partial(dlgdt1,C=C,t1=t1,t2=t2,L1=L1,L2=L2)
   dt1=generate(dfdt1,Dt1)
   
   Dt2=par['t2']['derivatives']
   dfdt2=functools.partial(dlgdt2,C=C,t1=t1,t2=t2,L1=L1,L2=L2)
   dt2=generate(dfdt2,Dt2)

   DL1=pL1['derivatives']
   dfdL1=functools.partial(dlgdL1,C=C,t1=t1,t2=t2,L1=L1,L2=L2)
   dL1=generate(dfdL1,DL1)

   DL2=pL2['derivatives']
   dfdL2=functools.partial(dlgdL2,C=C,t1=t1,t2=t2,L1=L1,L2=L2)
   dL2=generate(dfdL2,DL2)

   DC=par['C']['derivatives']
   dfdC=functools.partial(dlgdC,C=C,t1=t1,t2=t2,L1=L1,L2=L2)
   dC=generate(dfdC,DC)

   return Object(f,[dt1,dt2,dL1,dL2,dC])

def linearGrowthFixedSlope(par):
   #make sure we never return zero or negative
   alpha=par['alpha']['value']
   t1=par['t1']['value']
   t2=par['t2']['value']
   eps=1e-8

   pL1=logisticPar({'t0':par['t1'],'W':par['W1']})
   pL2=logisticPar({'t0':par['t2'],'W':par['W2']})
   L1=pL1['value']
   L2=pL2['value']
   f=functools.partial(lgFixedSlope,alpha=alpha,t1=t1,t2=t2,eps=eps,L1=L1,L2=L2)

   Dt1=par['t1']['derivatives']
   dfdt1=functools.partial(dlgFixedSlopedt1,alpha=alpha,L1=L1,L2=L2)
   dt1=generate(dfdt2,Dt1)
   
   Dt2=par['t2']['derivatives']
   dfdt2=functools.partial(dlgFixedSlopedt2,alpha=alpha,L1=L1,L2=L2)
   dt2=generate(dfdt2,Dt2)

   DL1=pL1['derivatives']
   dfdL1=functools.partial(dlgFixedSlopedL1,alpha=alpha,t1=t1,t2=t2,L1=L1,L2=L2)
   dL1=generate(dfdL1,DL1)

   DL2=pL2['derivatives']
   dfdL2=functools.partial(dlgFixedSlopedL2,alpha=alpha,t1=t1,t2=t2,L1=L1,L2=L2)
   dL2=generate(dfdL2,DL2)

   DA=par['alpha']['derivatives']
   dfdA=functools.partial(dlgFixedSlopedalpha,alpha=alpha,t1=t1,t2=t2,L1=L1,L2=L2)
   dA=generate(dfdA,DA)

   return Object(f,[dt1,dt2,dL1,dL2,dA])

def feps(t,A,tau):
   fA=fval(t,A)
   fTau=fval(t,tau)
   return fA*math.exp(-t/fTau)

def dfepsdA(t,A,tau):
   fTau=fval(t,tau)
   return math.exp(-t/fTau)

def dfepsdTau(t,A,tau):
   fTau=fval(t,tau)
   fA=fval(t,A)
   if fTau==0:
      return 0
   return fA*math.exp(-t/fTau)*(t/fTau/fTau)


def exp(par):
   A=par['A']['value']
   tau=par['tau']['value']
   f=functools.partial(feps,A=A,tau=tau)

   DA=par['A']['derivatives']
   dfdA=functools.partial(dfepsdA,A=A,tau=tau)
   dA=generate(dfdA,DA)

   DTau=par['tau']['derivatives']
   dfdTau=functools.partial(dfepsdTau,A=A,tau=tau)
   dTau=generate(dfdTau,DTau)
   
   return Object(f,[dA,dTau])

def Object(v,dArray):
   o={'value':v,'function':True}
   allKeys=[]
   for x in dArray:
      allKeys.extend(x.keys()) 
   allKeys=list(set(allKeys))
   o['derivatives']={x:sumDict(dArray,x) for x in allKeys}
   return o

def derivedObject(f,ders):
   o={'value':f}
   DD=[]
   for d in ders:
      df=d['df']
      D=d['D']
      DD.append({x:df*D[x] for x in D})
   allKeys=[]
   for x in DD:
      allKeys.extend(x.keys()) 
   allKeys=list(set(allKeys))
   o['derivatives']={x:sumValues(DD,x) for x in allKeys}
   return o
 

def sumDict(dArray,x):
   #print('sumFunctions for {}'.format(x))
   fs=[]
   for a in dArray:
      #print('\tChecking {}'.format(a))
      try:
         #check if a[x] is a function with a dummy parameter
         fs.append(a[x])
         #print('\tAdding [{}]'.format(v))
      except KeyError:
         #print('\t{} not in table'.format(x))
         continue
   return sumArray(fs)

def sumArray(fs,w=1):
   return lambda t,fs=fs,w=w: w*sum([to(f)(t) for f in fs])

def sumValues(dArray,x):
   s=0
   for a in dArray:
      try:
         s+=a[x]
      except KeyError:
         continue
   return s

def isFunction(vObject):
   if callable(vObject['value']):
      return True
   return False

def to(a):
   if callable(a):
      return a
   return lambda t,a=a:a
 
def contains(arr):
   for a in arr:
      if callable(a):
         return True
   return False

def multiply(t,a,b):
   #abstract actual type of a and b
   fa=fval(t,a)
   fb=fval(t,b)
   return fa*fb

def generate(df,D):
   return {x:functools.partial(multiply,a=D[x],b=df) for x in D}

def product(pA,pB):
   #return a product of two values with derivatives
   #print('{}*{}'.format(d['a'],d['b']))
   a=pA['value']
   DA=pA['derivatives']
   b=pB['value']
   DB=pB['derivatives']
   if any(['function' in pA,'function' in pB]):
      fa=to(a)
      fb=to(b)
      f=lambda t,a=fa,b=fb,:a(t)*b(t)
      dfdA=lambda t,b=fb: b(t)
      dfdB=lambda t,a=fa: a(t)
      dA=generate(dfdA,DA)
      dB=generate(dfdB,DB)
      return Object(f,[dA,dB])
   else:
      return derivedObject(a*b,[{'df':b,'D':DA},{'df':a,'D':DB}])

def power(pA,pN):
   #return pA to the power of pN, numpy.power(pA,pN), with derivatives
   a=pA['value']
   DA=pA['derivatives']
   n=pN['value']
   DN=pN['derivatives']
   if any(['function' in pA,'function' in pN]):
      fa=to(a)
      fn=to(n)
      f=lambda t,a=fa,n=fn:numpy.power(a(t),n(t))
      dfdA=lambda t,n=fn,f=f,a=fa:n(t)*f(t)/a(t)
      dfdN=lambda t,a=fa,f=f:numpy.log(a(t))*f(t)
      dA=generate(dfdA,DA)
      dN=generate(dfdN,DN)
      return Object(f,[dA,dN])
   else:
      f=numpy.power(a,n)
      return derivedObject(f,[{'df':n*f/a,'D':DA},{'df':f*numpy.log(a),'D':DN}])

def ratio(pA,pB):
   #return ratio pA/pB with all derivatives
   #print('{}/{}'.format(d['a'],d['b']))
   a=pA['value']
   DA=pA['derivatives']
   b=pB['value']
   DB=pB['derivatives']
   if any(['function' in pA,'function' in pB]):
      fa=to(a)
      fb=to(b)
      f=lambda t,a=fa,b=fb,:a(t)/b(t)
      dfdA=lambda t,f=f,a=fa: f(t)/a(t)
      dfdB=lambda t,f=f,b=fb: -f(t)/b(t)
      dA=generate(dfdA,DA)
      dB=generate(dfdB,DB)
      return Object(f,[dA,dB])
   else:
      return derivedObject(a/b,[{'df':1/b,'D':DA},{'df':-a/b/b,'D':DB}])

def add(pA,pB):
   #return sum of pA and pB with derivatives
#print('{}+{}'.format(d['a'],d['b']))
   a=pA['value']
   DA=pA['derivatives']
   b=pB['value']
   DB=pB['derivatives']
#even more generic -> df/dp=[df/dA*dA/dp+df/dB*dB/dp]
   if any(['function' in pA,'function' in pB]):
      fa=to(a)
      fb=to(b)
      f=lambda t,a=fa,b=fb,:a(t)+b(t)
      dfdA=lambda t: 1
      dfdB=lambda t: 1
      dA=generate(dfdA,DA)
      dB=generate(dfdB,DB)
      return Object(f,[dA,dB])
   else:
      return derivedObject(a+b,[{'df':1,'D':DA},{'df':1,'D':DB}])