addpath('../../fdaM')
addpath('SSE')
addpath('exptrans')
addpath('logtrans')

% The Chemostat equations represent a four-species Chemostat plus the resource 
% of Nitrogen_ There are two species of Algae with varying defenses against
% Rotifers_ The Rotifers themselves are divided into two class -- breeding
% and senescent, although these two are very tightly coupled_ 
%
% A full description of these equations can be found in the user manual_ 
% The five state variables for the equations are
% 
% N - nitrogen content in the Chemostat
% C1 - Algal type 1
% C2 - Algal type 2 
% B - Breeding Rotifers
% S - Senescent Rotifers
%
% The system has 16 parameters, also described in the user manual_ Notable
% features include that only the sums C1+C2 and B+S can be observed_ Further, 
% an unknown fraction of each is counted at each time_ This requires us to 
% set up a model for the observation process along with the ODE_ 


% First we generate some data 

fid  = fopen('ChemoData.txt','rt');
ChemoData = reshape(fscanf(fid,'%f'),2,101)';
fclose(fid);

%%%%% NEED TO ALSO FIX ChemoTime and ChemoPars


% The first two of these parameters give the fractions of Algae and Rotifers 
% that are counted_ The remaining parameters are all positive and using 
% their logged values is helpful_ 

logpars = [ChemoPars(1:2),log(ChemoPars(3:16))];

% Parameters 'p1' and 'p2' represent relative palatability of the two algal
% clones, as such only one can be estimated and we fix p2 = 0_ 

active = [1:2,5,7:16];     

% We'll choose a fairly large value of lambda_ 

lambda = 10000*ones(5,1); 

% We need some basis functions

rr = [min(ChemoTime) max(ChemoTime)];
knots = rr(1):0.5:rr(2);

%%%%%% UP TO HERE


bbasis = create_bspline_basis(rr,norder,knots);

% We now need to define two functions_ First, chemo_fun defines a map from the
% trajectory to its derivative

function r = chemo_funtimes, y, p, more = NULL) 

    r = y;
    p = exp(p);
    Q = p("p1") * y(:, "C1") + p("p2") * y(:, "C2");
    Qs = Q * exp(100 * (Q - p("Qstar")))/(1 + exp(100 * (Q - ...
        p("Qstar")))) + p("Qstar")/(1 + exp(100 * (Q - p("Qstar"))));
    r(:, "N") = p("delta") * (p("NI") - y(:, "N")) - p("rho") * ...
        y(:, "C1") * y(:, "N")/(p("KC1") + y(:, "N")) - p("rho") * ...
        y(:, "C2") * y(:, "N")/(p("KC2") + y(:, "N"));
    r(:, "C1") = y(:, "C1") * (p("XC") * p("rho") * y(:, "N")/(p("KC1") + ...
        y(:, "N")) - p("p1") * p("G") * (y(:, "B") + y(:, "S"))/(p("KB") + ...
        Qs) - p("delta"));
    r(:, "C2") = y(:, "C2") * (p("XC") * p("rho") * y(:, "N")/(p("KC2") + ...
        y(:, "N")) - p("p2") * p("G") * (y(:, "B") + y(:, "S"))/(p("KB") + ...
        Qs) - p("delta"));
    r(:, "B") = y(:, "B") * (p("XB") * p("G") * Q/(p("KB") + Qs) - ...
        (p("delta") + p("m") + p("lambda")));
    r(:, "S") = p("lambda") * y(:, "B") - (p("delta") + p("m")) * ...
        y(:, "S");        
end

% Second, we need a map from the trajectory to the observations_ We'll use 
% log observations in this case since the observations are positive and this
% helps to stabilize scales_ 

function r = chemo_obs(t,y,p,more)

  Algae = p(1)*(y(:,'C1')+y(:,'C2'));
  Rotifers = p(2)*(y(:,'B')+y(:,'S'));
  
  r = [Algae, Rotifers];

end


%%% From here we will set-up the objects needed to run profiling_ Below we give
% LS_setup the parameters, maps from trajectories to derivatives and to
% observations, along with data, times, basis functions, variable names and
% trade-off parameter lambda_ 
%
% Additionally, posproc and poslik indicate that we are modelling the log of the
% trajectory, so that we should exponentiate before mapping to derivatives and
% before mapping to observations_ 

out = LS_setup(pars = logpars, fn = chemo_fun, basisvals=bbasis,lambda=lambda,
               data = log(ChemoData),times=ChemoTime,posproc=TRUE,poslik=TRUE,
               names = ChemoVarnames,likfn=chemo_obs);

lik = out.lik;
proc = out.proc;

% Because we don't have direct observations of any state, we'll use a starting
% smooth obtained from generating some ODE solutions_ 

y0 = log(c(2,0.1,0.4,0.2,0.1));
names(y0) = ChemoVarnames;

odetraj = IntegrateForward(y0,ChemoTime,logpars,proc);

