%%Part 3 of the program accessed from "run_fitting.m". Read program notes
%%of that program before running this segment. 

%% MLE/LSQ METHODS ....        
%% gaussian distribution (Will derive it in class)
%% 
%% Written by Dhanoop Varghese, Sept. 2006. 
%% Version 2
%% This version allows negative values for mu grid for better fitting

%clear all;
%close all;
%clc;

%% Read in density function 
fden = load('fden.dat');
n = fden(:,1); 
f=fden(:,2);

%% Plot density function
figure(1);
loglog(n, f, 'ob', 'linewidth', 2);
set(gca, 'ylim', [1e-5 max(f)*5]);
hold on;

%% Find mean and standard deviation from read in data
mu0 = sum(f.*n)/sum(f);
sig0 = sqrt(sum(f.*(n-mu0).^2)/sum(f));
fprintf('Data statistics from input file:\n');
fprintf('mu0 = %e, sig0 = %e\n', mu0, sig0);

%% Read in distribution function
Fdist = load('Fdist.dat');
N = Fdist(:,1);

%% MLE fitting
fprintf('\n\nMLE fitting..\n');

%% Define grid for mean and standard deviation
clear mu; clear sig;
%mu = mu0*10.^[-1:0.1:1];
mu = mu0*[-5:0.1:1];
sig = sig0*10.^[-1:0.1:1];
fprintf('Grid boundaries for MLE:\n');
fprintf('mu_min = %e, mu_max = %e\n', min(mu), max(mu));
fprintf('sig_min = %e, sig_max = %e\n', min(sig), max(sig));

%% Evaluate joint probability
%% For each mu
for i=1:size(mu, 2)
    %% For each sigma
    for j=1:size(sig, 2)
        prob(i, j)=0;
        %% Area under density function should be 1
        %% Evaluate normalizing factor 'norm' for each (mu, sig)
        c1 = sqrt(pi)*sig(j);
        norm = quadl(@(x) exp(-((x-mu(i))/sig(j)).^2)/c1, min(n), mu(i)+10*sig(j));
        for k=1:size(N, 1)
            if (norm < 1e-2)
                prob(i, j) = -Inf;
                break;
            end
            c2 = ((N(k)-mu(i))/sig(j))^2;
            peval = exp(-c2)/(c1*norm);
            if (peval == 0)
                prob(i, j) = -Inf;
                break;
            else
                prob(i, j) = prob(i, j)+ log(peval);
            end
        end
    end
end

% Plot joint probability for the defined (mu, sig) grid
figure(2)
%mesh(log(sig), log(mu), prob);
mesh(log(sig), mu, prob);

% Find (<mu>, <sig>) for which joint probability maximizes
[I, J] = find(prob==max(max(prob)));
fprintf('\nParameters estimated using MLE:\n');
fprintf('mu = %e, sig = %e\n', mu(I), sig(J));

% Check if <mu>, <sig> lies on grid boundary
if ((I==1) || (I==size(mu,2)))
    fprintf('\nWARN: <mu>=%e found at grid boundary..\n', mu(I))
    fprintf('WARN: The estimated value might be wrong!\n');
    fprintf('WARN: Please resize the mu grid and re-run..\n');
    %return;
end

if ((J==1) || (J==size(sig,2)))
    fprintf('\nWARN: <sigma>=%e found at grid boundary..\n', sig(J));
    fprintf('WARN: The estimated value might be wrong!\n');
    fprintf('WARN: Please resize the sigma grid and re-run..\n');
    %return;
end

% Evaluate density function using estimated parameters
Nstep = 500;
Nmin = min(n);
Nmax = max(n);

c1 = sqrt(pi)*sig(J);
norm = quadl(@(x) exp(-((x-mu(I))/sig(J)).^2)/c1, min(n), mu(I)+10*sig(J));
        
