%--------------------------------------------------------------------------
% function AMPLITUDEVARESTIMATION calculates the maximum likelihood
% estimates of trial amplitudes, basis coefficients, noise covariance
% matrix, as well as the log-likelihood value associated with given
% spatiotemporal bases.
% 
% inputs: 
%         1) Data - has dimension of NxTxJ where N is the number of (good) sensor channels, T
%         is the number of time samples (covering time range of interest), and J is the number of trials
%         2) SpatialBases - has dimension of NxP, the P columns are spatial
%         basis vectors
%         3) TemporalBases - has dimension of TxL, the L columns are temporal
%         basis vectors
%
% outputs: 
%         CR-based, LDSR-constrained, and IR-constrained ML solutions are
%         given in 3 structures, CRsoln, LDSRsoln, and IRsoln, respectively. 
%         The fields are estimates of 
%           1) trial amplitude series (.amplitude) [1xJ] as a function of
%           trial index. For CR solution, the amplitude is a scalar
%           constant since CR assumes amplitude is constant across trials.
%           2) basis coefficients (.basiscoeff) [PxL] 
%           3) noise covariance matrix (.Rn) [NxN] 
%           4) log-likelihood value (.likelihood) 
%           5) signal waveform shape (.waveform) [NxT] 
%           6) state memory (.statememory), only for LDSR solution
%           7) driving noise variance (.drivingnoisevar), only for LDSR
%           solution
%           8) prior mean (.priormean), only for IR solution
%           9) prior variance (.priorvar), only for IR solution
% 
% usage: [LDSRsoln, IRsoln, CRsoln]  = amplitudeVarEstimation(Data,SpatialBases,TemporalBases)
% 
% Author: Tulaya Limpiti
% Last updated: 05/24/09
%--------------------------------------------------------------------------

function [LDSRsoln, IRsoln, CRsoln] = amplitudeVarEstimation(Data,SpatialBases,TemporalBases)

inverr = 1e-9;
J = size(Data,3)                      % number of trials
N = size(Data,1)                      % number of sensors
T = size(Data,2)                      % number of time samples
L = size(TemporalBases,2)

%--------------------------------------------------------------------------
% reformat data Y into different dimensions:
%--------------------------------------------------------------------------
% 1) y_j(t) = vec(Y_j) = amp_j*(c kron u)*b + vec(N_j)
%    and b = vec(B);

for kk = 1:J
    tmp = squeeze(Data(:,:,kk));
    Datavec(:,kk) = tmp(:);
end

% 2) concatenate all trials to form an NxJT matrix

Datastack = squeeze(Data(:,:,1));
for k = 2:J
    Datastack = [Datastack, squeeze(Data(:,:,k))];
end

%--------------------------------------------------------------------------
% Maximum likelihood (ML) estimation
%--------------------------------------------------------------------------

%----------------------------------------------------------------------
% Constant-response based ML solutions
%----------------------------------------------------------------------

CR = CRbased_MLsoln(Data,Datastack,SpatialBases,TemporalBases,inverr);

%----------------------------------------------------------------------
% ML solutions with LDSR constraint on trial amplitudes
%----------------------------------------------------------------------

% initialize the ECM algorithm with CR-based solutions
initLDSRf.Rn = CR.Rn;
initLDSRf.Q = CR.amp^2;
initLDSRf.A = CR.amp;
initLDSRf.B = CR.basiscoeff;
initLDSRf.x0 = CR.amp;
initLDSRf.P0 = CR.amp^2;

LDSR = LDSRbased_MLsoln(Data,Datastack,Datavec,initLDSRf,SpatialBases,TemporalBases,inverr);


%----------------------------------------------------------------------
% ML solutions with IR constraint on trial amplitudes
%----------------------------------------------------------------------

% initialize the ECM algorithm with CR-based solutions
initIRf.Rn = CR.Rn;
initIRf.xmu = CR.amp;
initIRf.xvar = CR.amp^2;
initIRf.B = CR.basiscoeff;

IR = IRbased_MLsoln(Data,Datastack,Datavec,initIRf,SpatialBases,TemporalBases,inverr);


%----------------------------------------------------------------------
% formatting output structures
%----------------------------------------------------------------------

% CR-based solutions
CRsoln.amplitude = CR.amp;
CRsoln.Rn = CR.Rn; 
CRsoln.basiscoeff = CR.basiscoeff; 
CRsoln.likelihood = CR.likelihood;
CRsoln.waveform = SpatialBases*CR.basiscoeff*TemporalBases';

% LDSR-constrained solutions
LDSRsoln.amplitude = LDSR.amp(2:end);
LDSRsoln.Rn = LDSR.Rn; 
LDSRsoln.basiscoeff = LDSR.basiscoeff; 
LDSRsoln.likelihood = LDSR.likelihood;
LDSRsoln.waveform = SpatialBases*LDSR.basiscoeff*TemporalBases';
LDSRsoln.statememory = LDSR.statememory;
LDSRsoln.drivingnoisevar = LDSR.drivingvar;

% IR-constrained solutions
IRsoln.amplitude = IR.amp;
IRsoln.Rn = IR.Rn; 
IRsoln.basiscoeff = IR.basiscoeff; 
IRsoln.likelihood = IR.likelihood;
IRsoln.waveform = SpatialBases*IR.basiscoeff*TemporalBases';
IRsoln.priormean = IR.priormean;
IRsoln.priorvar = IR.priorvar;