% Chromosome Fisher
% this code is used to fish for molecules that span full or near full
% chromosomes

% we can go for a full fish or a partial fish
%  Let's place any changable parameters here at the start?
%0. Load and prep theory
%   -Load and Bitmask the theory and expand bitmask
%1. Select which chromosome to fish for
%   - Extract the size from bT
%2. Fish molecules from barcodeGen
%   - mask peaks/reps expand masks
%   - Set a threshold and fish mols that that go +- the threshold
%3. Local alignment
%   - Maybe we want w= 70% chromsome size. Maybe we fish for all
%   chromosomes within a size, not just a chromosome...
%4. Choose the best matches 
%   - Maybe calculate Stouffer? or weighted score
%   - Sort if nothing works

%% Parameters

bGen=barcodeGen_allmerged_scer;
chr=3; % Which chromosome to choose 1:16
chrfr= 0.7; % fraction of the chromosome to look for (size)
mol_threshold= 0.4; % mols will be as small as chrfr but as large as this times more

masks_on=0;

sF = 0.85:0.01:1.15; 
timestamp = datestr(clock(), 'yyyy-mm-dd_HH_MM_SS');

%% Load and prep theory
bT=bT_newpsf_scer;

bitmask_buff=2;

if masks_on==1
% Mask peaks
bT=mask_peaks( bT,3.5);

% Apply mask
for k=1:length(bT)
%     if isfield(cGen{1},'repmask')
%     cGen{k}.rawBitmask= cGen{k}.rawBitmask & ~cGen{k}.repmask;
%     end
    if isfield(bT{k},'peakBitmask')
    bT{k}.rawBitmask= bT{k}.rawBitmask & bT{k}.peakBitmask;
    end
end
end

% Expand mask
bT = expand_mask(bT, bitmask_buff,1); % last argument adds mask (6px) on the edges

%% Chromosome selection

%Get chromosome from chr and bT and its longest continuous region
ref_chr=bT{chr};
if sum(ref_chr.rawBitmask==0)==0 %if no bitmask, then just take whole thing
ref_chr_size=length(ref_chr.rawBarcode);
else
ref_chr_size=max(diff(find(ref_chr.rawBitmask==0)));  % This weird thing here finds the largest continuous bitmask 
end

w=ceil(ref_chr_size*chrfr); % overlap window should be a fraction of the largets chromosome
%w=200

%% Fish molecules from barcodeGen

%Mask peaks/reps expand masks
if masks_on==1
% Mask peaks
bGen=mask_peaks( bGen,3.5);
% Rep mask
[~,~,~,bGen]=periodic_finder(bGen);

Apply mask
for k=1:length(bT)
    if isfield(cGen{1},'repmask')
    cGen{k}.rawBitmask= cGen{k}.rawBitmask & ~cGen{k}.repmask;
    end
    if isfield(bT{k},'peakBitmask')
    bT{k}.rawBitmask= bT{k}.rawBitmask & bT{k}.peakBitmask;
    end
end
end
%Expand mask
bGen = expand_mask(bGen, bitmask_buff,0); % last argument adds mask (6px) on the edges

%Fish mols that that go +- the threshold
% size check
bGen_sizefiltered={};
l=0;
for i=1:length(bGen)

    if max(diff(find(bGen{i}.rawBitmask==0)))>=(w*max(sF)) && max(diff(find(bGen{i}.rawBitmask==0)))<=(ref_chr_size*(1+mol_threshold)) % This weird thing here makes sure all barcodes have a continuous bitmask thats longer that w and shorter than the mol threshold
    l=l+1;
    bGen_sizefiltered{l}=bGen{i};
    bGen_sizefiltered{l}.index=l;
    end

end

bGen=bGen_sizefiltered;

%% Local alignment
%   - Maybe we want w= 70% chromsome size. Maybe we fish for all
%   chromosomes within a size, not just a chromosome...
cGen =[bGen bT]; %if a super short chromosome is used here, it can go wrong

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

    if max(diff(find([0 cGen{i}.rawBitmask 0]==0)))>=(w*max(sF)) && max(diff(find([0 cGen{i}.rawBitmask 0]==0)))<=(ref_chr_size*(1+mol_threshold))
    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;

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 bGen{c}.rawBitmask 0]==0)))/max(sF)*0.7);
[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 rerun without parfor

 %% Choose the best matches 

% Best is sorted according to local score (not stouffer)
best_chromosomes=[];
for r=1:numExpBar
 [~,sortedIds,~,~,~,~,~,~,~,stscore]=sorted_scores(minioS(r,:),alphaN1,alphaN2,alphaNu1,alphaNu2);
 [~,maxind]=max(stscore);

best_chromosomes(r)=sortedIds(1); % best local score
best_chromosomes(r)=sortedIds(maxind); %best stouffer score
end

%%
threshold=0;

chr=0; % Which chromosome to choose 1,3,5!,6,8!,9,10!,11!
% Plot all matches for each chromosome all together.
[f,consensus_chr] = plot_per_chromosome(cGen, minioS,best_chromosomes,3,chr,0,10,threshold,1);
%plot_per_chromosome(barcodeGen, overlapStruct,best_chromosomes,layerwidth, chromosome,highlight_overlap,font_size,threshold,consensus_line)

