// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using System.Collections.Generic;
using System.Collections.Immutable;
using Microsoft.ML.Runtime;

namespace Microsoft.ML.Trainers.FastTree
{
 /// <summary>
 /// A container base class for exposing <see cref="InternalRegressionTree"/>'s and
 /// <see cref="InternalQuantileRegressionTree"/>'s attributes to users.
 /// This class should not be mutable, so it contains a lot of read-only members.
 /// </summary>
 public abstract class RegressionTreeBase
 {
 /// <summary>
 /// <see cref="RegressionTreeBase"/> is an immutable wrapper over <see cref="_tree"/> for exposing some tree's
 /// attribute to users.
 /// </summary>
 private readonly InternalRegressionTree _tree;

 /// <summary>
 /// See <see cref="LeftChild"/>.
 /// </summary>
 private readonly ImmutableArray<int> _lteChild;
 /// <summary>
 /// See <see cref="RightChild"/>.
 /// </summary>
 private readonly ImmutableArray<int> _gtChild;
 /// <summary>
 /// See <see cref="NumericalSplitFeatureIndexes"/>.
 /// </summary>
 private readonly ImmutableArray<int> _numericalSplitFeatureIndexes;
 /// <summary>
 /// See <see cref="NumericalSplitThresholds"/>.
 /// </summary>
 private readonly ImmutableArray<float> _numericalSplitThresholds;
 /// <summary>
 /// See <see cref="CategoricalSplitFlags"/>.
 /// </summary>
 private readonly ImmutableArray<bool> _categoricalSplitFlags;
 /// <summary>
 /// See <see cref="LeafValues"/>.
 /// </summary>
 private readonly ImmutableArray<double> _leafValues;
 /// <summary>
 /// See <see cref="SplitGains"/>.
 /// </summary>
 private readonly ImmutableArray<double> _splitGains;

 /// <summary>
 /// <see cref="LeftChild"/>[i] is the i-th node's child index used when
 /// (1) the numerical feature indexed by <see cref="NumericalSplitFeatureIndexes"/>[i] is less than or equal
 /// to the threshold <see cref="NumericalSplitThresholds"/>[i], or
 /// (2) the categorical features indexed by <see cref="GetCategoricalCategoricalSplitFeatureRangeAt(int)"/>'s
 /// returned value with nodeIndex=i is NOT a sub-set of <see cref="GetCategoricalSplitFeaturesAt(int)"/> with
 /// nodeIndex=i.
 /// Note that the case (1) happens only when <see cref="CategoricalSplitFlags"/>[i] is true and otherwise (2)
 /// occurs. A non-negative returned value means a node (i.e., not a leaf); for example, 2 means the 3rd node in
 /// the underlying <see cref="_tree"/>. A negative returned value means a leaf; for example, -1 stands for the
 /// <see langword="~"/>(-1)-th leaf in the underlying <see cref="_tree"/>. Note that <see langword="~"/> is the
 /// bitwise complement operator in C#; for details, see
 /// https://docs.microsoft.com/en-us/dotnet/csharp/language-reference/operators/bitwise-complement-operator.
 /// </summary>
 public IReadOnlyList<int> LeftChild => _lteChild;

 /// <summary>
 /// <see cref="RightChild"/>[i] is the i-th node's child index used when the two conditions, (1) and (2),
 /// described in <see cref="LeftChild"/>'s document are not true. Its return value follows the format
 /// used in <see cref="LeftChild"/>.
 /// </summary>
 public IReadOnlyList<int> RightChild => _gtChild;

 /// <summary>
 /// <see cref="NumericalSplitFeatureIndexes"/>[i] is the feature index used the splitting function of the
 /// i-th node. This value is valid only if <see cref="CategoricalSplitFlags"/>[i] is false.
 /// </summary>
 public IReadOnlyList<int> NumericalSplitFeatureIndexes => _numericalSplitFeatureIndexes;

 /// <summary>
 /// <see cref="NumericalSplitThresholds"/>[i] is the threshold on feature indexed by
 /// <see cref="NumericalSplitFeatureIndexes"/>[i], where i is the i-th node's index
 /// (for example, i is 1 for the 2nd node in <see cref="_tree"/>).
 /// </summary>
 public IReadOnlyList<float> NumericalSplitThresholds => _numericalSplitThresholds;

