%% load data, set time ranges
close all; clear all; clc;
BLOCKPATH = uigetdir; %choose tank and block folder
STREAM_STORE1 = 'x405S';
STREAM_STORE2 = 'x465S';
cd(BLOCKPATH);
savename=BLOCKPATH(end-16:end);
data = TDTbin2mat(BLOCKPATH, 'TYPE', {'epocs', 'scalars', 'streams'});
% TRANGE = [-20, 50]; % window size [start time relative to epoc onset, window duration]
display([savename ' loaded']);

%Pre-process behavioral timestamps, insert back in to TDT, extract around filtered events
%Time stamp data - hex, binary, event
% 0011110111111111	15871	Lick+Mag+CSp
% 0011111111111111	16383	Lick+Mag
% 0111101111111111	31743	Lick+CSm
% 0111110111011111	32223	Lick+CSp+Pump
% 0111110111111111	32255	Lick+CSp
% 0111111111111111	32767	Lick
% 1011011111111111	47103	Mag+CSe
% 1011101111111111	48127	Mag+CSm
% 1011110111011111	48607	Mag+CSp+Pump
% 1011110111111111	48639	Mag+CSp
% 1011111111111111	49151	Mag
% 1111011111111111	63487	CSe
% 1111101111111111	64511	CSm
% 1111110111011111	64991	Pump+CSp
% 1111110111111111	65023	CSp
% 1111111111011111	65503	Pump
% 1111111111111111	65535	Offset

%grab event timestamps
CSp=data.epocs.PtAB.onset(data.epocs.PtAB.data == 65023,1);
CSm=data.epocs.PtAB.onset(data.epocs.PtAB.data == 64511,1);
CSe=data.epocs.PtAB.onset(data.epocs.PtAB.data == 63487,1);
Mag=data.epocs.PtAB.onset(data.epocs.PtAB.data == 49151 | data.epocs.PtAB.data == 48639 | data.epocs.PtAB.data == 48639 | data.epocs.PtAB.data == 48607,1);
Pump=data.epocs.PtAB.onset(data.epocs.PtAB.data == 64991 | data.epocs.PtAB.data == 48607 | data.epocs.PtAB.data == 32223,1);
UnQmag=data.epocs.PtAB.onset(data.epocs.PtAB.data == 49151,1);
UnQmag(:,2)=[100;diff(UnQmag(:,1))]; %save first one
UnQmag=UnQmag(UnQmag(:,2) > 10,1); %remove anything within 10 sec inter mag interval

%clean up timestamps
cleanCSp=[CSp(1,1) 100]; %keep first trial
cleanCSp(2:size(CSp,1),1:2)=[CSp(2:end,1) diff(CSp)]; %all CSp nodes plus ITIs
cleanCSp(cleanCSp(:,2) < 10,:)=[]; %remove trials with IQI < 10

cleanCSm=[CSm(1,1) 100]; %keep first trial
cleanCSm(2:size(CSm,1),1:2)=[CSm(2:end,1) diff(CSm)]; %all CSm nodes plus ITIs
cleanCSm(cleanCSm(:,2) < 10,:)=[]; %remove trials with IQI < 10

%grab first head entry after CS-
firstCSmmag=[];
for i = 1:size(cleanCSm,1) %reward head entries
    if isempty(Mag(find(Mag > cleanCSm(i,1),1,'first'),1)) == 0
        firstCSmmag(end+1,1:3)=[Mag(find(Mag > cleanCSm(i,1),1,'first'),1) Mag(find(Mag > cleanCSm(i,1),1,'first'),1)-cleanCSm(i,1) cleanCSm(i,1)];
    end
end
firstCSmmag(firstCSmmag(:,2) > 15,:)=[];%remove nodes > 15 sec after CS-

%grab first head entry after CS+
firstCSpmag=[];
for i = 1:size(cleanCSp,1) %reward head entries
    if isempty(Mag(find(Mag > cleanCSp(i,1),1,'first'),1)) == 0
        firstCSpmag(end+1,1:3)=[Mag(find(Mag > cleanCSp(i,1),1,'first'),1) Mag(find(Mag > cleanCSp(i,1),1,'first'),1)-cleanCSp(i,1) cleanCSp(i,1)];
    end
end
savefirstCSpmag=firstCSpmag(:,1); %save all of them for later with uncued analysis
firstCSpmag(firstCSpmag(:,2) > 20,:)=[];%remove nodes > 10 sec after CS+

%clean up timestamps
cleanCSe=[CSe(1,1) 100]; %keep first trial
cleanCSe(2:size(CSe,1),1:2)=[CSe(2:end,1) diff(CSe)]; %all CSe nodes plus ITIs
cleanCSe(cleanCSe(:,2) < 10,:)=[]; %remove trials with ITI < 10

%grab first head entry after CSe
firstCSemag=[];
for i = 1:size(cleanCSe,1) %reward head entries
    if isempty(Mag(find(Mag > cleanCSe(i,1),1,'first'),1)) == 0
        firstCSemag(end+1,1:3)=[Mag(find(Mag > cleanCSe(i,1),1,'first'),1) Mag(find(Mag > cleanCSe(i,1),1,'first'),1)-cleanCSe(i,1) cleanCSe(i,1)];
    end
