function main
clc, clear all, close all

global pc kpc Mpc Gpc s minute hr day yr kyr Myr Gyr
global H0 Omega_r0 Omega_m0 Omega_Lambda0

%% Units
% Distance
Gpc = 1;
pc  = 1e-9*Gpc;
kpc = 1e+3*pc;
Mpc = 1e+6*pc;
km  = 3.241e-14*pc;
% Time
Gyr    = 1;
yr     = 1e-9*Gyr;
kyr    = 1e+3*yr;
Myr    = 1e+6*yr;
day    = yr/365.25;
hr     = day/24;
minute = hr/60;
s      = minute/60;

%% Constants
H0 = 70*km/(s*Mpc);
Omega_r0      = 8.5*10^(-5);
Omega_m0      = 0.3;
Omega_Lambda0 = 0.7;

%% Start computation
solve_friedmann
end


function solve_friedmann()
global pc kpc Mpc Gpc s minute hr day yr kyr Myr Gyr
global H0 Omega_r0 Omega_m0 Omega_Lambda0 Omega_0
global t_cum a_cum

% The Friedmann equation
Omega_0 = Omega_r0 + Omega_m0 + Omega_Lambda0;
adot = @(a) a*H0*sqrt(Omega_r0/(a + eps)^4 + Omega_m0/(a + eps)^3 + Omega_Lambda0 + (1 - Omega_0)/a^2);

% Integration limits
t_start = 1*yr;
t_end  = 100*Gyr;

% Find initial value for scale factor
try
    a_start = fzero(@(a) adot(a)*t_start - a, [eps, 1]);
catch
    a_start = 0.01;
end

% Initial plotting
plot(inf, inf, 'r', 'linewidth', 2)
hold on
plot(inf, inf, 'g', 'linewidth', 2)
plot(inf, inf, 'b', 'linewidth', 2)
plot(inf, inf, 'k', 'linewidth', 2)

% Evolve scale factor
rkf45(adot, a_start, t_start, t_end)

% Find present
index = find(a_cum > 1, 1) + 2;
t_0 = spline(a_cum(1:index), t_cum(1:index), 1);

% Plot
if a_cum(end) < 1
    % Big Crunch
    ylims = ylim;
    text_y = mean([ylims(1), 1]);
    hzalign = 'left';
else
    set(gca, 'xscale', 'log')
    set(gca, 'yscale', 'log')
    xlim([t_cum(1), t_cum(end)])
    ylim([a_cum(1), a_cum(end)])
    ylims = ylim;
    text_y = 10^mean([log10(ylims(1)), log10(1)]);
    hzalign = 'right';
end
% Present time
xlims = xlim;
plot([xlims(1), t_0], [1, 1], 'k:')
if Omega_r0 ~= 0 || Omega_m0  ~= 0
    plot([t_0, t_0], [ylims(1), 1], 'k:')
    text(t_0, text_y, ['$$\,t_0 = ', num2str(t_0, '%.2f'), '\, \mathrm{Gyr}$$ '],            ...
                                                     'interpreter', 'latex',         ...
                                                     'horizontalalignment', hzalign, ...
                                                     'fontsize', 12)
else
    set(gca, 'xtick', [])
end
% Axis labels
xlabel('$$t \, [\mathrm{Gyr}]$$', 'interpreter', 'latex', 'fontsize', 12)
ylabel('$$a(t)$$', 'interpreter', 'latex', 'fontsize', 12)
% Legend
legend('Radiation', 'Matter', 'Cosmological Constant', 'Curvature', 'location', 'best')
end

function rkf45(adot, f_start, t_start, t_end)
global t_cum a_cum
global Omega_r0 Omega_m0 Omega_Lambda0 Omega_0
persistent a21 a31 a41 a51 a61 a32 a42 a52 a62 a43 a53 a63 a54 a64 a65
persistent b1 b3 b4 b5 b6
persistent d1 d3 d4 d5
persistent turned
if isempty(a21)
    a21 = 1/4;
    a31 = 3/32;        a32 = 9/32;
    a41 = 1932.0/2197; a42 = -7200/2197;  a43 = 7296/2197;
    a51 = 439/216;     a52 = -8;          a53 = 3680/513;    a54 = -845/4104;
    a61 = -8/27;       a62 = 2;           a63 = -3544/2565;  a64 = 1859/4104;  a65 = -11/40;
    b1  = 16/135;      b3  = 6656/12825;  b4  = 28561/56430; b5  = -9/50;      b6  = 2/55;
    d1  = 25/216;      d3  = 1408/2565;   d4  = 2197/4104;   d5  = -1/5;