 /// <summary>
 /// Determine the types of splitting function. If <see cref="CategoricalSplitFlags"/>[i] is true, the i-th
 /// node's uses categorical splitting function. Otherwise, traditional numerical split is used.
 /// </summary>
 public IReadOnlyList<bool> CategoricalSplitFlags => _categoricalSplitFlags;

 /// <summary>
 /// <see cref="LeafValues"/>[i] is the learned value at the i-th leaf.
 /// </summary>
 public IReadOnlyList<double> LeafValues => _leafValues;

 /// <summary>
 /// Return categorical thresholds used at node indexed by nodeIndex. If the considered input feature does NOT
 /// matche any of values returned by <see cref="GetCategoricalSplitFeaturesAt(int)"/>, we call it a
 /// less-than-threshold event and therefore <see cref="LeftChild"/>[nodeIndex] is the child node that input
 /// should go next. The returned value is valid only if <see cref="CategoricalSplitFlags"/>[nodeIndex] is true.
 /// </summary>
 public IReadOnlyList<int> GetCategoricalSplitFeaturesAt(int nodeIndex)
 {
 if (nodeIndex < 0 || nodeIndex >= NumberOfNodes)
 throw Contracts.Except($"The input index, {nodeIndex}, is invalid. Its valid range is from 0 (inclusive) to {NumberOfNodes} (exclusive).");

 if (_tree.CategoricalSplitFeatures == null || _tree.CategoricalSplitFeatures[nodeIndex] == null)
 return new List<int>(); // Zero-length vector.
 else
 return _tree.CategoricalSplitFeatures[nodeIndex];
 }
 /// <summary>
 /// Return categorical thresholds' range used at node indexed by nodeIndex. A categorical split at node indexed
 /// by nodeIndex can consider multiple consecutive input features at one time; their range is specified by
 /// <see cref="GetCategoricalCategoricalSplitFeatureRangeAt(int)"/>. The returned value is always a 2-element
 /// array; its 1st element is the starting index and its 2nd element is the endining index of a feature segment.
 /// The returned value is valid only if <see cref="CategoricalSplitFlags"/>[nodeIndex] is true.
 /// </summary>
 public IReadOnlyList<int> GetCategoricalCategoricalSplitFeatureRangeAt(int nodeIndex)
 {
 if (nodeIndex < 0 || nodeIndex >= NumberOfNodes)
 throw Contracts.Except($"The input node index, {nodeIndex}, is invalid. Its valid range is from 0 (inclusive) to {NumberOfNodes} (exclusive).");

 if (_tree.CategoricalSplitFeatureRanges == null || _tree.CategoricalSplitFeatureRanges[nodeIndex] == null)
 return new List<int>(); // Zero-length vector.
 else
 return _tree.CategoricalSplitFeatureRanges[nodeIndex];
 }

 /// <summary>
 /// The gains obtained by splitting data at nodes. Its i-th value is computed from to the split at the i-th node.
 /// </summary>
 public IReadOnlyList<double> SplitGains => _splitGains;

 /// <summary>
 /// Number of leaves in the tree. Note that <see cref="NumberOfLeaves"/> does not take non-leaf nodes into account.
 /// </summary>
 public int NumberOfLeaves => _tree.NumLeaves;

 /// <summary>
 /// Number of nodes in the tree. This doesn't include any leaves. For example, a tree with node0->node1,
 /// node0->leaf3, node1->leaf1, node1->leaf2, <see cref="NumberOfNodes"/> and <see cref="NumberOfLeaves"/> should
 /// be 2 and 3, respectively.
 /// </summary>
 // A visualization of the example mentioned in this doc string.
 // node0
 // / \
 // node1 leaf3
 // / \
 // leaf1 leaf2
 // The index of leaf starts with 1 because interally we use "-1" as the 1st leaf's index, "-2" for the 2nd leaf's index, and so on.
 public int NumberOfNodes => _tree.NumNodes;

