function bbox_reg = spp_train_bbox_regressor(imdb, roidb, spp_model, varargin)
% bbox_reg = spp_train_bbox_regressor(imdb, roidb, spp_model, varargin)
%   Trains a bounding box regressor on the image database imdb
%   for use with the SPP model spp_model. The regressor is trained
%   using ridge regression.
%
%   Keys that can be passed in:
%
%   min_overlap     Proposal boxes with this much overlap or more are used
%   layer           The CNN layer features to regress from (either 5, 6 or 7)
%   lambda          The regularization hyperparameter in ridge regression
%   robust          Throw away examples with loss in the top [robust]-quantile
%   binarize        Binarize features or leave as real values >= 0
%
% Adapted from spp code written by Ross Girshick
% AUTORIGHTS
% ---------------------------------------------------------
% Copyright (c) 2014, Shaoqing Ren
% 
% This file is part of the SPP code and is available 
% under the terms of the Simplified BSD License provided in 
% LICENSE. Please retain this notice and LICENSE if you use 
% this file (or any portion of it) in your project.
% ---------------------------------------------------------
% Copyright (c) 2014, Ross Girshick
% 
% This file is part of the R-CNN code and is available 
% under the terms of the Simplified BSD License provided in 
% LICENSE. Please retain this notice and LICENSE if you use 
% this file (or any portion of it) in your project.
% ---------------------------------------------------------

ip = inputParser;
ip.addRequired('imdb',       @iscell);
ip.addRequired('spp_model',  @isstruct);
ip.addParamValue('min_overlap', 0.5,   @isscalar);
ip.addParamValue('layer',       5,     @isscalar);
ip.addParamValue('lambda',      4000,  @isscalar);
ip.addParamValue('robust',      0,     @isscalar);
ip.addParamValue('binarize',    false, @islogical);

ip.parse(imdb, spp_model, varargin{:});
opts = ip.Results;
opts = rmfield(opts, 'spp_model');
opts = rmfield(opts, 'imdb');
opts.cache_name = spp_model.cache_name;

fprintf('\n\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n');
fprintf('Training options:\n');
disp(opts);
fprintf('~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n');

imdbs_name = cell2mat(cellfun(@(x) x.name, imdb, 'UniformOutput', false));
conf = spp_config('sub_dir', fullfile(opts.cache_name, imdbs_name));
clss = spp_model.classes;
num_clss = length(clss);

% ------------------------------------------------------------------------
% Get the average norm of the features
opts.feat_norm_mean = spp_feature_stats(imdb, roidb, opts.layer, spp_model);
fprintf('average norm = %.3f\n', opts.feat_norm_mean);
% ------------------------------------------------------------------------

% ------------------------------------------------------------------------
% Get all positive examples
imdbs_name = cell2mat(cellfun(@(x) x.name, imdb, 'UniformOutput', false));
save_file = sprintf('./cachedir/%s/%s/bbox_regressor_XY_layer_%d_overlap_%.5g.mat', ...
                    spp_model.cache_name, imdbs_name, opts.layer, opts.min_overlap);
try
  load(save_file);
  fprintf('Loaded saved positives from ground truth boxes\n');
catch
  [X, Y, O, C] = get_examples(spp_model, imdb, roidb, opts);
  save(save_file, 'X', 'Y', 'O', 'C', '-v7.3');
end
for i = 1:num_clss
  fprintf('%14s has %6d samples\n', spp_model.classes{i}, length(find(C == i)));
end
X = spp_poolX_to_fcX(X, opts.layer, spp_model, false);
X = spp_scale_features(X, opts.feat_norm_mean);
% ------------------------------------------------------------------------

% use ridge regression solved by cholesky factorization
method = 'ridge_reg_chol';

models = cell(num_clss, 1);
for i = 1:num_clss
  fprintf('Training regressors for class %s (%d/%d)\n', ...
      spp_model.classes{i}, i, num_clss);
  I = find(O > opts.min_overlap & C == i);
  Xi = X(:, I); 
  if opts.binarize
    Xi = single(Xi > 0);
  end
  Yi = Y(I,:); 
  Oi = O(I); 
  Ci = C(I);

  % add bias feature
  Xi = cat(1, Xi, ones(1, size(Xi,2), class(Xi)));

  % Center and decorrelate targets
  mu = mean(Yi);
  Yi = bsxfun(@minus, Yi, mu);
  S = Yi'*Yi / size(Yi,1);
  [V, D] = eig(S);
  D = diag(D);
  T = V*diag(1./sqrt(D+0.001))*V';
  T_inv = V*diag(sqrt(D+0.001))*V';
  Yi = Yi * T;

  models{i}.mu = mu;
  models{i}.T = T;
  models{i}.T_inv = T_inv;

  models{i}.Beta = [ ...
    solve_robust(Xi, Yi(:,1), opts.lambda, method, opts.robust) ...
    solve_robust(Xi, Yi(:,2), opts.lambda, method, opts.robust) ...
    solve_robust(Xi, Yi(:,3), opts.lambda, method, opts.robust) ...
    solve_robust(Xi, Yi(:,4), opts.lambda, method, opts.robust)];
end
bbox_reg.models = models;
bbox_reg.training_opts = opts;
save([conf.cache_dir 'bbox_regressor_final'], 'bbox_reg');


% ------------------------------------------------------------------------
function [X, Y, O, C] = get_examples(spp_model, imdb, roidb, opts)
% ------------------------------------------------------------------------
num_classes = length(spp_model.classes);

