%------------------------------------------------------------------------
% NAME:     qpoem
%
%           x
%
% FORMAT:   [X,G,y_corrected,y_fitted] = qpoem( Q, y )
%             or
%           Dy = qpoem( Q ) 
%
% OUT:      x
% IN:       -
% OPTIONAL: -
% TEMPLATE: -
%------------------------------------------------------------------------

% HISTORY: 2001.01.04  Started by by Patrick Eriksson.


function [X,G,y_corrected,y_fitted,Xiter] = qpoem( Q, y )



%### Determine calculation options ############################################
%
%= Set 0 as first guess for the options
%
nonlin     = 0;
return_Dy  = 0;  
do_post    = 0;
recalc_abs = 0;
logspecies = 0;
do_xiter   = 0;
%
out(1,1);
%
if (nargin<2) | isempty(y)
  %
  return_Dy  = 1;  
  %
  out(2,'ONLY CALCULATING Dy');
  %
else
  %
  nonlin     = Q.OEM_NONLIN_ON;
  %
  if nargout > 1
    do_post = 1;
  end
  %
  if nonlin 
    %
    if Q.SPECIES_POSITIVE_ON
      logspecies = 1;
    end
    %
    if Q.TEMPERATURE_DO == 3
      recalc_abs = 1;
    end
  end
  %
  out(2,'STARTING RETRIEVAL');
end
%
if nargout > 4
  do_xiter = 1;
end



%### Some printing ############################################################
out(2,0);
if nonlin
  if Q.OEM_GA_START_VALUE == 0
    out(2,'Method         : OEM, Gauss-Newton');
  else
    out(2,'Method         : OEM, Marquardt-Levenberg');
  end
  out(2,sprintf('Start gamma    : %.1f', Q.OEM_GA_START_VALUE ));
  out(2,sprintf('Factor OK      : %.1f', Q.OEM_GA_FAC_WHEN_OK ));
  out(2,sprintf('Factor not OK  : %.1f', Q.OEM_GA_FAC_WHEN_NOT_OK ));
  out(2,sprintf('Lower treshold : %.2f', Q.OEM_GA_LOWER_TRESHOLD ));
  out(2,sprintf('Stop value     : %.3f', Q.OEM_STOP ));
  %
  if recalc_abs
    out(2,'   Absorption will be recalculated for each iteration.');
  else
    out(2,'   Absorption will be scaled by species amounts.');
  end
else
  out(2,'Method : Linear OEM');
end
%
if ~return_Dy
  if do_post
    out(2,'   Spectrum, cost etc. will be calculated after retrieval.');
  else
    out(2,'   No post-retrieval calculations will be performed.');
  end
end



%### Create temporary dir #####################################################
%
tmpdir    = temporary_directory( Q.TMP_AREA );


%### Run ARTS for first iteration #############################################
%
do_nonlin = nonlin | do_post;
[ fy, H, Hd, Kx, kx_names, kx_index, kx_aux, cfile, basename ] = ...
                               qpoem_iter1( Q, tmpdir, do_nonlin, recalc_abs );
%
%= Some sizes
xa     = kx_aux(:,2);
nrq    = size( kx_names, 1 );
nx     = length( xa );



%### Load the covariance matrices #############################################
%
load( [Q.OUT,'.sxinv'], '-mat' );
load( [Q.OUT,'.dxinv'], '-mat' );
load( [Q.OUT,'.seinv'], '-mat' );
%
if size(Kx,2) ~= size(Sxinv,1)
  error('Sizes of Kx and Sx do not match (recalculation needed?).')
end 



%### Return here if only Dy shall be calculated ###############################
%
if  return_Dy
  %
  KtSeinv = Kx'*Seinv;
  X       = inv( Sxinv + KtSeinv*Kx ) * KtSeinv;
  %
  return   %>>>>> Exit >>>>>
end