 internal RegressionTreeBase(InternalRegressionTree tree)
 {
 _tree = tree;

 _lteChild = ImmutableArray.Create(_tree.LteChild, 0, _tree.NumNodes);
 _gtChild = ImmutableArray.Create(_tree.GtChild, 0, _tree.NumNodes);

 _numericalSplitFeatureIndexes = ImmutableArray.Create(_tree.SplitFeatures, 0, _tree.NumNodes);
 _numericalSplitThresholds = ImmutableArray.Create(_tree.RawThresholds, 0, _tree.NumNodes);
 _categoricalSplitFlags = ImmutableArray.Create(_tree.CategoricalSplit, 0, _tree.NumNodes);
 _leafValues = ImmutableArray.Create(_tree.LeafValues, 0, _tree.NumLeaves);
 _splitGains = ImmutableArray.Create(_tree.SplitGains, 0, _tree.NumNodes);
 }
 }

 /// <summary>
 /// A container class for exposing <see cref="InternalRegressionTree"/>'s attributes to users.
 /// This class should not be mutable, so it contains a lot of read-only members. Note that
 /// <see cref="RegressionTree"/> is identical to <see cref="RegressionTreeBase"/> but in
 /// another derived class <see cref="QuantileRegressionTree"/> some attributes are added.
 /// </summary>
 public sealed class RegressionTree : RegressionTreeBase
 {
 internal RegressionTree(InternalRegressionTree tree) : base(tree) { }
 }

 /// <summary>
 /// A container class for exposing <see cref="InternalQuantileRegressionTree"/>'s attributes to users.
 /// This class should not be mutable, so it contains a lot of read-only members. In addition to
 /// things inherited from <see cref="RegressionTreeBase"/>, we add <see cref="GetLeafSamplesAt(int)"/>
 /// and <see cref="GetLeafSampleWeightsAt(int)"/> to expose (sub-sampled) training labels falling into
 /// the leafIndex-th leaf and their weights.
 /// </summary>
 public sealed class QuantileRegressionTree : RegressionTreeBase
 {
 /// <summary>
 /// Sample labels from training data. <see cref="_leafSamples"/>[i] stores the labels falling into the
 /// i-th leaf.
 /// </summary>
 private readonly double[][] _leafSamples;

 /// <summary>
 /// Sample labels' weights from training data. <see cref="_leafSampleWeights"/>[i] stores the weights for
 /// labels falling into the i-th leaf. <see cref="_leafSampleWeights"/>[i][j] is the weight of
 /// <see cref="_leafSamples"/>[i][j].
 /// </summary>
 private readonly double[][] _leafSampleWeights;

 /// <summary>
 /// Return the training labels falling into the specified leaf.
 /// </summary>
 /// <param name="leafIndex">The index of the specified leaf.</param>
 /// <returns>Training labels</returns>
 public IReadOnlyList<double> GetLeafSamplesAt(int leafIndex)
 {
 if (leafIndex < 0 || leafIndex >= NumberOfLeaves)
 throw Contracts.Except($"The input leaf index, {leafIndex}, is invalid. Its valid range is from 0 (inclusive) to {NumberOfLeaves} (exclusive).");

 // _leafSample always contains valid values assigned in constructor.
 return _leafSamples[leafIndex];
 }

 /// <summary>
 /// Return the weights for training labels falling into the specified leaf. If <see cref="GetLeafSamplesAt"/>
 /// and <see cref="GetLeafSampleWeightsAt"/> use the same input, the i-th returned value of this function is
 /// the weight of the i-th label in <see cref="GetLeafSamplesAt"/>.
 /// </summary>
 /// <param name="leafIndex">The index of the specified leaf.</param>
 /// <returns>Training labels' weights</returns>
 public IReadOnlyList<double> GetLeafSampleWeightsAt(int leafIndex)
 {
 if (leafIndex < 0 || leafIndex >= NumberOfLeaves)
 throw Contracts.Except($"The input leaf index, {leafIndex}, is invalid. Its valid range is from 0 (inclusive) to {NumberOfLeaves} (exclusive).");

 // _leafSampleWeights always contains valid values assigned in constructor.
 return _leafSampleWeights[leafIndex];
 }

 internal QuantileRegressionTree(InternalQuantileRegressionTree tree) : base(tree)
 {
 tree.ExtractLeafSamplesAndTheirWeights(out _leafSamples, out _leafSampleWeights);
 }
 }
}