% This is the wrapper for fitting exercise 1, the zero-erosion one. 
% Makes Figure 6 in the paper. 
% 
% Greg Balco - Berkeley Geochronology Center - 2018
% Not licensed for redistribution

clear all; close all;

% Get data
d = create_BCO_data();

% Atmospheric pressure
h = antatm(2183);

% Surface production rate

P10 = 33.2; % This is default in v3 calculators with CRONUS "Primary" data set and St scaling
%P10 = 37.3; % Same thing, LSDn scaling...note that this basically makes no difference. 
P21 = P10.*4; % Ratio is basically the average of lots of studies. 

p.P0 = P21;

% Generate muon fluxes and stopping rates at all sample depths
mc.Natoms = 1.003e22; % For Si in SiO2
mc.sigma0 = 1; % Dummy
mc.k_neg = 1; % Dummy
for a = 1:length(d.N21m);
    m = P_mu_total_alpha1(d.z(a),h,mc,'yes');    
    d.stub_mfast(a) = m.P_fast;
    d.stub_mneg(a) = m.P_neg;
end;

d.stub_mfast = d.stub_mfast'; % ?
d.stub_mneg = d.stub_mneg';

% Other constants for negative muon capture production

p.fC = 0.296; % Heisinger 2002b Table 2 for Si in SiO2
p.fD = 0.6559; % Same

% Decay constants for nucleogenic production
p.l232 = -log(0.5)./1.405e10;
p.l235 = -log(0.5)./7.038e8;
p.l238 = -log(0.5)./4.468e9;

% Uncertainty expansions
p.dnfactor = 0.0; % Define uncert on nucleogenic Ne-21, e.g., 0.1 is 10%
p.dmfactor = 1; % Increase measurement uncertainty to test. Not used. 

% Initial guess 

t0 = 4; % Ma
tc0 = 200; % Ma;
Lsp0 = 100; % g/cm2
La0 = 5; % g/cm2
fa0 = 4; % albedo loss fraction
fstar0 = 2; % fraction in 10^-3
sigma00 = 0.1; % millibarns

x0 = [t0 tc0 Lsp0 La0 fa0 fstar0 sigma00];
xmin = [0 0 0 0 0 0 0];
xmax = [Inf Inf Inf 20 5 Inf Inf];

%% Try optimizing

opts = optimset('fmincon');
opts = optimset(opts,'display','iter','tolfun',1e-7);

[optx,fval] = fmincon(@(x) BCO_objective_zeroerosion(x,d,p,0),x0,[],[],[],[],xmin,xmax,[],opts);

opt.t = optx(1).*1e6;
opt.tc = optx(2).*1e6;
opt.Lsp = optx(3);
opt.La = optx(4);
opt.fa = optx(5).*0.01;
opt.fstar = optx(6).*1e-3;
opt.sigma0 = optx(7).*1e-27;

% Report results
disp(['texp = ' sprintf('%0.2f',optx(1)) ' Ma; tc = ' sprintf('%0.1f',optx(2)) ' Ma']);
disp(['Lsp = ' sprintf('%0.2f',optx(3)) '; La = ' sprintf('%0.2f',optx(4)) '; fa = ' sprintf('%0.3f',opt.fa)]);
disp(['fstar = ' sprintf('%0.4f',opt.fstar) '; sigma0 = ' sprintf('%0.4f',optx(7)) ' millibarns']);

%% Plot results

% Here change nucleogenic Ne-21 uncertainty if you want to test its effect
% on model fit statistic.
%p.dnfactor = 0.25; % Define uncert on nucleogenic Ne-21, e.g., 0.1 is 10%

% Get diagnostics
optd = BCO_objective_zeroerosion(optx,d,p,1);

% Generate model curve

pz = [0:100 110:10:6000];
for a = 1:length(pz);
    m = P_mu_total_alpha1(pz(a),h,mc,'yes');    
    pstub_mfast(a) = m.P_fast;
    pstub_mneg(a) = m.P_neg;
end;

%% Plot non-nucleogenic concentrations
dd = d;
dd.N21m = d.N21m - optd.N21nuc;
dd.delN21m = sqrt(d.delN21m.^2 + optd.delN21nuc.^2);
plotaxes = plot_to_show_model(dd);

pred.N21sp = opt.t.*p.P0.*exp(-pz./opt.Lsp) - opt.t.*p.P0.*opt.fa.*exp(-pz./opt.La);
pred.N21fast = opt.t.*pstub_mfast.*opt.sigma0;
pred.N21neg = opt.t.*pstub_mneg.*opt.fstar.*p.fC.*p.fD;
pred.N21 = pred.N21sp + pred.N21fast + pred.N21neg;

% Main axis

axes(plotaxes(1));
% Plot nucleogenic
plot(optd.N21nuc,dd.z,'o','markersize',10,'color',[0.75 0.75 0.75]);
% Plot model
plot(pred.N21sp,pz,'k--');
plot(pred.N21fast,pz,'k--');
plot(pred.N21neg,pz,'k--');
plot(pred.N21,pz,'k');

axes(plotaxes(2));
% Plot model
plot(pred.N21,pz,'k');