end

% Tolerences
delta = 1e-8;
epsilon = 1e-9;
% The maximum and minimum step size
Deltat = t_end - t_start;
h_max  = 0.01*Deltat;
h_min  = 5*eps;
% Initial values
h = h_max*epsilon;
i = length(t_cum);
a = f_start;
t = t_start;
while t < t_end
    % The embedded Runge–Kutta–Fehlberg 4(5) step
    k1 = h*adot(a);
    k2 = h*adot(a + a21*k1);
    k3 = h*adot(a + a31*k1 + a32*k2);
    k4 = h*adot(a + a41*k1 + a42*k2 + a43*k3);
    k5 = h*adot(a + a51*k1 + a52*k2 + a53*k3 + a54*k4);
    k6 = h*adot(a + a61*k1 + a62*k2 + a63*k3 + a64*k4 + a65*k5);
    f5 = a                 +  b1*k1          +  b3*k3 +  b4*k4 +  b5*k5 + b6*k6;
    f4 = a                 +  d1*k1          +  d3*k3 +  d4*k4 +  d5*k5;
    % Check for turning point
    if imag(f5)
        %disp('turning point')
        if isempty(turned) || ~turned
            adot = @(a) -adot(a);
            turned = true;
            continue
        else
            break
        end
    end
    % The error estimate
    error = abs(f5 - f4) + eps;
    % The local tolerence
    tolerence = (epsilon*abs(f5) + delta)*sqrt(h/Deltat);
    if h == h_min && max(a_cum) >= 1
        % Integration error (either physical (Big Crunch) or numerical)
        %disp('break (possible Big Crunch)')
        break
    elseif error < tolerence || h == h_min
        % Step accepted
        t = t + h;
        a = f5;
        i = i + 1;
        % Test or Big Crunch
        if a <= 0
            % Big Crunch
            %disp('Big Crunch')
            break
        end
        % If necessary, t_cum and a_cum get resized (doubled)
        if i > length(t_cum)
            t_cum = [t_cum; zeros(size(t_cum))];
            a_cum = [a_cum; zeros(size(a_cum))];
        end
        t_cum(i) = t;
        a_cum(i) = a;
        % Plot
        if i > 1
            rgb = [Omega_r0/(a + eps)^4, Omega_m0/(a + eps)^3, Omega_Lambda0];
            max_rgb = max(rgb);
            rgb = rgb/ max_rgb;
            omegaK = abs(1 - Omega_0)/(a + eps)^2;
            if omegaK > 0.25*max_rgb
                rgb = rgb*max_rgb/(4*(omegaK + eps));
            end
            if any(rgb > 1)
                rgb = rgb/max(rgb);
            elseif any(rgb < 0)
               rgb = zeros(1, 3); 
            end
            if Omega_r0 == 0 && Omega_m0 == 0 && t < t_start*1000
                % No Big Bang
                plot(t_cum((i - 1):i), a_cum((i - 1):i), '--', 'color', rgb, 'linewidth', 2)
            else
                plot(t_cum((i - 1):i), a_cum((i - 1):i), 'color', rgb, 'linewidth', 2)
            end
            % Positive curvature domination
            if omegaK > 0.01*max_rgb && Omega_0 > 1
                if t_end < t*2
                    t_end = t*2;
                    Deltat = t_end - t_start;
                    h_max  = 0.01*Deltat;
                end
            end
        end
    end
    % Updating step size
    h = h*0.95*(tolerence/error)^0.25;
    if h > h_max
        h = h_max;
    end
    if h < h_min
        h = h_min;
    end
    if t + h > t_end
        h = t_end - t;
    end
    if i > 0 && mod(i, 1000) == 0
       disp(['a(', num2str(t), ' Gyr) = ', num2str(a)]) 
    end
end
t_cum = t_cum(1:i);
a_cum = a_cum(1:i);
end