cls_counts = zeros(num_classes, 1);
for idb = 1:length(imdb)
    for i = 1:length(imdb{idb}.image_ids)
      tic_toc_print('%s: counting %d/%d\n', ...
                    procid(), i, length(imdb{idb}.image_ids));

      d = roidb{idb}.rois(i);
      [max_ov cls] = max(d.overlap, [], 2);
      sel_ex = find(max_ov >= 0.5);
      cls = cls(sel_ex);
      for j = 1:length(sel_ex)
        cls_counts(cls(j)) = cls_counts(cls(j)) + 1;
      end
    end
end
total = sum(cls_counts);
feat_dim = size(spp_model.cnn.layers(1).weights{1},2);
% features
X = zeros(feat_dim, total, 'single');
% target values
Y = zeros(total, 4, 'single');
% overlap amounts
O = zeros(total, 1, 'single');
% classes
C = zeros(total, 1, 'single');
cur = 1;

for idb = 1:length(imdb)
    for i = 1:length(imdb{idb}.image_ids)
      tic_toc_print('%s: pos features %d/%d\n', ...
                    procid(), i, length(imdb{idb}.image_ids));
      d = roidb{idb}.rois(i);
      d.feat = spp_load_cached_poolX_features(spp_model.spp_pooler, spp_model.feat_cache{idb}, ...
          imdb{idb}.name, imdb{idb}.image_ids{i}, d.boxes);

      sel_gt = find(d.class > 0);
      gt_boxes = d.boxes(sel_gt, :);
      gt_classes = d.class(sel_gt);

      max_ov = max(d.overlap, [], 2);
      sel_ex = find(max_ov >= opts.min_overlap);
      ex_boxes = d.boxes(sel_ex, :);

      X(:, cur+(0:length(sel_ex)-1)) = d.feat(:, sel_ex);

      for j = 1:size(ex_boxes, 1)
        ex_box = ex_boxes(j, :);
        ov = boxoverlap(gt_boxes, ex_box);
        [max_ov, assignment] = max(ov);
        gt_box = gt_boxes(assignment, :);
        cls = gt_classes(assignment);

        src_w = ex_box(3) - ex_box(1) + eps;
        src_h = ex_box(4) - ex_box(2) + eps;
        src_ctr_x = ex_box(1) + 0.5*src_w;
        src_ctr_y = ex_box(2) + 0.5*src_h;

        gt_w = gt_box(3) - gt_box(1) + eps;
        gt_h = gt_box(4) - gt_box(2) + eps;
        gt_ctr_x = gt_box(1) + 0.5*gt_w;
        gt_ctr_y = gt_box(2) + 0.5*gt_h;

        dst_ctr_x = (gt_ctr_x - src_ctr_x) * 1/src_w;
        dst_ctr_y = (gt_ctr_y - src_ctr_y) * 1/src_h;
        dst_scl_w = log(gt_w / src_w);
        dst_scl_h = log(gt_h / src_h);

        target = [dst_ctr_x dst_ctr_y dst_scl_w dst_scl_h];

        if 0
          % debugging visualizations and checks
          im = imread(imdb.image_at(i));
          showboxesc(im, gt_box, 'g', '-');
          showboxesc([], ex_box, 'r', '-');
          hold on;
          plot(gt_ctr_x, gt_ctr_y, 'gd');
          plot(src_ctr_x, src_ctr_y, 'rd');
          hold off;
          fprintf('target = [%.3f %.3f %.3f %.3f]\n', target(1), target(2), target(3), target(4));
          fprintf('cls = %s\n', spp_model.classes{cls});

          % check that we can correctly reconstruct the gt_box from the
          % gold-standard target
          pred_ctr_x = (target(1) * src_w) + src_ctr_x;
          pred_ctr_y = (target(2) * src_h) + src_ctr_y;
          pred_w = exp(target(3)) * src_w;
          pred_h = exp(target(4)) * src_h;
          pred_box = [pred_ctr_x - 0.5*pred_w, pred_ctr_y - 0.5*pred_h, ...
                      pred_ctr_x + 0.5*pred_w, pred_ctr_y + 0.5*pred_h];
          disp(pred_box);
          disp(gt_box);
          assert(sum(abs(pred_box - gt_box)) < 0.0001);

          pause;
        end

        assert(cur <= total);
        Y(cur, :) = target;
        O(cur)    = max_ov;
        C(cur)    = cls;
        cur = cur + 1;
      end
    end
end

% ------------------------------------------------------------------------
function [x, losses] = solve_robust(A, y, lambda, method, qtile)
% ------------------------------------------------------------------------
[x, losses] = solve(A, y, lambda, method);
fprintf('loss = %.3f\n', mean(losses));
if qtile > 0
  thresh = quantile(losses, 1-qtile);
  I = find(losses < thresh);
  [x, losses] = solve(A(I,:), y(I), lambda, method);
  fprintf('loss (robust) = %.3f\n', mean(losses));
end


% ------------------------------------------------------------------------
function [x, losses] = solve(A, y, lambda, method)
% ------------------------------------------------------------------------

%tic;
switch method
case 'ridge_reg_chol'
  % solve for x in min_x ||Ax - y||^2 + lambda*||x||^2
  %
  % solve (A'A + lambdaI)x = A'y for x using cholesky factorization
  % R'R = (A'A + lambdaI)
  % R'z = A'y  :  solve for z  =>  R'Rx = R'z  =>  Rx = z
  % Rx = z     :  solve for x
  R = chol(A*A' + lambda*eye(size(A,1)));
  z = R' \ (A*y);
  x = R \ z;
case 'ridge_reg_inv'
  % solve for x in min_x ||Ax - y||^2 + lambda*||x||^2
  x = inv(A*A' + lambda*eye(size(A,1)))*A*y;
case 'ls_mldivide'
  % solve for x in min_x ||Ax - y||^2
  if lambda > 0
    warning('ignoring lambda; no regularization used');
  end
  x = A'\y;
end
%toc;
losses = 0.5 * (A'*x - y).^2;
