%% load data and set some parameters

datafile = '20080516_R1.mat'; % data file name
load([pwd filesep 'Data' filesep datafile]); % load data
load([pwd filesep 'stimulus.mat']);
trial = 1; % trial number

w = 0.6; % window length in second(s)
n = round(w/stimulus(trial).frame); % window length in frames

Nx = stimulus(trial).param.x/stimulus(trial).param.dx; % number of independent stimulus regions in horizontal direction 
Ny = stimulus(trial).param.y/stimulus(trial).param.dy; % number of independent stimulus regions in vertical direction
m = Nx*Ny; %number of stimulus dimensions

time = (1:stimulus(trial).Nframes-1)'*stimulus(trial).frame+stimulus(trial).onset; % stimulus timings
T = length(time)-1; % number of stimulus frames

%% first, to save time, use the STA to approximate the receptive field of the neuron
cellnum = 3;
observedSpikes = histc(spikes{cellnum,trial},time);

STA = zeros(m,n);
spikeCountsShift = observedSpikes/sum(observedSpikes);
for j=1:n % we loop over time offsets as before
    STA(:,j) = r*spikeCountsShift;
    spikeCountsShift = circshift(spikeCountsShift,[-1 0]); %shift things one time bin into the past and repeat
end


figure(2);clf;
subplot(1,2,1);
imagesc(STA);
keepWin = 168:200; % look at the STA to get a rough sense of where our cell's receptive field is
subplot(1,2,2);
imagesc(STA(keepWin,:));

%% make a temporally shifted version of the stimulus, rstack

L = length(keepWin);
rstack=zeros(n*L,length(r)-n+1);
for t = 1:n
    disp(t)
    inds = (1:L)+L*(t-1);
    rstack(inds,:) = r(keepWin,t:end-n+t);
end

%% add a spike-history term

w2 = 0.5;
n2 = round(w2/stimulus(trial).frame);

spikeStack = zeros(n2,length(observedSpikes));
for t = 2:n2+1
    spikeStack(t-1,t:end-n2+t) = observedSpikes(1:end-n2+1);
end

%% model-fitting!

inds = 100:60000; %pick a range of data to fit over (save some data for testing!)

% possible distributions: 'binomial', 'gamma', 'inverse gaussian', 'normal' (the default), and 'poisson'
mdl     = fitglm([rstack(:,inds); spikeStack(:,n+inds)]',observedSpikes(n+inds)','distribution','poisson');
mdl0    = fitglm(rstack(:,inds)',observedSpikes(n+inds)','distribution','poisson'); %without spike-history model

%%
% with spike-history term:
offset          = mdl.Coefficients.Estimate(1);
stimFilter      = reshape(mdl.Coefficients.Estimate(1+(1:L*n)),L,[]);
historyFilter   = mdl.Coefficients.Estimate(1+L*n+(1:n2));

% without spike-history term
offset0         = mdl0.Coefficients.Estimate(1);
stimFilter0     = reshape(mdl0.Coefficients.Estimate(1+(1:L*n)),L,[]);

%% simulate individual trials:
inds = 65001:66000;
spstore = zeros(10,n2+length(inds));
for rep = 1:10
    disp(rep)
    spikes = [zeros(1,n2) zeros(size(inds))];
    count=n2;
    for t = inds
        count=count+1;
        input   = mdl.Coefficients.Estimate(1) + mdl.Coefficients.Estimate(1+(1:L*n))' * rstack(:,t);
        spHist  = mdl.Coefficients.Estimate(1+L*n+(1:n2))' * fliplr(spikes(count-n2+1:count))';
        drive = exp(input + spHist);
        spikes(count) = poissrnd(drive);
    end
    spstore(rep,:) = spikes;
end
%%
figure(1);clf;
plot(observedSpikes(n+inds),'k');
hold on;
plot(mean(spstore),'r');

%% predict activity directly from the model

inds = 65001:65500; %range of data to test over
testValues  = [rstack(:,inds); spikeStack(:,n+inds)]; %inputs to the full model
testValues0 = rstack(:,inds); %inputs to the model without a spike-history filter

% predict spiking of our cell from the fit model:
predictedSpikes = predict(mdl,testValues');
% sanity check- the following should give the same result for
% predictedSpikes, assuming you used a poisson linking function:
% predictedSpikes     = exp(offset + stimFilter(:)'*rstack(:,inds) + historyFilter'*spikeStack(:,n+inds));


% predict spiking of our cell from the spike-history-less model:
predictedSpikes_0 = predict(mdl0,testValues0');
% another sanity check, for predictedSpikes_0:
% predictedSpikes_0   = exp(offset0 + stimFilter0(:)'*rstack(:,inds));


figure(1);clf
subplot(3,3,1:3);
plot(observedSpikes(n+inds),'k');
hold on;
plot(predictedSpikes,'r');
legend('observed spiking','predicted spiking');
subplot(3,3,4:6);
plot(observedSpikes(n+inds),'k');
hold on;
plot(predictedSpikes_0,'r');
legend('observed spiking','predicted spiking, no history filter');

subplot(3,3,7);
imagesc((1:n)*stimulus(trial).frame,[],fliplr(stimFilter0));
title('Fit stimulus filter, no history term');
xlabel('time (seconds)');
colorbar;

subplot(3,3,8);
imagesc((1:n)*stimulus(trial).frame,[],fliplr(stimFilter));
title('Fit stimulus filter, with history term');
xlabel('time (seconds)');
colorbar;

subplot(3,3,9);
plot((1:n2)*stimulus(trial).frame,historyFilter);
hold on
plot([0 n2]*stimulus(trial).frame,[0 0],'k');
title('Fit post-spike filter');
xlabel('time (seconds)');
xlim([0 n2]*stimulus(trial).frame)