end
firstCSemag(firstCSemag(:,2) > 20,:)=[];%remove nodes > 10 sec after CSe

%this might be a more accurate node for rewards --> if CS+ entry before
%reward, use time of pump onset, otherwise use time of CS+ entry after pump
firstCSppump=[];
for i = 1:size(firstCSpmag,1) %reward head entries
    if firstCSpmag(i,2) >= 8.5 %onset of pump after CSp
         firstCSppump(end+1,1:3)=firstCSpmag(i,1:3);
    else
        firstCSppump(end+1,1:3)=[Pump(find(Pump > firstCSpmag(i,3),1,'first'),1) 8.5 firstCSpmag(i,3)];
    end
end

firstCSppumpQ=firstCSppump(:,3);
firstCSemagQ=firstCSemag(:,3);
firstCSpmagQ=firstCSpmag(:,3);
firstCSmmagQ=firstCSmmag(:,3);

%code for uncued port entries
tmp=[];
for i = 1:size(savefirstCSpmag,1)
    if isempty(find(savefirstCSpmag(i,1) == UnQmag)) == 0
        tmp(end+1,:)=find(UnQmag == savefirstCSpmag(i,1));
    else
        tmp(end+1,1)=NaN;
    end
end

UnQmag(isnan(tmp(:,1)) == 0,:)=[];%remove nodes after CSp that animal found reward

%insert cleaned up timestamps back into data so we can use TDTfilter around them
data.epocs.PtAB.data(end+1:end+(size(cleanCSp,1))) = 99; 
data.epocs.PtAB.onset(end+1:end+(size(cleanCSp,1))) = cleanCSp(:,1);
data.epocs.PtAB.offset(end+1:end+(size(cleanCSp,1))) = cleanCSp(:,1)+10;

data.epocs.PtAB.data(end+1:end+(size(cleanCSe,1))) = 27; 
data.epocs.PtAB.onset(end+1:end+(size(cleanCSe,1))) = cleanCSe(:,1);
data.epocs.PtAB.offset(end+1:end+(size(cleanCSe,1))) = cleanCSe(:,1)+10;

data.epocs.PtAB.data(end+1:end+(size(cleanCSm,1))) = 88; 
data.epocs.PtAB.onset(end+1:end+(size(cleanCSm,1))) = cleanCSm(:,1);
data.epocs.PtAB.offset(end+1:end+(size(cleanCSm,1))) = cleanCSm(:,1)+10;

data.epocs.PtAB.data(end+1:end+(size(firstCSmmag,1))) = 92; 
data.epocs.PtAB.onset(end+1:end+(size(firstCSmmag,1))) = firstCSmmag(:,1);
data.epocs.PtAB.offset(end+1:end+(size(firstCSmmag,1))) = firstCSmmag(:,1)+10;

data.epocs.PtAB.data(end+1:end+(size(firstCSpmag,1))) = 47; 
data.epocs.PtAB.onset(end+1:end+(size(firstCSpmag,1))) = firstCSpmag(:,1);
data.epocs.PtAB.offset(end+1:end+(size(firstCSpmag,1))) = firstCSpmag(:,1)+10;

data.epocs.PtAB.data(end+1:end+(size(firstCSemag,1))) = 62; 
data.epocs.PtAB.onset(end+1:end+(size(firstCSemag,1))) = firstCSemag(:,1);
data.epocs.PtAB.offset(end+1:end+(size(firstCSemag,1))) = firstCSemag(:,1)+10;

data.epocs.PtAB.data(end+1:end+(size(firstCSppump,1))) = 84; 
data.epocs.PtAB.onset(end+1:end+(size(firstCSppump,1))) = firstCSppump(:,1);
data.epocs.PtAB.offset(end+1:end+(size(firstCSppump,1))) = firstCSppump(:,1)+10;

%cue insertions needed to match behaviors
data.epocs.PtAB.data(end+1:end+(size(firstCSppumpQ,1))) = 564; 
data.epocs.PtAB.onset(end+1:end+(size(firstCSppumpQ,1))) = firstCSppumpQ(:,1);
data.epocs.PtAB.offset(end+1:end+(size(firstCSppumpQ,1))) = firstCSppumpQ(:,1)+10;

data.epocs.PtAB.data(end+1:end+(size(firstCSemagQ,1))) = 987; 
data.epocs.PtAB.onset(end+1:end+(size(firstCSemagQ,1))) = firstCSemagQ(:,1);
data.epocs.PtAB.offset(end+1:end+(size(firstCSemagQ,1))) = firstCSemagQ(:,1)+10;

data.epocs.PtAB.data(end+1:end+(size(firstCSpmagQ,1))) = 123; 
data.epocs.PtAB.onset(end+1:end+(size(firstCSpmagQ,1))) = firstCSpmagQ(:,1);
data.epocs.PtAB.offset(end+1:end+(size(firstCSpmagQ,1))) = firstCSpmagQ(:,1)+10;