%### Special stuff: non-linear ################################################
if nonlin

  gamma  = Q.OEM_GA_START_VALUE;

  %=== Build cfile for calculation of Kx
  template = which( 'oem_iter2_kx.tmplt' );
  %
  QE            = [];
  QE.RECALC_ABS = recalc_abs;
  %
  qtool( Q, tmpdir, template, QE );
  cfile_kx  = [ cfile, '-FOR_KX' ];
  copyfile( cfile, cfile_kx );

  %=== Index for variables where log will be used
  %
  if logspecies
    ilogsp = [];
    for i = 1:nrq
      group = qp_kinfo( kx_names{i} );
      if strcmp( group, 'Species' )
        ilogsp = [ ilogsp; (kx_index(i,1):kx_index(i,2))' ];
      end
    end
    %
    xa(ilogsp) = log( xa(ilogsp) );
    %
  end


%### Linear ###################################################################
else
  %
  gamma = 0;
  %
end



%### Special stuff: non-linear or post calculations ###########################
if nonlin | do_post

  %=== Build cfile for calculation of FY 
  template = which( 'oem_iter2_y.tmplt' );
  %
  QE            = [];
  QE.RECALC_ABS = recalc_abs;
  %
  qtool( Q, tmpdir, template, QE );
  cfile_y  = [ cfile, '-FOR_Y' ];
  copyfile( cfile, cfile_y );

  %=== Start cost
  cost_x = 0;
  a      = y - fy;
  cost_y = a' * Seinv * a;
  cost   = cost_x + cost_y;

  %=== Printing
  out(2,0);
  out(1,sprintf(' Iter              Total    Profile    Spectrum   Conver-'));
  out(1,sprintf('  nr    Gamma      cost      cost       cost       gence '));
  out(1,0);
  out(1,sprintf('  %2.0f      -      %.2e     0.00     %.2e       -',...
                                                         0, cost_y, cost_y ));

  %=== Load needed variables
  %
  %= Absorption
  if ~recalc_abs
    Abs  = read_artsvar( basename, 'abs');
    Absp = read_artsvar( basename, 'abs_per_tg');
  else
    Abs  = [];
    Absp = [];
  end
  %
  %= Zenith angles
  if Q.POINTING_DO == 3
    za_pencil = read_artsvar( basename, 'za_pencil');
  else
    za_pencil = [];
  end
  %
  %= Temparature profile
%  if Q.TEMPERATURE_DO == 3
%    t_abs = read_artsvar( basename, 't_abs');
%  else
%    t_abs = [];
%  end
end



%### Read ARTS variables always needed ########################################
%
p_abs = read_artsvar( basename, 'p_abs');
vmrs  = read_artsvar( basename, 'vmrs');




%### Loop until convergence ###################################################
%
%      ||          ||          ||          ||          ||          ||       
%      ||          ||          ||          ||          ||          ||       
%      ||          ||          ||          ||          ||          ||       
%      ||          ||          ||          ||          ||          ||       
%    \ || /      \ || /      \ || /      \ || /      \ || /      \ || /
%     \||/        \||/        \||/        \||/        \||/        \||/
%      \/          \/          \/          \/          \/          \/
%
%
iteration   = 0;
if nonlin
  max_iter  = Q.OEM_MAX_ITER;
else
  max_iter  = 1;
end
%
x           = xa;
y_corrected = y;
%
if do_xiter
  Xiter     = [];
end
%
G           = [];
G.converged = 0;
%
while ~G.converged & ( iteration < max_iter )

  iteration = iteration + 1;

  %=== Calculate new KX if iteration > 1
  %
  % [!! Check this later
  % The Kx to apply for the the log-retrieval of species is Kx*x (see Eq. 45
  % of Paper A in my thesis), but this is achieved automatically as Kx is
  % calculated in fractions of x (and not xa). ]
  %
  if iteration > 1
    %
    fprintf('ERROR, ERROR! I have discovered a bug when doing non-linear\n')
    fprintf('inversions and including parameters beside species.\n');
    fprintf('No time to fix this now. Patrick Eriksson (2001-08-31).\n');
    error('See above');
    %
    copyfile( cfile_kx, cfile );
    qp_arts( Q, cfile ); 
    %
    Kx = qp_Kx( Q, basename, H, Hd, fy );
    %
    if ~logspecies % WRONG, WRONG
      Kx = Kx ./ (ones(size(Kx,1),1)*x');
    end
    %
  end


  %=== Make inversion
  %
  % For OEM-Gauss-Newton, the error covariance matrix (Sd) is calculated 
  % as a step in the solution to have the error estimation ready.
  % For OEM-ML, Gaussian elimination is used as this is faster (the last 
  % iteration is always OEM-GN.
  %
  KtSeinv = Kx'*Seinv;
  %
  ok_iter = 0;
  %
  while ~ok_iter

    if ~nonlin
      %= If linear, set ok flag for iteration to true
      ok_iter     = 1;
    end

    %= A vector needed in different places
    if iteration == 1
      wdy = KtSeinv * ( y_corrected - fy );
    else
      wdy = KtSeinv * ( y_corrected - fy ) - Sxinv * ( x - xa );
    end

    %= New x
    if gamma == 0
      Sd = inv( Sxinv + KtSeinv*Kx ); 
      x_new = x + Sd * wdy;
    else
      dx  = (Sxinv + KtSeinv*Kx + gamma*Dxinv) \ wdy;
      x_new = x + dx;
    end

    if nonlin | do_post

      %= Map X onto absorption etc. yc is correction term for the spectrum.
      yc = map_x( x_new, basename, Kx, kx_names, kx_index, kx_aux, ...
                   recalc_abs, logspecies, p_abs, vmrs, Abs, Absp, za_pencil );
      %
      y_corrected = y - yc;
  
      %= New fy
      copyfile( cfile_y, cfile );
      qp_arts( Q, cfile ); 
      fy_new = read_artsvar( basename, 'y' );
      fy_new = H*fy_new;

      %= New cost
      a        = x_new - xa;
      cost_x   = a' * Sxinv * a;
      a        = y_corrected - fy_new;
      cost_y   = a' * Seinv * a;
      cost_new = cost_x + cost_y;
    end

    if nonlin

      %= Is OK?
      if cost_new < cost
        ok_iter = 1;
      else
        out(1,sprintf('FAILED %.1e   %.2e', gamma, cost_new ));
        if gamma > 0
          gamma = gamma * Q.OEM_GA_FAC_WHEN_NOT_OK;
        else
          gamma = Q.OEM_GA_LOWER_TRESHOLD * Q.OEM_GA_FAC_WHEN_NOT_OK;
        end
      end
    end
  end % while


  if nonlin | do_post  % Note that gamma=0 when do_post=1

    %=== If gamma=0, check convergence
    if gamma==0 
      %
      conv = (( x_new - x )'*wdy) / nx;
      %
      out(1,sprintf('  %2.0f   %.1e   %.2e   %.2e   %.2e   %.2e',...
			   iteration, gamma, cost_new, cost_x, cost_y, conv ));
      %
      if nonlin & ( conv < Q.OEM_STOP )
        G.converged = 1;
      end 
      
    %=== Or, update gamma
    else
      out(1,sprintf('  %2.0f   %.1e   %.2e   %.2e   %.2e       -',...
				 iteration, gamma, cost_new, cost_x, cost_y ));
      %
      gamma = gamma / Q.OEM_GA_FAC_WHEN_OK;
      %
      if gamma < Q.OEM_GA_LOWER_TRESHOLD
	gamma = 0;
      end
    end

    fy   = fy_new;
    cost = cost_new;

  end

  x = x_new;
  
  if do_xiter
    Xiter = [Xiter, x];
  end

end

if logspecies
  x(ilogsp) = exp( x(ilogsp) );
end

%
%      /\          /\          /\          /\          /\          /\
%     /||\        /||\        /||\        /||\        /||\        /||\
%    / || \      / || \      / || \      / || \      / || \      / || \
%      ||          ||          ||          ||          ||          ||       
%      ||          ||          ||          ||          ||          ||       
%      ||          ||          ||          ||          ||          ||       
%      ||          ||          ||          ||          ||          ||       
%##############################################################################





%### Fill X structure #########################################################
%
X        = cell( nrq, 1 );
ispecies = 0;
%
%= Read needed ARTS variables that can been changed
z_abs = read_artsvar( basename ,'z_abs' );
%
for i = 1:nrq

  [group,name] = qp_kinfo(  kx_names{i} );
  X{i}.name = name;       

  ind = kx_index(i,1):kx_index(i,2);

  %=== Species
  %
  if strcmp( group, 'Species' )
    %
    ispecies     = ispecies + 1;
    X{i}.p       = kx_aux(ind,1);
    X{i}.z       = interpp( p_abs, z_abs, X{i}.p );
    X{i}.apriori = interpp( p_abs, vmrs{ispecies}, X{i}.p );
    X{i}.x       = x(ind) .* X{i}.apriori;
    X{i}.xfrac   = x(ind);

  %=== Sensor scalars such as pointing off-set
  %
  elseif strcmp( group, 'Sensor scalar' )
    %
    X{i}.apriori = 0;
    X{i}.x       = x(ind);

  %=== Temperature and continuum fit
  %
  elseif strcmp( group, 'Temperature' ) | strcmp( group, 'Continuum' )
    %
    X{i}.p       = kx_aux(ind,1);
    X{i}.z       = interpp( p_abs, z_abs, X{i}.p );
    X{i}.apriori = kx_aux(ind,2);
    X{i}.x       = x(ind);

  %=== Polynomial baseline fits
  %
  elseif strcmp( group, 'Baseline' )
    %
    X{i}.p       = kx_aux(ind,1);
    X{i}.z       = kx_aux(ind,1);
    X{i}.apriori = kx_aux(ind,2);
    X{i}.x       = x(ind);

  else
    error(sprintf( 'Unknown retrieval identity: %s', X{i}.name ));
  end
end



%### Other return varaibles ##################################################
%
if do_post
  %
  y_fitted = fy;
  %
  G.COST   = cost;
  G.COST_X = cost_x;
  G.COST_Y = cost_y;
end


%=== Delete the temporary directory
delete_tmp_dir( Q.TMP_AREA, tmpdir );



out(1,-1);
if out(1)
  fprintf('\n');
end
%_____________________________________________________________________________





%#############################################################################
%###
%### Sub-functions
%###
%#############################################################################


function yc = map_x( x, basename, Kx, kx_names, kx_index, kx_aux, recalc_abs, logspecies, p_abs, vmrs, Abs, Absp, za )


  yc       = 0;
  ispecies = 0;
  nf       = size( Abs, 1 );
  

  for i = 1:size(kx_names,1)
  
    [group,name,par1] = qp_kinfo( kx_names{i} );

    ind = kx_index(i,1):kx_index(i,2);
  
    %=== Species
    %
    if strcmp( group, 'Species' )
      %
      ispecies = ispecies + 1;
      %
      xp = interpp( kx_aux(ind,1), x(ind), p_abs );
      %
      if logspecies
        xp = exp( xp );
      end
      %
      vmrs{ispecies} = vmrs{ispecies} .* xp; 
      %
      if ~recalc_abs
        %
        dAbs = Absp{ispecies} .* ( ones(nf,1) * (xp-1)' );
        Absp{ispecies} = Absp{ispecies} + dAbs; 
        Abs = Abs + dAbs;
        %
      end


    %=== Temperature
    %
    elseif strcmp( group, 'Temperature' )
      %
      t_abs = interpp( kx_aux(ind,1), x(ind), p_abs );
      write_artsvar( basename, 't_abs', t_abs, 'b' );


    %=== Pointing off-set
    %
    elseif strcmp( name, 'Pointing: off-set' )
      %
      za = za + x(ind);
      write_artsvar( basename, 'za_pencil', za, 'b' );


    %=== Polynomial baseline fits
    %
    elseif strncmp( group, 'Baseline', 8 )
      %
      yc = yc + Kx(:,ind) * x(ind);
      

    else
       error(sprintf( 'Unknown retrieval identity: %s', name ));
    end
  end

  %=== Save profile/absorption data
  %
  write_artsvar( basename, 'vmrs', vmrs, 'b' );
  %
  if ~recalc_abs
    write_artsvar( basename, 'abs', Abs, 'b' );
    write_artsvar( basename, 'abs_per_tg', Absp, 'b' );
  end
  
return