for i=1:Nstep+1
    Neval(i) = Nmin + (Nmax-Nmin)*(i-1)/Nstep;
    c2 = ((Neval(i)-mu(I))/sig(J))^2;
    feval(i) = exp(-c2)/(c1*norm);
end

figure(1);
loglog(Neval, feval, 'b', 'linewidth', 2);

%% LSQ fitting..
clear mu; clear sig; 
fprintf('\n\nLSQ Method..\n');

%% Define grid for mean and standard deviation
%mu = mu0*10.^[-1:0.1:1];
mu = mu0*[-5:0.1:1];
sig = sig0*10.^[-1:0.1:1];
fprintf('Grid boundaries for LSQ:\n');
fprintf('mu_min = %e, mu_max = %e\n', min(mu), max(mu));
fprintf('sig_min = %e, sig_max = %e\n', min(sig), max(sig));

%% Evaluate LSQ error 
%% for each mu
for i=1:size(mu, 2)
    %% For each sig
    for j=1:size(sig, 2)
        error(i, j)=0;
        %% Area under density function should be 1
        %% Evaluate normalizing factor 'norm' for each (mu, sig)
        c1 = sqrt(pi)*sig(j);
        norm = quadl(@(x) exp(-((x-mu(i))/sig(j)).^2)/c1, min(n), mu(i)+10*sig(j));
        for k=1:size(n, 1)
            if (f(k)==0)
                continue;
            end
            if (norm == 0)
                error(i, j) = Inf;
                break;
            else
                c1 = sqrt(pi)*sig(j);
                c2 = ((n(k)-mu(i))/sig(j))^2;
                error(i, j) = error(i, j)+ (f(k)-exp(-c2)/(c1*norm))^2/f(k)^2;
            end
        end
    end
end

%% Plot LSQ error for the defined grid
figure(3)
%mesh(log(sig), log(mu), log(error));
mesh(log(sig), mu, log(error));

%% Find (<mu>, <sig>) for which LSQ error minimizes
[I, J] = find(error==min(min(error)));
fprintf('\nParameters estimated using LSQ:\n');
fprintf('mu = %e, sig = %e\n', mu(I), sig(J));

%% Check if (<mu>, <sig>) lies on grid boundary
if ((I==1) || (I==size(mu,2)))
    fprintf('\nWARN: mu=%e found at grid boundary..\n', mu(I))
    fprintf('WARN: The estimated value might be wrong!\n');
    fprintf('WARN: Please resize the mu grid and re-run..\n');
    %return;
end

if ((J==1) || (J==size(sig,2)))
    fprintf('\nWARN: sigma=%e found at grid boundary..\n', sig(J));
    fprintf('WARN: The estimated value might be wrong!\n');
    fprintf('WARN: Please resize the sigma grid and re-run..\n');
    %return;
end

%% Evaluate density function for estimated (mu0, sig0)
Nstep = 500;
Nmin = min(n);
Nmax = max(n);

c1 = sqrt(pi)*sig(J);
norm = quadl(@(x) exp(-((x-mu(I))/sig(J)).^2)/c1, min(n), mu(I)+10*sig(J));
        
for i=1:Nstep+1
    Neval(i) = Nmin + (Nmax-Nmin)*(i-1)/Nstep;
    c1 = sqrt(pi)*sig(J);
    c2 = ((Neval(i)-mu(I))/sig(J))^2;
    feval(i) = exp(-c2)/(c1*norm);
end

figure(1);
loglog(Neval, feval, '--r', 'linewidth', 2);
set(gca, 'linewidth', 2, 'fontsize', 14);
xlabel('N -->');
ylabel('f -->');
legend('MC', 'MLE', 'LSQ');

figure(2);
set(gca, 'linewidth', 2, 'fontsize', 14);
xlabel('log(\sigma)');
ylabel('\mu');
zlabel('log(Prob)');

figure(3);
set(gca, 'linewidth', 2, 'fontsize', 14);
xlabel('log(\sigma)');
ylabel('\mu');
zlabel('LSQ Error');

figure(1);