%% Get place cells OML
    % REPEAT ALL SES PQ fixOffsetTrialsDay WAS BAD
clearvars;
dirData = 'A:\OptoMECLEC\';
% animals = {'OML5','OML8','OML10','OML11'};
% sessions = {{'day6'},{'day7','day8','day9'},{'day9'},{'day10'}};
animals = {'OML5','OML3','OML7','OML8','OML10','OML11'};
sessions = {{'day4','day5','day6','day8','day9','day13'},...
            {'day11','day12','day13','day17'},...
            {'day7','day9'},...
            {'day6','day7','day8','day9'},...
            {'day7','day9'},...
            {'day10'}};

    %% Get postions for trials 
    for rat = 1:length(animals)
        for ses = 1:length(sessions{rat})
            %try
            basepath = [dirData animals{rat} '\' sessions{rat}{ses}];
            cd(basepath);
            basename = bz_BasenameFromBasepath(pwd);
            if exist([basename '.tracking.behavior.mat'])
               load([basename '.tracking.behavior.mat']);
            else
                [tracking] = catPositionsDay(basepath);  
            end
            if exist('cheesbTrialsDayOffset.mat')  
                load('cheesbTrialsDayOffset.mat');
            else
                load('cheesbTrialsDay.mat');
                load('sessionsSeq.mat')
                load([basename '.MergePoints.events.mat']);
                cheesbTrialsDay = fixOffsetTrialsDay(cheesbTrialsDay,sessionsSeq,MergePoints,1);
                clear sessionsSeq MergePoints              
            end
            
           if ~exist([basename '.posTrials.mat'])
            % trial structure
            if isfield(cheesbTrialsDay,'preProbe')
               posTrials{1}.int(:,1) = cheesbTrialsDay.preProbe.start-10;
               posTrials{1}.int(:,2) = cheesbTrialsDay.preProbe.end+10;
            else
                posTrials{1}.int  = [];
            end
            if isfield(cheesbTrialsDay,'postProbe')
               posTrials{5}.int(:,1) = cheesbTrialsDay.postProbe.start-10;
               posTrials{5}.int(:,2) = cheesbTrialsDay.postProbe.end+10;
            else
               posTrials{5}.int  = [];
            end
            if ~isempty(cheesbTrialsDay.trials.start)
               posTrials{2}.int(:,1) = cheesbTrialsDay.trials.start(1:10,:)-10;
               posTrials{2}.int(:,2) = cheesbTrialsDay.trials.end(1:10,:)+10;  
               posTrials{3}.int(:,1) = cheesbTrialsDay.trials.start(11:20,:)-10;
               posTrials{3}.int(:,2) = cheesbTrialsDay.trials.end(11:20,:)+10;                 
               posTrials{4}.int(:,1) = cheesbTrialsDay.trials.start(21:end,:)-10;
               posTrials{4}.int(:,2) = cheesbTrialsDay.trials.end(21:end,:)+10;                  
            else
               posTrials{2}.int  = [];  posTrials{3}.int  = []; posTrials{4}.int  = [];
            end
            
           posXY = [tracking.timestamps tracking.position.x tracking.position.z];
           % remove outlayers
           for i = 1:size(posXY,1) % Wmaze3
               if posXY(i,2) > 1.5 || posXY(i,2) < -1 || ...
                  posXY(i,3) > 1 || posXY(i,3) < -1
                  posXY(i,2:3) = nan(1,2);
               end
           end
           % normalize positons 
           temp2 = ZeroToOne(posXY(:,2));temp3 = ZeroToOne(posXY(:,3));
           posXY(:,2) = temp2; posXY(:,3) = temp3; clear temp2 temp3;
           figure; plot(posXY(:,2),posXY(:,3));  
           
           % velocity 
           Psm = [posXY(:,1) Smooth(posXY(:,2),25,'kernel','rectangular') Smooth(posXY(:,3),25,'kernel','rectangular')];           
           Vsm = LinearVelocity(Psm); 
           
           % movement periods 
            Vth = 0.05; %decide Vthreshold           
            [Mperiods,mob] = MovementPeriods(Vsm,Vth,1);           
            
           % restrict pos,vel, mob to trials
           for t =1:5
               for i = 1:length(posTrials{t}.int)
                   if ~isempty(posTrials{t}.int)
                       posTrials{t}.pos = Restrict(posXY,posTrials{t}.int); 
                       posTrials{t}.vel = Restrict(Vsm,posTrials{t}.int); 
                       posTrials{t}.mob = Restrict(mob,posTrials{t}.int);
                   else
                       posTrials{t}.pos = [];posTrials{t}.vel = [];posTrials{t}.mob = [];
                   end
               end
           end
           
           % restrict pos to movement
           for t =1:5
               for i = 1:length(posTrials{t}.int)
                   if ~isempty(posTrials{t}.int)
                       posTrials{t}.posMv = Restrict(posTrials{t}.pos,Mperiods); 
                   else
                       posTrials{t}.posMv = [];
                   end
               end
           end           
           
           figure;
            plot(posTrials{2}.pos(:,1),posTrials{2}.pos(:,2)/2,'b');hold on;
            plot(posTrials{2}.pos(:,1),posTrials{2}.vel(:,2),'r');hold on;
            plot(posTrials{2}.pos(:,1),posTrials{2}.mob(:,2)/2.5,'g');
            ylim([0 0.5]);title('velocity trials');
            saveas(gcf,'mobPeriodsTask.fig');
           
           figure;
           for t = 1:5
               try
               if ~isempty(posTrials{t}.int)
               subplot(1,5,t); 
               plot(posTrials{t}.pos(:,2),posTrials{t}.pos(:,3));hold on;
               plot(posTrials{t}.posMv(:,2),posTrials{t}.posMv(:,3),'r');hold on;
               ylim([0 1]);xlim([0 1]);
               end
               catch
               end
           end
           %suptitle([animals{rat} sessions{rat}{ses}]);
           saveas(gcf,'pos.fig');
           
           figure;
           try
           for i = 1:size(posTrials{1}.int,1)
                subplot(5,10,i);
                posT = Restrict(posTrials{1}.pos,posTrials{1}.int(i,:));
                posMvT = Restrict(posTrials{1}.posMv,posTrials{1}.int(i,:));
                plot(posT(:,2),posT(:,3));hold on;
                plot(posMvT(:,2),posMvT(:,3),'r');hold on;
                ylim([0 1]);xlim([0 1]);                
           end   
           catch
           end
           try
           for i = 1:length(posTrials{2}.int)
                subplot(5,10,i+10);
                posT = Restrict(posTrials{2}.pos,posTrials{2}.int(i,:));
                posMvT = Restrict(posTrials{2}.posMv,posTrials{2}.int(i,:));
                plot(posT(:,2),posT(:,3));hold on;
                plot(posMvT(:,2),posMvT(:,3),'r');hold on;
                ylim([0 1]);xlim([0 1]);                
           end 
           catch
           end
           try           
           for i = 1:length(posTrials{3}.int)
                subplot(5,10,i+20);
                posT = Restrict(posTrials{3}.pos,posTrials{3}.int(i,:));
                posMvT = Restrict(posTrials{3}.posMv,posTrials{3}.int(i,:));
                plot(posT(:,2),posT(:,3));hold on;
                plot(posMvT(:,2),posMvT(:,3),'r');hold on;
                ylim([0 1]);xlim([0 1]);                
           end  
           catch
           end
           try           
           for i = 1:length(posTrials{4}.int)
                subplot(5,10,i+30);
                posT = Restrict(posTrials{4}.pos,posTrials{4}.int(i,:));
                posMvT = Restrict(posTrials{4}.posMv,posTrials{4}.int(i,:));
                plot(posT(:,2),posT(:,3));hold on;
                plot(posMvT(:,2),posMvT(:,3),'r');hold on;
                ylim([0 1]);xlim([0 1]);                
           end  
           catch
           end
           try           
           for i = 1:size(posTrials{5}.int,1)
                subplot(5,10,i+40);
                posT = Restrict(posTrials{5}.pos,posTrials{5}.int(i,:));
                posMvT = Restrict(posTrials{5}.posMv,posTrials{5}.int(i,:));
                plot(posT(:,2),posT(:,3));hold on;
                plot(posMvT(:,2),posMvT(:,3),'r');hold on;
                ylim([0 1]);xlim([0 1]);                
           end  
           catch
           end
           saveas(gcf,'posTrials.fig');
           
           save([basename '.posTrials.mat'],'posTrials');
           %clear posXY cheesbTrialsDay tracking Vsm Psm 
           close all; 
           clear posTrials cheesbTrialsDay posT posMvT
           end
           
           disp([animals{rat} '\' sessions{rat}{ses} ' DONE']);
%             catch
%            warning([animals{rat} '\' sessions{rat}{ses} ' ERROR']);     
%             end
        end
    end
    
    %% get spikes - move out
    for rat = 1%:length(animals)
        for ses = 1 %:length(sessions{rat})
            try
            basepath = [dirData animals{rat} '\' sessions{rat}{ses}];
            cd(basepath);
            basename = bz_BasenameFromBasepath(pwd);
            load([basename '.posTrials.mat']);
            
            kilosortFolder = dir('*Kilosort*');
            
            [spikes] = bz_LoadPhy('basepath',basepath,'kilosort_path',kilosortFolder.name,'fs',30000,'nChannels',64,...
                'noPrompts',true,'getWaveforms',false);
            clear spikes;
            disp(['%%%%%%%%%%%%%%%% ' dirData animals{rat} '\' sessions{rat}{ses} ' spk DONE']);
            catch
            disp(['%%%%%%%%%%%%%%%% ' dirData animals{rat} '\' sessions{rat}{ses} ' spk PROBLEM']);
            end
        end
    end
    
%%  Calculate PCs
    for rat = 1:length(animals)
        for ses = 1:length(sessions{rat})
            try
            basepath = [dirData animals{rat} '\' sessions{rat}{ses}];
            cd(basepath);
            basename = bz_BasenameFromBasepath(pwd);
            load([basename '.posTrials.mat']);
            load([basename '.spikes.cellinfo.mat']);
           
            for t = 1:5 % with added segments linpos
                for i = 1:length(spikes.times)
                    if ~isempty(posTrials{t}.int)  
                        [map{i}{t},stats{i}{t}] = FiringMap(posTrials{t}.posMv,spikes.times{i},...
                        'minTime',0.01,'nBins',[100 100],'threshold',0.2,'minSize',200,'minPeak',2,'smooth',3); 
                    else
                        map{i}{t} = []; stats{i}{t} = [];
                    end
                end
            end             
            save([basename '.PCs.mat'],'map','stats');
            
            % plot
            mkdir('PCs');
            for i = 1:length(spikes.times)
                for t = 1:5
                    subplot(1,5,t);
                    if ~isempty(map{i}{t})
                    PlotColorMap(map{i}{t}.rate,map{i}{t}.time,'threshold',0.01); 
                    title(stats{i}{t}.peak(1));caxis([0 Inf]);set(gca,'XTickLabel',[],'YTickLabel',[]);
                    end
                end
                set(gcf,'Position',[500 500 700 300]);
                suptitle([basename ' cell' num2str(i)]);
                saveas(gcf,[pwd '\PCs\cell_' num2str(i) '.png'],'png'); close all;   
            end
            
%             figure;
%             for t = 1:5
%                 subplot(1,5,t);
%                 PlotColorMap(map{1}{t}.time,[],'threshold',0.05);
%             end
            
            %BW = roipoly
            
            clear map stats posTrials spikes
            disp(['%%%%%%%%%%%%%%%% ' dirData animals{rat} '\' sessions{rat}{ses} ' PCs DONE']);
            catch
            warning(['%%%%%%%%%%%%%%%% ' dirData animals{rat} '\' sessions{rat}{ses} ' PCs PROBLEM']);                
            end
        end
    end
         
%% PCs for stim,sham and prbSt
dirSave = 'A:\OptoMECLEC\PCs\';
sesType = {[2 1 2 2 1 4],[2 1 2 4],[1],[3 1 3],[3 3],[3]};% 1-sham, 2-stimMEC, 3-stimLEC, 4-prbSt
% plot 
    for rat = 1:length(animals)
        for ses = 1:length(sessions{rat})
            try
                basepath = [dirData animals{rat} '\' sessions{rat}{ses}];
                cd(basepath);
                basename = bz_BasenameFromBasepath(pwd);
                load([basename '.PCs.mat']);

                if sesType{rat}(ses) == 1
                   for unit = 1:length(stats)
                      if (~isempty(stats{unit}{1}) && stats{unit}{1}.peak(1)>0) || stats{unit}{2}.peak(1)>0 || stats{unit}{3}.peak(1)>0 || stats{unit}{4}.peak(1)>0 || (~isempty(stats{unit}{5}) && stats{unit}{5}.peak(1)>0)
                         figure;
                         for t = 1:5
                             subplot(1,5,t);
                             if ~isempty(map{unit}{t})
                             PlotColorMap(map{unit}{t}.rate,map{unit}{t}.time,'threshold',0.01); 
                             title(stats{unit}{t}.peak(1));caxis([0 Inf]);set(gca,'XTickLabel',[],'YTickLabel',[]);
                             end
                          end
                             set(gcf,'Position',[500 500 700 300]);
                             suptitle([animals{rat} sessions{rat}{ses} ' cell' num2str(unit)]);
                             saveas(gcf,[dirSave 'spatTask_sham\' animals{rat} sessions{rat}{ses} '_cell' num2str(unit) '.png'],'png'); close all;   
                      end
                   end
                   
                elseif sesType{rat}(ses) == 2
                   for unit = 1:length(stats)
                      if (~isempty(stats{unit}{1}) && stats{unit}{1}.peak(1)>0) || stats{unit}{2}.peak(1)>0 || stats{unit}{3}.peak(1)>0 || stats{unit}{4}.peak(1)>0 || (~isempty(stats{unit}{5}) && stats{unit}{5}.peak(1)>0)
                         figure;
                         for t = 1:5
                             subplot(1,5,t);
                             if ~isempty(map{unit}{t}.rate)
                             PlotColorMap(map{unit}{t}.rate,map{unit}{t}.time,'threshold',0.01); 
                             title(stats{unit}{t}.peak(1));caxis([0 Inf]);set(gca,'XTickLabel',[],'YTickLabel',[]);
                             end
                          end
                             set(gcf,'Position',[500 500 700 300]);
                             suptitle([animals{rat} sessions{rat}{ses} ' cell' num2str(unit)]);
                             saveas(gcf,[dirSave 'spatTask_stimMEC\' animals{rat} sessions{rat}{ses} '_cell' num2str(unit) '.png'],'png'); close all;   
                      end
                   end
                    
                elseif sesType{rat}(ses) == 3
                   for unit = 1:length(stats)
                      if (~isempty(stats{unit}{1}) && stats{unit}{1}.peak(1)>0) || stats{unit}{2}.peak(1)>0 || stats{unit}{3}.peak(1)>0 || stats{unit}{4}.peak(1)>0 || (~isempty(stats{unit}{5}) && stats{unit}{5}.peak(1)>0)
                         figure;
                         for t = 1:5
                             subplot(1,5,t);
                             if ~isempty(map{unit}{t})
                             PlotColorMap(map{unit}{t}.rate,map{unit}{t}.time,'threshold',0.01); 
                             title(stats{unit}{t}.peak(1));caxis([0 Inf]);set(gca,'XTickLabel',[],'YTickLabel',[]);
                             end
                          end
                             set(gcf,'Position',[500 500 700 300]);
                             suptitle([animals{rat} sessions{rat}{ses} ' cell' num2str(unit)]);
                             saveas(gcf,[dirSave 'spatTask_stimLEC\' animals{rat} sessions{rat}{ses} '_cell' num2str(unit) '.png'],'png'); close all;   
                      end
                   end
                    
                elseif sesType{rat}(ses) == 4
                   for unit = 1:length(stats)
                      if (~isempty(stats{unit}{1}) && stats{unit}{1}.peak(1)>0) || stats{unit}{2}.peak(1)>0 || stats{unit}{3}.peak(1)>0 || stats{unit}{4}.peak(1)>0 || (~isempty(stats{unit}{5}) && stats{unit}{5}.peak(1)>0)
                         figure;
                         for t = 1:5
                             subplot(1,5,t);
                             if ~isempty(map{unit}{t})
                             PlotColorMap(map{unit}{t}.rate,map{unit}{t}.time,'threshold',0.01); 
                             title(stats{unit}{t}.peak(1));caxis([0 Inf]);set(gca,'XTickLabel',[],'YTickLabel',[]);
                             end
                          end
                             set(gcf,'Position',[500 500 700 300]);
                             suptitle([animals{rat} sessions{rat}{ses} ' cell' num2str(unit)]);
                             saveas(gcf,[dirSave 'spatTask_prbSt\' animals{rat} sessions{rat}{ses} '_cell' num2str(unit) '.png'],'png'); close all;   
                      end
                   end
                    
                end
            
              disp(['%%%%%%%%%%%%%%%% ' dirData animals{rat} '\' sessions{rat}{ses} ' PCs DONE']);  
            catch
              warning(['%%%%%%%%%%%%%%%% ' dirData animals{rat} '\' sessions{rat}{ses} ' PCs PROBLEM']);  
            end
            
        end
    end
    
% mapStats 
    for rat = 1:length(animals)
        for ses = 1:length(sessions{rat})
                basepath = [dirData animals{rat} '\' sessions{rat}{ses}];
                cd(basepath);
                basename = bz_BasenameFromBasepath(pwd);
                load([basename '.PCs.mat'],'map');

                for unit = 1:length(map)
                    for t = 1:5
                        if length(map{unit}) >= t && ~isempty(map{unit}{t})
                        mapStats{unit}{t} = MapStats(map{unit}{t},'threshold',0.2,'minSize',200,'minPeak',2);  
                        [info1,info2,spars,coef,select] = PlaceCellInfo(map{unit}{t}.rate,map{unit}{t}.count,map{unit}{t}.time);
                        mapStats{unit}{t}.info1 = info1; clear info1;
                        mapStats{unit}{t}.info2 = info2; clear info2;
                        mapStats{unit}{t}.spars = spars; clear spars;
                        mapStats{unit}{t}.coef = coef; clear coef;
                        mapStats{unit}{t}.select = select; clear select;                      
                     end
                    end
                end
                save([basename '.PCstats.mat'],'mapStats');
 
        end
    end
    
        
                       
                       
                       


                
            