data.epocs.PtAB.data(end+1:end+(size(firstCSmmagQ,1))) = 741; 
data.epocs.PtAB.onset(end+1:end+(size(firstCSmmagQ,1))) = firstCSmmagQ(:,1);
data.epocs.PtAB.offset(end+1:end+(size(firstCSmmagQ,1))) = firstCSmmagQ(:,1)+10;

data.epocs.PtAB.data(end+1:end+(size(UnQmag,1))) = 966; 
data.epocs.PtAB.onset(end+1:end+(size(UnQmag,1))) = UnQmag(:,1);
data.epocs.PtAB.offset(end+1:end+(size(UnQmag,1))) = UnQmag(:,1)+10;
display([savename ' behavior filtered']);

% overview downsampled dF
n=10;
b = arrayfun(@(i) mean(data.streams.x65S.data(i:i+n-1)),1:n:length(data.streams.x65S.data)-n+1)';
a = arrayfun(@(i) mean(data.streams.x05S.data(i:i+n-1)),1:n:length(data.streams.x05S.data)-n+1)';
b=b(500:end-500); a=a(500:end-500); %remove 5 sec at beginning and end for better fit overview

blstmp = polyfit(a, b, 1);
Y_fit_alltmp = blstmp(1) .* a + blstmp(2);
Y_dF_alltmp = b-Y_fit_alltmp;
%
subplot(2,2,1)
plot(b,'-g');
subplot(2,2,2)
plot(a,'-b');
subplot(2,2,3)
plot(Y_fit_alltmp,'-g');
subplot(2,2,4)
plot(Y_dF_alltmp,'-g');hold on; 
print('-djpeg','-r300',[savename '-all']);
close all


%filter the neural data around the bout onset and offset times
cleanCSpdata=TDTfilter(data, 'PtAB', 'VALUES', 99, 'TIME', [-10 40]);
cleanCSedata=TDTfilter(data, 'PtAB', 'VALUES', 27, 'TIME', [-10 40]);
cleanCSmdata=TDTfilter(data, 'PtAB', 'VALUES', 88, 'TIME', [-10 40]);

firstCSpmagQdata=TDTfilter(data, 'PtAB', 'VALUES', 123, 'TIME', [-10 40]);
firstCSppumpQdata=TDTfilter(data, 'PtAB', 'VALUES', 564, 'TIME', [-10 40]);
firstCSemagQdata=TDTfilter(data, 'PtAB', 'VALUES', 987, 'TIME', [-10 40]);
firstCSmmagQdata=TDTfilter(data, 'PtAB', 'VALUES', 741, 'TIME', [-10 40]);

firstCSpmagdata=TDTfilter(data, 'PtAB', 'VALUES', 47, 'TIME', [-20 50]);
firstCSemagdata=TDTfilter(data, 'PtAB', 'VALUES', 62, 'TIME', [-20 50]);
firstCSmmagdata=TDTfilter(data, 'PtAB', 'VALUES', 92, 'TIME', [-20 50]);
firstCSppumpdata=TDTfilter(data, 'PtAB', 'VALUES', 84, 'TIME', [-20 50]);
unqmagdata=TDTfilter(data, 'PtAB', 'VALUES', 966, 'TIME', [-20 50]);
display([savename ' neural extracted around behaviors']);

PhotometryAnalysis(cleanCSpdata,[-10 30],[savename '-AllCSp'],[-3.1 -.1],[-6 -3]);
PhotometryAnalysis(cleanCSedata,[-10 30],[savename '-AllCSe'],[-3.1 -.1],[-6 -3]);
PhotometryAnalysis(cleanCSmdata,[-10 30],[savename '-AllCSm'],[-3.1 -.1],[-6 -3]);

PhotometryAnalysis(firstCSpmagQdata,[-10 30],[savename '-HitQCSp'],[-3.1 -.1],[-6 -3]);
PhotometryAnalysis(firstCSppumpQdata,[-10 30],[savename '-PumpQCSp'],[-3.1 -.1],[-6 -3]);
PhotometryAnalysis(firstCSemagQdata,[-10 30],[savename '-HitQCSe'],[-3.1 -.1],[-6 -3]);
PhotometryAnalysis(firstCSmmagQdata,[-10 30],[savename '-HitQCSm'],[-3.1 -.1],[-6 -3]);

PhotometryAnalysis(firstCSpmagdata,[-20 30],[savename '-MagCSp'],[-13.1 -10.1],[-16 -13]);
PhotometryAnalysis(firstCSppumpdata,[-20 30],[savename '-Pump'],[-13.1 -10.1],[-16 -13]);
PhotometryAnalysis(firstCSemagdata,[-20 30],[savename '-MagCSe'],[-13.1 -10.1],[-16 -13]);
PhotometryAnalysis(firstCSmmagdata,[-20 30],[savename '-MagCSm'],[-13.1 -10.1],[-16 -13]);
PhotometryAnalysis(unqmagdata,[-20 30],[savename '-MagUnQ'],[-3.1 -.1],[-6 -3]);
display('Program Ended');