axes(plotaxes(3));
% Plot model
plot(pred.N21,pz,'k');
%set(gca,'xlim',[1e4 1e9])

%% Plot normalized residuals on same figure

set1 = find(d.source == 1 | d.source == 3);
set2 = find(d.source == 2);

set1 = find(d.lab == 1);
set2 = find(d.lab == 2);

% Lower 

axes(plotaxes(4));
for a = 1:length(d.z);
    yy = [d.z(a) d.z(a)];
    xx = [optd.miss(a)-1 optd.miss(a)+1];
    plot(xx,yy,'k'); hold on;
end;

slab = 'v^';
for a = 1:length(d.z);
    if d.source(a) == 1; 
        cc = [0 0 0]; 
    elseif d.source(a) == 2;
        cc = [1 1 1];
    elseif d.source(a) == 3;
        cc = [0.5 0.5 0.5];
    end;
    plot(optd.miss(a),d.z(a),['k' slab(d.lab(a))],'markerfacecolor',cc);
end;

set(gca,'ydir','reverse','ylim',[0 6000]);
set(gca,'xlim',[-4 4]);
plot([0 0],get(gca,'ylim'),'k');
grid on;
xlabel('Normalized residual');
set(gca,'yticklabel',[]);

% Middle 

axes(plotaxes(5)); cla;
for a = 1:length(d.z);
    yy = [d.z(a) d.z(a)];
    xx = [optd.miss(a)-1 optd.miss(a)+1];
    plot(xx,yy,'k'); hold on;
end;

slab = 'v^';
for a = 1:length(d.z);
    if d.source(a) == 1; 
        cc = [0 0 0]; 
    elseif d.source(a) == 2;
        cc = [1 1 1];
    elseif d.source(a) == 3;
        cc = [0.5 0.5 0.5];
    end;
    plot(optd.miss(a),d.z(a),['k' slab(d.lab(a))],'markerfacecolor',cc);
end;

set(gca,'ydir','reverse','ylim',[0 500]);
set(gca,'xlim',[-4 4]);
plot([0 0],get(gca,'ylim'),'k');
set(gca,'xticklabel',[],'yticklabel',[]);
grid on;

% Upper

axes(plotaxes(6)); cla;
for a = 1:length(d.z);
    yy = [d.z(a) d.z(a)];
    xx = [optd.miss(a)-1 optd.miss(a)+1];
    plot(xx,yy,'k'); hold on;
end;

slab = 'v^';
for a = 1:length(d.z);
    if d.source(a) == 1; 
        cc = [0 0 0]; 
    elseif d.source(a) == 2;
        cc = [1 1 1];
    elseif d.source(a) == 3;
        cc = [0.5 0.5 0.5];
    end;
    plot(optd.miss(a),d.z(a),['k' slab(d.lab(a))],'markerfacecolor',cc);
end;

set(gca,'ydir','reverse','ylim',[0 50]);
set(gca,'xlim',[-4 4]);
plot([0 0],get(gca,'ylim'),'k');
set(gca,'xticklabel',[],'yticklabel',[]);
grid on;


%% Plot normalized residuals on second figure

set1 = find(d.source == 1 | d.source == 3);
set2 = find(d.source == 2);

set1 = find(d.lab == 1);
set2 = find(d.lab == 2);

figure;
for a = 1:length(d.z);
    xx = [d.z(a) d.z(a)];
    yy = [optd.miss(a)-1 optd.miss(a)+1];
    plot(xx,yy,'k'); hold on;
end;

slab = 'v^';
for a = 1:length(d.z);
    if d.source(a) == 1; 
        cc = [0 0 0]; 
    elseif d.source(a) == 2;
        cc = [1 1 1];
    elseif d.source(a) == 3;
        cc = [0.5 0.5 0.5];
    end;
    plot(d.z(a),optd.miss(a),['k' slab(d.lab(a))],'markerfacecolor',cc);
end;

set(gca,'xscale','log');
p1 = plot(get(gca,'xlim'),[0 0],'k'); grid on;
xlabel('Depth (g cm^{-2})'); ylabel('Normalized residual');
set(gca,'ylim',[-4 4])



%% Plot non-normalized residuals

miss1 = optd.N21p-d.N21m;
dmiss1 = sqrt(optd.delN21p.^2 + d.delN21m.^2);

figure; 
 
for a = 1:length(d.z);
    xx = [d.z(a) d.z(a)];
    yy = [miss1(a)-dmiss1(a) miss1(a)+dmiss1(a)];
    plot(xx,yy,'k'); hold on;
end;

for a = 1:length(d.z);
    if d.source(a) == 1; 
        cc = [0 0 0]; 
    elseif d.source(a) == 2;
        cc = [1 1 1];
    elseif d.source(a) == 3;
        cc = [0.5 0.5 0.5];
    end;
    plot(d.z(a),miss1(a),['k' slab(d.lab(a))],'markerfacecolor',cc);
end;

plot(get(gca,'xlim'),[0 0],'k'); grid on;
set(gca,'xscale','log');
set(gca,'ylim',[-40e6 40e6]);
xlabel('Depth (g cm^{-2})'); ylabel('Residual (Matoms g^{-1})');


