from pylab import *
import numpy as np
import pandas as pd
import dill
import copy
dill.settings['recurse'] = True
rng = np.random.default_rng(1)

def rate_cov(coef1, cov1, coef2, cov2):
    return exp(-5*cov1**2-5*cov2**2+2.)

def rate_base(t):
    return 2*t**1.5

n_sub = 100000
n_int = [1,5,10]
coef1 = 1.5
coef2 = 3.5

t = linspace(0,1,10**4)
dt = t[1]-t[0]

a1, w1, b1 = rng.uniform(0,1,n_sub), rng.uniform(5,10,n_sub), rng.uniform(0,1,n_sub)
cov1 = a1[:,newaxis]*cos(2*pi*w1[:,newaxis]*t + pi*b1[:,newaxis])
a2, w2, b2 = rng.uniform(0,1,n_sub), rng.uniform(20,30,n_sub), rng.uniform(0,1,n_sub)
cov2 = a2[:,newaxis]*cos(2*pi*w2[:,newaxis]*t + pi*b2[:,newaxis])

cov_func = {i:[[a1[i],w1[i],b1[i]],[a2[i],w2[i],b2[i]]] for i in range(n_sub)}

r = rate_cov(coef1,cov1,coef2,cov2) * rate_base(t)
spk = r*dt > rng.random(r.shape)

for n_i in n_int:
    print(n_i)
    colname = ['id','t0','t1','event','cov1','cov2']
    df = pd.DataFrame(zeros((0,6)),columns=colname)
    i = 0
    for id in range(n_sub):
        idx = where(spk[id]==True)
        event = 0 if len(idx[0])==0 else 1
        t1 = t[-1] if len(idx[0])==0 else t[idx][0]
        interval = linspace(0,t1,n_i+1)
        z = [ [id]*n_i,interval[:-1],interval[1:],[0]*(n_i-1)+[event],
              a1[id]*cos(2*pi*w1[id]*interval[1:]+pi*b1[id]),
              a2[id]*cos(2*pi*w2[id]*interval[1:]+pi*b2[id]) ]
        df = pd.concat([df,pd.DataFrame(array(z).T,columns=colname)],axis=0,ignore_index=True)
    
    data = {'df':df,'func':cov_func}
    dill.dump(data, open('data/nonlinear/data_nsub'+str(n_sub)+'nint'+str(n_i)+'.dill','wb'))

print(df.groupby('event').count())

