%% This code can be used to make  Fig 5 and Suppl Fig 5
% Takes a collection of barcodes (bGen) or bargroups
% and compares it against a collection of ref barcodes (bT)
% you need the theory files already, can generate using hca tools
% this contains all for local alignment and plot
% Can also be used to see how bargroups match into theories, just supply cGenAll in place of bGen 

%the uncommented code runs Fig 5

w =150;% minimum overlap length for local alignement
per_stretch=5;

cGenAll=[cGen_final cGen_final_pombe];

%Let's add a proper bitmask here...
for c=1:length(cGenAll)
   cGenAll{c}.rawBitmask(isnan(cGenAll{c}.rawBarcode))=false;
end
cGenAll = keeplongestbitmask(cGenAll); 
tokeep=1:length(cGenAll);

to_run="bargroups"; % bargroups only
%to_run="merged"; %barcodes and bargroups
%to_run="barcodes"; %barcodes only
bGen=barcodeGen_mix; % irrelevant if running bargroup only 

timestamp = datestr(clock(), 'yyyy-mm-dd_HH_MM_SS');
sF = (1-(per_stretch+8)/100):0.01:(1+per_stretch/100); 

%% Load theory
% addpath(genpath('/Users/moacarlsson/Desktop/Yeast reference files'))
% fastas=dir('/Users/moacarlsson/Desktop/Yeast reference files/*fasta');
% addpath(genpath('C:/Users/luisma/OneDrive - Chalmers/yeastdata/ReferenceScer'))
% fastas=dir('C:/Users/luisma/OneDrive - Chalmers/yeastdata/ReferenceScer/*fasta');
%addpath(genpath('C:\Users\luisma\OneDrive - Chalmers\S.pombe genome proyect\Spombe_Genome\ref_fastas'))
%fastas=dir('C:\Users\luisma\OneDrive - Chalmers\S.pombe genome proyect\Spombe_Genome\ref_fastas\*fasta');

%for i=1:length(fastas)
%fastaFile{1,i}=strjoin({fastas(i).folder,'/',fastas(i).name},"");
%end

%quickly calculate theory for given single fasta file (fastaFile, nmbp, nmPerPx, psffac, islinear)
%[theoryStructRev_newpsf_scer,theoryStruct_newpsf_scer,bT_newpsf_scer] = prep_thry_for_local_comp(fastaFile, 0.25, 1,254,1); % I think stretch is closer to 0.2375 not 0.25

%% Data Prep
% make bitmask logical and add the chromosome name and theory check
%bT=bT_newpsf_scer;
%bT=bT_newpsf_scer_split12;% version where scer chr12 is split
%bT=bT_newpsf_pombe;
bT=[bT_newpsf_scer_split12 bT_newpsf_pombe]; %all chromosomes 


for i=1:length(bT)
bT{i}.rawBitmask = logical(bT{i}.rawBitmask);
bT{i}.theory=i;
end
% Expand mask
bitmask_buff=0;
bT = expand_mask(bT, bitmask_buff,1); % last argument adds mask (6px) on the edges


for i=1:length(bGen)
    bGen{i}.index=i;
end

if (to_run == "bargroups" || to_run == "merged" )
for i=1:length(cGenAll)

    cGenAll{i}.bargroupid=i;

end
mini_cGenAll=cGenAll(tokeep);
if isempty(mini_cGenAll)
mini_cGenAll=cGenAll;
end


if (to_run == "merged")
 [bGen_grouped,bGen_trimmed,allIdx,bGen_excluded]  = barcode_bargroup_cleaner(bGen, mini_cGenAll);
bGen=bGen_grouped;
else
bGen=mini_cGenAll;
end

end

for i = 1:numel(bGen)
    bm = bGen{i}.rawBitmask;
    firstOne = find(bm == 1, 1, 'first');
    lastOne  = find(bm == 1, 1, 'last');
    bm(firstOne:lastOne) = 1;
    bGen{i}.rawBitmask = bm;
end

% For same w add theory as last barcode
cGen =[bGen bT];

l=0;
cGen_sizefiltered={};
numTheoryBar=0;
for i=1:length(cGen)

    if max(diff(find([0 cGen{i}.rawBitmask 0]==0)))>=(w*max(sF)) % This weird thing here makes sure all barcodes have a continuous bitmask thats longer that w 
    l=l+1;
    cGen_sizefiltered{l}=cGen{1,i};
    if isfield(cGen_sizefiltered{l},'theory')
    numTheoryBar=numTheoryBar+1;
    end
    end
