// See www.openfst.org for extensive documentation on this weighted
// finite-state transducer library.
//
// Conversions to/from PowerWeight and SparsePowerWeight.

#ifndef FST_POWER_WEIGHT_MAPPERS_H_
#define FST_POWER_WEIGHT_MAPPERS_H_

namespace fst {

// Converts FromWeight to ToPowerWeight (with one component).
// ToPowerWeight may be PowerWeight or SparsePowerWeight.
template <class FromWeight_, class ToPowerWeight>
class ToPowerWeightMapper {
 public:
  using FromWeight = FromWeight_;
  using ToWeight = ToPowerWeight;
  using Index = typename ToPowerWeight::Index;

  explicit ToPowerWeightMapper(Index index = 0) : index_(index) {}

  ToPowerWeight operator()(const FromWeight &w) const {
    return ToPowerWeight(index_, w.Value());
  }

 private:
  const Index index_;
};

// Converts FromPowerWeight to ToWeight. Uses only one component.
// FromPowerWeight may be PowerWeight or SparsePowerWeight.
template <class FromPowerWeight, class ToWeight_>
class FromPowerWeightMapper {
 public:
  using FromWeight = FromPowerWeight;
  using ToWeight = ToWeight_;
  using Index = typename FromPowerWeight::Index;

  explicit FromPowerWeightMapper(Index index = 0) : index_(index) {}

  ToWeight operator()(const FromPowerWeight &w) const {
    return ToWeight(w.Value(index_));
  }

 private:
  const Index index_;
};

// Projects to one dimension of the weight vector, filling the indices
// with `default_weight`.
// PowerWeightT may be PowerWeight or SparsePowerWeight.
template <class PowerWeightT>
class ProjectPowerWeightMapper {
 public:
  using FromWeight = PowerWeightT;
  using ToWeight = PowerWeightT;
  using Index = typename PowerWeightT::Index;
  using ComponentWeight = typename PowerWeightT::Weight;

  explicit ProjectPowerWeightMapper(
      Index from_index = 0, Index to_index = 0,
      const ComponentWeight &default_weight = ComponentWeight::Zero())
      : from_index_(from_index),
        to_index_(to_index),
        default_weight_(default_weight) {}

  PowerWeightT operator()(const PowerWeightT &w) const {
    return PowerWeightT(to_index_, w.Value(from_index_), default_weight_);
  }

 private:
  const Index from_index_;
  const Index to_index_;
  const ComponentWeight default_weight_;
};

// Applies a transformation function to each weight vector.
// PowerWeightT may be PowerWeight or SparsePowerWeight.
template <class PowerWeightT, class TransformFn_>
class TransformPowerWeightMapper {
 public:
  using TransformFn = TransformFn_;
  using FromWeight = PowerWeightT;
  using ToWeight = PowerWeightT;

  explicit TransformPowerWeightMapper(
      const TransformFn &transform = TransformFn())
      : transform_(transform) {}

  PowerWeightT operator()(const PowerWeightT &w) const { return transform_(w); }

 private:
  TransformFn transform_;
};

}  // namespace fst

#endif  // FST_POWER_WEIGHT_MAPPERS_H_
