// Copyright 2014 Google Inc. All Rights Reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
// Author: nevena@google.com (Nevena Lazic)

#include "lossmin/losses/inner-product-loss-function.h"
#include <math.h>

namespace lossmin {

float InnerProductLossFunction::LossCurvature(
    const InstanceSet &instances) const {
  float data_curvature = (instances.cwiseProduct(instances) *
                          Weights::Ones(instances.cols())).maxCoeff();
  return curvature_ * data_curvature;
}

void InnerProductLossFunction::PerCoordinateCurvature(
    const InstanceSet &instances, VectorXf *per_coordinate_curvature) const {
  *per_coordinate_curvature =
      VectorXf::Ones(instances.rows()).transpose() *
            instances.cwiseProduct(instances) / instances.rows();
  *per_coordinate_curvature *= curvature_;
}

float InnerProductLossFunction::ExampleLoss(
    const Weights &weights, const InstanceSet &instances,
    const LabelSet &labels, int example) const {
  float inner_product = instances.innerVector(example).dot(weights);
  return InnerProductExampleLoss(inner_product, labels.coeff(example, 0));
}

void InnerProductLossFunction::AddExampleGradient(
    const Weights &weights, const InstanceSet &instances,
    const LabelSet &labels, int example, float weights_scale,
    float example_scale, Weights *gradient) const {
  float inner_product = instances.innerVector(example).dot(weights);
  if (weights_scale != 1.0f) inner_product *= weights_scale;
  float inner_product_gradient =
      InnerProductExampleGradient(inner_product, labels.coeff(example, 0));
  if (example_scale != 1.0f) inner_product_gradient *= example_scale;

  if (synchronous_update()) {
    std::lock_guard<std::mutex> lock(gradient_update_mutex_);
    for (InstanceSet::InnerIterator it(instances, example); it; ++it) {
      gradient->coeffRef(it.index()) += inner_product_gradient * it.value();
    }
  } else {
    for (InstanceSet::InnerIterator it(instances, example); it; ++it) {
      gradient->coeffRef(it.index()) += inner_product_gradient * it.value();
    }
  }
}

void InnerProductLossFunction::ExampleGradient(
    const Weights &weights, const InstanceSet &instances,
    const LabelSet &labels, int example, float weights_scale,
    float example_scale,
    std::vector<std::pair<int, float>> *example_gradient) const {
  float inner_product = instances.innerVector(example).dot(weights);
  if (weights_scale != 1.0f) inner_product *= weights_scale;
  float inner_product_gradient =
      InnerProductExampleGradient(inner_product, labels.coeff(example, 0));
  if (example_scale != 1.0f) inner_product_gradient *= example_scale;

  example_gradient->resize(instances.row(example).nonZeros());
  int i = 0;
  for (InstanceSet::InnerIterator it(instances, example); it; ++it) {
    example_gradient->at(i++) =
        std::make_pair(it.index(), inner_product_gradient * it.value());
  }
}

// Assigns labels to 'instances' given 'weights'.
void InnerProductLossFunction::PredictLabels(
    const Weights &weights, const InstanceSet &instances,
    LabelSet *labels) const {
  // Compute inner products.
  VectorXf inner_products = instances * weights;

  // Assign labels by calling InnerProductAssignLabel coefficientwise on
  // 'inner_products'.
  std::function<float(float)> assign_label_ptr =
      std::bind(&InnerProductLossFunction::InnerProductPredictLabel,
                this, std::placeholders::_1);
  *labels = inner_products.unaryExpr(assign_label_ptr);
}

}  // namespace lossmin