from thunder import RegressionModel, PCA, Colorize
from numpy import amax

%matplotlib inline

import matplotlib.pyplot as plt
import seaborn as sns
sns.set_context('notebook')

data, params = tsc.loadExampleEC2('zebrafish-optomotor-response')

img_mean = data.seriesMean().pack()

img_mean.shape

sns.set_context("poster")
sns.set_style('white')
plt.imshow(amax(img_mean,2), cmap='gray', clim=(0,1500));

img_single = data.select(100).pack()

sns.set_context("poster")
sns.set_style('white')
plt.imshow(img_single[:,:,10], cmap='gray', clim=(0,1500));

y = data.mean()

sns.set_context('notebook')
sns.set_style('darkgrid')
plt.plot(y);

y = data.detrend().normalize().mean()

sns.set_context('notebook')
sns.set_style('darkgrid')
plt.plot(y);

model = RegressionModel.load(params, 'linear')

data_filtered = data.detrend().normalize()
params = model.fit(data_filtered)
params.index

betas = params.select('betas')
betas.cache();

pca = PCA(k=3).fit(betas)

sns.set_context("notebook")
sns.set_palette("husl", 3)
plt.plot(pca.comps.T);

imgs = pca.scores.pack()

imgs.shape

maps = Colorize("polar", scale=1000).images(imgs[0:2])

sns.set_context("poster")
sns.set_style("white")
plt.imshow(maps[:,:,2,:]);

sns.set_context("poster")
sns.set_style("white")
plt.imshow(maps[:,:,10,:]);

sns.set_context("poster")
sns.set_style("white")
plt.imshow(amax(maps, 2));

pts = pca.scores.subset(nsamples=500, thresh=0.0003)

sns.set_style("darkgrid")
sns.set_context("notebook")
clrs = Colorize("polar", scale=1000).points(pts[:,0:2])
a = plt.scatter(pts[:, 0], pts[:, 1], s=150, c=clrs, alpha=0.6, edgecolor='black', linewidth=0.2)
xlim([-0.003, 0.002]); ylim([-0.004, 0.004]);

recon = map(lambda x: (x[0] * pca.comps[0, :] + x[1] * pca.comps[1, :]).tolist(), pts)
plt.gca().set_color_cycle(clrs)
plt.plot(asarray(recon).T);
plt.ylim(-0.0007,0.0007);