end

cGen=cGen_sizefiltered;
numExpBar=length(cGen)-numTheoryBar;


%% Can we do a different w for each cGen element?
%numExpBar=numel(bGen);
%numTheoryBar=numel(bT);

bTm=cGen(end-numTheoryBar+1:end);
beGen=cGen(1:numExpBar);
template = struct( ...
    'score', [], ...
    'fullscore', [], ...
    'overlaplen', [], ...
    'pA', [], ...
    'pB', [], ...
    'or', [], ...
    'bestBarStretch', [], ...
    'lenA', [], ...
    'lenB', [], ...
    'partialScore', [], ...
    'partialLength', [], ...
    'h', [] );
minioS = repmat(template, length(beGen), numTheoryBar);
failedIdx = false(length(beGen),1);
tic

parfor c = 1:numExpBar
%for c= bad
try
tGen=[];
tGen{1}=beGen{c};
tGen=[tGen bTm];
w=round(max(diff(find([0 beGen{c}.rawBitmask 0]==0)))/max(sF)*0.7);
w=200;
[tempoS] = calc_overlap_mp(tGen,sF, w,timestamp);
fprintf('barcode: %d\n', c);
tempoS=tempoS(1,2:end);
if ~isempty(tempoS)
    minioS(c,:) = tempoS;
else
    failedIdx(c) = true;
end
catch ME
    warning('Error at c = %d', c);
    failedIdx(c) = true;
end
end
toc

bad=find(failedIdx)'
% if bad is not empty, then re run them without parfor!
% fixing w=200 makes it harder to fail but can be done variable

%% Best is sorted according to local score (not stouffer)
% maybe ccdif no, but stscore or something else
ccdif_thres=0;
s=[];

stscores=[];
pvalLs=[];
pvalLOs=[];
Vals=[];
ccdifs=[];
best_chromosomes=[];
for r=1:numExpBar
 [sortedVals,sortedIds,pscores,fullScores,overlaplen,partialScore,partialLength,pvalL,pvalLO,stscore]=sorted_scores(minioS(r,:),alphaN1,alphaN2,alphaNu1,alphaNu2);
 [~,maxind]=max(stscore);
 
 %how does this look?
 stscores(r)=max(stscore);
 pvalLs(r)=min(pvalL); 
 pvalLOs(r)=min(pvalLO);
Vals(r)=max(sortedVals);
ccdifs(r)=sortedVals(1)-sortedVals(2);
% WHAT TO DO HERE???
if ccdifs(r)<=ccdif_thres
best_chromosomes(r)= 0 ;% if ccdif < 0 kick them out!
else
best_chromosomes(r)=sortedIds(1); % best local score
%best_chromosomes(r)=sortedIds(maxind); %best stouffer score
end
end



font_size=8;
see_local=1; %0 no local, % 1local and global % 2 only local
chr_to_plot=0 ; %(0=all)
consensus_line=1; % 1 for full , 2 for local
layerwidth=5;
threshold=-3;%4.5;% -7.88
horizontalgap=600;
unmask=0;
%% Plot all matches for each chromosome all together.

[f, chromosomeConsensus_bargroup] = plot_per_chromosome(cGen,minioS ,best_chromosomes,layerwidth,chr_to_plot,see_local,font_size,threshold,consensus_line);% other arguments are offset, chromosome (0=all), highlight (to highlight local or not ) 1/0
%% --- Run this one to use standsrd window size for all chr ---
xlim([-100 1700]); % for scer chromosomes
xlim([-100 5600]); % for spombe chromosomes
%xlim([3000 3500]); % for zoom ins for Suppl Fig5
%xlim([4000 5000]); % for zoom ins for Suppl Fig5
%ylim([-40 20]);
ylim([-8 8]);
ax = gca;
ax.Units = 'normalized';
%ax.Position = [0.1 0.15 0.65 0.75];
ax.Position = [0.1 0.15 0.92 0.75];
ax.PositionConstraint = 'innerposition';
ax = gca;
ax.YAxis.Visible = 'off';

lgd = legend('Location','eastoutside');
set(gcf, 'Units','centimeters')
%set(gcf, 'Position',[5 5 70 10])% a good figure size
set(gcf, 'Position',[5 5 80 10])% a good figure size

%set(gcf, 'Position',[5 5 40 5])% a good figure size
set(gcf,'Color','white')
