%% Experimental data, and Data fitting for aiptasia

% aiptasia experimental data set
aiptasiasting = [0.748195044826219
0.709884911506533
0.634023778642186
0.458957804587173
0.541324947349551
0.450834061791692
];
full_aipsting = [
      0.7905    0.3488    0.5882    1.0000    1.0000    0.5603    1.0000    0.2168    0.9773    1.0000
      0.7905    0.9482    0.9188    0.5705    0.2865    0.4267    0.6833    1.0000    1.0000    0.4744
      0.7230    1.0000    1.0000    0.5577    0.3021    0.3879    0.5113    0.4172    0.4602    0.9808
      0.3716    0.7684    0.4622    0.5577    0.4844    0.9310         0    0.4802    0.5341         0
      0.4189    0.8801    0.8768    0.3397    0.9167    1.0000         0    0.0606    0.9205         0
      1.0000    0.3597    0.7255    0.6667         0    0.4095         0    0.6084    0.7386         0
]';
error_aip = 0.1*(max(full_aipsting) - min(full_aipsting));

% nematostella experimental data set
stellasting = [0.080278691
0.237643952
0.398757602
0.440352101
0.690318779
0.613852814];
full_stellasting = [
    0	0	0.123809524	0.221518987	0.052941176	0.039772727	0.132231405	0.253968254	0	0	0.058823529
    0	0.033898305	0.666666667	0.367088608	0.447058824	0.034090909	0.619834711	0.148148148	0.297297297	0	0
    0.884615385	0.449152542	0.752380952	0.335443038	0.388235294	0.715909091	0.553719008	0.306878307	0	0	0
    0.153846154	0.525423729	1	0.101265823	0.382352941	0.448863636	0.537190083	0.978835979	0	0.284722222	0.431372549
    0.666666667	0.13559322	0.438095238	0.417721519	0.711764706	1	0.636363636	0.587301587	1	1	1
    1	1	0.419047619	1	1	0	1	1	0	0.333333333	0
    ]';
error_stella = .1* (max(full_stellasting) - min(full_stellasting));

%hours when the data points are measured
timestarvation = [0 4 20 24 48 75]';

figure
hold on
plot(timestarvation, aiptasiasting, 'o-','linewidth',3, 'LineStyle', 'none')
errorbar(timestarvation, aiptasiasting, error_aip, '*-','linewidth',3, 'LineStyle', 'none')
errorbar(timestarvation, stellasting, error_stella, '*-','linewidth',3, 'LineStyle', 'none')
set(gca,'fontsize',20,'linewidth',2);
ylim([0 1])