DEfd   = smooth_basis(odetraj.times,odetraj.states,fdPar(bbasis,int2Lfd(2),1e-6);
coefs0 = DEfd.fd.coef;

obsvals = lik.more.fn(ChemoTime,odetraj.states,ChemoPars,lik.more.more);

%% We can now call the inner optimization -- this just tries to fined the 
% coefficients that both fit the data and the differential equation well_ 

res = inneropt(coefs=coefs0, pars=logpars, times=ChemoTime, data=log(ChemoData),
               lik=lik, proc=proc, in_meth='optim',control_in=list(trace=2));

% We'll form the trajectory and also the appropriate sum of exponentiated
% states to compare to the data_ 

traj    = lik.bvals %*% res.coefs
obstraj = lik.more.fn(ChemoTime,traj,logpars,lik.more.more)

% Plot these against the data

X11()
par(mfrow=c(2,1))
plot(obstraj(:,1),type='l',ylab='Chlamy',xlab='',cex_lab=1_5,cex_axis=1_5)
points(log(ChemoData(:,1)))
plot(obstraj(:,2),type='l',ylab='Brachionus',xlab='days',cex_lab=1_5,cex_axis=1_5)
points(log(ChemoData(:,2)))

% The function outeropt now allows parameter estimation to take place through 
% profiling procedure_ 

res2 = outeropt(pars=logpars,times=ChemoTime,data=log(ChemoData),coef=res.coefs,
    lik=lik,proc=proc,active=active,in_meth='optim',out_meth='nlminb')

% We'll extract the resulting parameters and coefficients_ 

npars = res2.pars
C = res2.coefs

% And obtain an estimated trajectory and the exponentiated sum to comprare
% to the data_ 

traj = lik.bvals%*%C
ptraj = lik.more.fn(ChemoTime,traj,npars,lik.more.more)

% Lets have a look at how much we changed our parameters on the original 
% scale_ 

new_pars = npars
new_pars(3:16) = exp(new_pars(3:16))

print(ChemoPars)
print(new_pars)
print(new_pars/ChemoPars)

% Now we can produce a set of diagnostic plots_ 

% Firstly, a representation of the trajectory compared to the data_ 

X11()
par(mfrow=c(2,1))
plot(ChemoTime,ptraj(:,1),type='l',ylab='Chlamy',xlab='',cex_lab=1_5)
points(ChemoTime,log(ChemoData(:,1)))
plot(ChemoTime,ptraj(:,2),type='l',ylab='Brachionus',xlab='days',cex_lab=1_5)
points(ChemoTime,log(ChemoData(:,2)))


% Now we'll plot both the derivative of the trajectory and the value of the
% differential equation right hand side at each point_ This represents the 
% fit to the model_ 

traj2 = proc.bvals.bvals%*%C
dtraj2 = proc.bvals.dbvals%*%C

colnames(traj2) = ChemoVarnames
ftraj2 = proc.more.fn(proc2.more.qpts,traj2,npars,proc.more.more)


X11()
par(mfrow=c(5,1),mai=c(0.3,0.6,0.1,0.1))
for(i = 1:5){
  plot(proc.more.qpts,dtraj2(:,i),type='l',xlab='',ylab=ChemoVarnames(i),
    cex_lab=1.5,ylim=c(-0.,5,0.5))
  lines(proc.more.qpts,ftraj2(:,i),col=2,lty=2)
  abline(h=0)
}
legend('topleft',legend=c('Smooth','Model'),lty=1:2,col=1:2,cex=1.2)

% Solving the differential equation from the estiamted initial condition
% of the trajectory allows us to compare the qualitative behavior of 
% our estimate to that of the differential equation_ 

y0 = traj(1,:);
names(y0) = ChemoVarnames;
odetraj = IntegrateForward(y0,ChemoTime,npars,proc);


X11()
par(mfrow=c(2,1))
matplot(ChemoTime,traj,col=1,type='l',lwd=3,cex_lab=1_5,cex_axis=1.5,
    ylab='',cex_main=1.5,main='Reconstructed Trajectories')
legend(x='topright',legend=ChemoVarnames,lwd=3,col=1,lty=1:5)
matplot(ChemoTime,odetraj.states,col=1,type='l',lwd=3,cex_axis=1.5,cex_lab=1.5,
    ylab='',cex_main=1.5,main='ODE Solution')

% We can also compare the pattern of observations predicted by the differential
% equation and that estimated by our methods_ 

otraj = lik.more.fn(ChemoTime,odetraj.states,npars,lik.more.more);

X11()
par(mfrow=c(2,1))
matplot(ChemoTime,ptraj,type='l',lwd=2,xlab='days',cex_lab=1.5,ylab='',
  cex_axis=1.5,cex_main=1.5,main='Predicted Observations -- Smooth')
matplot(ChemoTime,log(ChemoData),add=TRUE,pch = c(1,2))
legend('topright',legend=c('Algae','Rotifers'),pch=1:2,col=1:2)
matplot(ChemoTime,otraj,type='l',lwd=2,xlab='days',cex_lab=1.5,ylab='',
  cex_axis=1_5,cex_main=1.5,main='Predicted Observations -- ODE')
legend('topright',legend=c('Algae','Rotifers'),lty=1:2,col=1:2,lwd=2)