%data fitting for aiptasia (to find out how the cost changes with state for
%aiptasia
maxperc = 1;
x1 = timestarvation/max(timestarvation)*maxperc;
poly = polyfit(x1, aiptasiasting, 3);
x2 = (1:max(timestarvation))/max(timestarvation)*maxperc;

k = linspace(0,1,100);
spline_aiptasia = poly(1)*k.^3+poly(2)*k.^2+poly(3)*k.^1 + poly(4);

plot(k/maxperc*max(timestarvation), spline_aiptasia,'linewidth',3)
title('fitting the experimental curve for aiptasia')



%% MDP model

% 1: nematostella
% 2: aiptasia

model = 1;

% parameters:
cost_changewithstate = 1;
d     = 100 ;    % starvation state corresponding to death. while d=0 : well fed.
nm    = 40;     % total # nematocytes
g    = 0.9;     % iteration magnitude = discount factor. the closer to 1, the longer the horizon -> never dies !! default 0.9
death = -1;      % me. penalty upon death

% define MDP
s     = 1 : d;                                      % starvation state. 0=well fed. d=dead.
starv = linspace(0,1,length(s));                  % starvation state, normalized between 0 and 1
n     = 0 : nm;  a     = n / nm;               % actions = # nematocysts discharged

% models for cost and probability to catch
alfa  = 1*50/d;
beta  = 0.8;
delta = 2;   
p     = beta*a.*(delta-a);  %.3*atan(n/nm);

%B is how the cost (fitting) changes with state
spline_aiptasia = poly(1)*starv.^3+poly(2)*starv.^2+poly(3)*starv.^1 + poly(4);
B = 2*(2-g)/g * (1-spline_aiptasia)./(1/beta/g-spline_aiptasia.^2);

if cost_changewithstate == 0
    c     =a.*alfa;              % cost of discharge: linear with number nematocytes shot
else
    % alfa = starv;   % linear
    alfa = B*50/d;
    % alfa = -(starv-1).^2 + 1; %parabolic 1
    % alfa = 0.05*starv.^2; %parabolic 2
    % alfa = 1*ones(size(starv)); %constant
    % alfa = -cos(starv*pi)+1;        %cosine
    c  = a'.* alfa;
end

% models for reward
% r = 1-s/d;   % linear
% r     =  -3*(s/d).^4 + 3 ; % parabolic 
% r     = exp(-10*s/d);
r     = 5*cos(s/d*pi/2); % cosine 
% r     = -10*atan(starv-1); % arctan 

r(end) = death;
r_diff = r(3:end) - r(1:end-2);

% analytical results for nematostella
if cost_changewithstate == 0
    a_theory_nem = max(delta/2 - alfa*(g-1)./ (2*beta*r_diff), 0);
    K = (2-g)./alfa/g;
    A = -1/g/beta + 2*K;
    a_theory_aip = (K - sqrt(K.^2-A))*ones(size(starv));
else
    a_theory_nem = max(delta/2 - alfa(2:end-1)*(g-1)./ (2*beta*r_diff), 0);
    K = (2-g)./alfa/g;
    A = -1/g/beta + 2*K;
    a_theory_aip = K - sqrt(K.^2-A);
end

figure
subplot(3,2,1)
if model == 1
    plot(starv,r,'linewidth',2)
elseif model == 2
        line([-1,0],[1,1],'linewidth',2)
        hold on
        line([0,1],[-1,-1],'linewidth',2)
        xlabel('starvation')
        ylabel('reward')
end
title('assumption')

subplot(3,2,3)
plot(a,c,'linewidth',2)
xlabel('# nematocytes discharged')
ylabel('cost of discharge')
title('assumption')

subplot(3,2,5)
plot(a,p,'linewidth',2)
xlabel('# nematocytes discharged')
ylabel('probability to kill prey')
title('assumption')
% value iteration
Vo = -10*ones(size(s)); % initialize

% s = d -> death. transition d-1 -> d with prob 1-p(in), reaward death - c(in)

hold all
delta = 1 ;
theta = .0001 ;
i=0;
while delta > theta
    i = i+1;
    ddd=0;
    Q=zeros(length(n),length(s));
    for in = 1:length(a)
        if (model==1)
            Q(in, 1:end-1) = p(in)*(r([1 1:end-2])-c(in) + g*Vo([1 1:end-2])) + (1-p(in))*(r(2:end)-c(in)+g*Vo(2:end));
            if cost_changewithstate == 1
                Q(in,1:end-1) = p(in)*(r([1 1:end-2])-c(in, [1 1:end-2]) + g*Vo([1 1:end-2])) + (1-p(in))*(r(2:end)-c(in,2:end)+g*Vo(2:end)); % r = s-s'
            end            
            Q(in,end)     = 0;
        elseif (model==2)
            Q(in,1:end-1) = p(in)*(1-c(in)+ g*Vo(1:end-1)) + (1-p(in))*(-c(in)+death); % r = s-s'
            if cost_changewithstate == 1
                Q(in,1:end-1) = p(in)*(1-c(in,1:end-1)+ g*Vo(1:end-1)) + (1-p(in))*(-c(in,1:end-1)+death); % r = s-s'
            end
                Q(in,end)     =  death - c(in); % absorbing state
        end
    end

    if model == 1
        [V numeric_nematostella] = max(Q);
    elseif model == 2
        [V numeric_aiptasia] = max(Q);
    end
    delta = max(abs(V-Vo));

    Vo = V;

    subplot(3,2,2), plot(starv,V,'linewidth',2)
    hold all
    xlabel('starvation')
    ylabel('expected future reward')
    title(['prediction @ iteration ',num2str(i)])

    subplot(3,2,6), hold all;
    plot(i,ddd,'o-')
    xlabel('iteration')
    ylabel('max update V')

    drawnow
end

if model == 1
    subplot(3,2,4), plot(starv,a(numeric_nematostella),'o-','linewidth',2)
    ylim([0 1])
elseif model == 2
    subplot(3,2,4), plot(starv,a(numeric_aiptasia),'o-','linewidth',2)
    ylim([0 1])
end
hold all
xlabel('starvation')
ylabel('optimal # discharged nematocytes')
title(['prediction @ iteration ',num2str(i)])


figure
if model == 1
    plot(starv(2:2:end-1), a_theory_nem(1:2:end),'linewidth',2)
    hold all
    plot(starv(1:2:end), a(numeric_nematostella(1:2:end)),'o','linewidth',2)
elseif model == 2
    plot(starv(1:2:end-1), a_theory_aip(1:2:end),'linewidth',2)
    hold all
    plot(starv(1:2:end), a(numeric_aiptasia(1:2:end)),'o','linewidth',2)
end
set(gca,'fontsize',20,'linewidth',2);
legend("Theoretical", "Numerical")
ylim([0 1])
xticks([0 0.5 1])
yticks([0 0.5 1])
pbaspect([2 1 1])

%% Experimental comparison

figure
errorbar(timestarvation, aiptasiasting, error_aip, 'o-','linewidth',3, 'LineStyle', 'none')
hold on
errorbar(timestarvation, stellasting, error_stella, '*-','linewidth',3, 'LineStyle', 'none')
set(gca,'fontsize',18,'linewidth',2);
ylim([0 1])

y1 = a(numeric_aiptasia(1:end/2));
y_aip = y1(1:end*maxperc);
y2 = a(numeric_nematostella(1:end/2));
% y2 = a_steady(1:2:end);
y_nem = y2(1:end*maxperc);
hold on
plot(linspace(0,75,length(y_aip)), y_aip,'linewidth',3)
plot(linspace(0,75,length(y_nem)), y_nem,'linewidth',3)
legend("experiment aiptasia", "experiment nematostella", "numerical aiptasia", "numerical nematostella")




