// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "ink/strokes/internal/stroke_input_modeler/sliding_window_input_modeler.h"

#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "fuzztest/fuzztest.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/types/span.h"
#include "ink/brush/brush_family.h"
#include "ink/geometry/angle.h"
#include "ink/geometry/type_matchers.h"
#include "ink/strokes/input/fuzz_domains.h"
#include "ink/strokes/input/stroke_input.h"
#include "ink/strokes/input/stroke_input_batch.h"
#include "ink/strokes/internal/modeled_stroke_input.h"
#include "ink/strokes/internal/stroke_input_modeler.h"
#include "ink/types/duration.h"

namespace ink::strokes_internal {
namespace {

using ::absl_testing::IsOk;
using ::testing::AllOf;
using ::testing::Each;
using ::testing::Field;
using ::testing::Ge;
using ::testing::IsEmpty;
using ::testing::Le;
using ::testing::Not;
using ::testing::SizeIs;

TEST(SlidingWindowInputModelerTest, ConstantVelocityRawInputs) {
  StrokeInputModeler modeler;
  modeler.StartStroke(
      BrushFamily::SlidingWindowModel{
          .window_size = Duration32::Millis(20),
          .upsampling_period = Duration32::Millis(5),
      },
      /* brush_epsilon = */ 0.01);

  // Extend the stroke with a bunch of inputs that move at a constant velocity
  // of 1000 stroke units per second.
  StrokeInputBatch inputs;
  for (int i = 0; i < 100; ++i) {
    StrokeInput input = StrokeInput{
        .position = {static_cast<float>(i), 0},
        .elapsed_time = Duration32::Millis(i),
    };
    ASSERT_THAT(inputs.Append(input), IsOk());
  }
  modeler.ExtendStroke(inputs, {}, Duration32::Millis(100));

  // All modeled inputs (even the first and last one) should have a modeled
  // velocity of about 1000 stroke units per second, and a modeled acceleration
  // of (roughly) zero.
  EXPECT_THAT(modeler.GetModeledInputs(), SizeIs(Ge(100)));
  EXPECT_THAT(modeler.GetModeledInputs(),
              Each(Field("velocity", &ModeledStrokeInput::velocity,
                         VecNear({1000, 0}, 0.001))));
  EXPECT_THAT(modeler.GetModeledInputs(),
              Each(Field("acceleration", &ModeledStrokeInput::acceleration,
                         VecNear({0, 0}, 0.5))));
}

TEST(SlidingWindowInputModelerTest, Orientation) {
  StrokeInputModeler modeler;
  modeler.StartStroke(
      BrushFamily::SlidingWindowModel{
          .window_size = Duration32::Millis(10),
          .upsampling_period = Duration32::Millis(5),
      },
      /* brush_epsilon = */ 0.01);

  // Extend the stroke with a bunch of inputs with an orientation of 10°, then
  // a bunch of inputs with an orientation of 350°.
  StrokeInputBatch inputs;
  for (int i = 0; i < 100; ++i) {
    StrokeInput input = StrokeInput{
        .position = {static_cast<float>(i), 0},
        .elapsed_time = Duration32::Millis(i),
        .orientation = i < 50 ? Angle::Degrees(10) : Angle::Degrees(350),
    };
    ASSERT_THAT(inputs.Append(input), IsOk());
  }
  modeler.ExtendStroke(inputs, {}, Duration32::Millis(100));

  // All modeled inputs should have an orientation (roughly) between ±10° when
  // normalized about zero; they shouldn't be naively averaged between 10° and
  // 350° to get ~180°.
  EXPECT_THAT(modeler.GetModeledInputs(), SizeIs(Ge(100)));
  for (const ModeledStrokeInput& modeled_input : modeler.GetModeledInputs()) {
    ASSERT_NE(modeled_input.orientation, StrokeInput::kNoOrientation);
    EXPECT_THAT(
        modeled_input.orientation.NormalizedAboutZero(),
        AllOf(Ge(Angle::Degrees(-10.0001)), Le(Angle::Degrees(10.0001))));
  }
}

TEST(SlidingWindowInputModelerTest, Upsampling) {
  StrokeInputModeler modeler;
  modeler.StartStroke(
      BrushFamily::SlidingWindowModel{
          .window_size = Duration32::Millis(10),
          .upsampling_period = Duration32::Millis(1),
      },
      /* brush_epsilon = */ 0.01);

  // Extend the stroke with three raw inputs, spaced 10 ms apart.
  absl::StatusOr<StrokeInputBatch> inputs = StrokeInputBatch::Create({
      {.position = {0, 0}, .elapsed_time = Duration32::Millis(0)},
      {.position = {100, 0}, .elapsed_time = Duration32::Millis(10)},
      {.position = {100, 100}, .elapsed_time = Duration32::Millis(20)},
  });
  ASSERT_THAT(inputs, IsOk());
  modeler.ExtendStroke(*inputs, {}, Duration32::Millis(20));

  // Since the upsampling period is 1 ms, we should end up with 21 modeled
  // inputs.
  absl::Span<const ModeledStrokeInput> modeled = modeler.GetModeledInputs();
  ASSERT_THAT(modeled, SizeIs(21));
  // The modeled positions should move in a smooth curve near the corner of the
  // L shaped formed by the raw inputs.
  EXPECT_THAT(modeled[0].position, PointNear({0.0, 0.0}, 0.1));
  EXPECT_THAT(modeled[1].position, PointNear({10.0, 0.0}, 0.1));
  EXPECT_THAT(modeled[2].position, PointNear({20.0, 0.0}, 0.1));
  EXPECT_THAT(modeled[3].position, PointNear({30.0, 0.0}, 0.1));
  EXPECT_THAT(modeled[4].position, PointNear({40.0, 0.0}, 0.1));
  EXPECT_THAT(modeled[5].position, PointNear({50.0, 0.0}, 0.1));
  EXPECT_THAT(modeled[6].position, PointNear({59.5, 0.5}, 0.1));
  EXPECT_THAT(modeled[7].position, PointNear({68.0, 2.0}, 0.1));
  EXPECT_THAT(modeled[8].position, PointNear({75.5, 4.5}, 0.1));
  EXPECT_THAT(modeled[9].position, PointNear({82.0, 8.0}, 0.1));
  EXPECT_THAT(modeled[10].position, PointNear({87.5, 12.5}, 0.1));
  EXPECT_THAT(modeled[11].position, PointNear({92.0, 18.0}, 0.1));
  EXPECT_THAT(modeled[12].position, PointNear({95.5, 24.5}, 0.1));
  EXPECT_THAT(modeled[13].position, PointNear({98.0, 32.0}, 0.1));
  EXPECT_THAT(modeled[14].position, PointNear({99.5, 40.5}, 0.1));
  EXPECT_THAT(modeled[15].position, PointNear({100.0, 50.0}, 0.1));
  EXPECT_THAT(modeled[16].position, PointNear({100.0, 60.0}, 0.1));
  EXPECT_THAT(modeled[17].position, PointNear({100.0, 70.0}, 0.1));
  EXPECT_THAT(modeled[18].position, PointNear({100.0, 80.0}, 0.1));
  EXPECT_THAT(modeled[19].position, PointNear({100.0, 90.0}, 0.1));
  EXPECT_THAT(modeled[20].position, PointNear({100.0, 100.0}, 0.1));
}

TEST(SlidingWindowInputModelerTest, PruneModeledInputsWithinEpsilon) {
  StrokeInputModeler modeler;
  modeler.StartStroke(
      BrushFamily::SlidingWindowModel{
          .window_size = Duration32::Millis(10),
          .upsampling_period = Duration32::Millis(5),
      },
      /*brush_epsilon=*/0.01);

  // Extend the stroke with three raw inputs, spaced 100 ms apart, but all with
  // the same position.
  absl::StatusOr<StrokeInputBatch> inputs = StrokeInputBatch::Create({
      {.position = {12, 34}, .elapsed_time = Duration32::Millis(0)},
      {.position = {12, 34}, .elapsed_time = Duration32::Millis(100)},
      {.position = {12, 34}, .elapsed_time = Duration32::Millis(200)},
  });
  ASSERT_THAT(inputs, IsOk());
  modeler.ExtendStroke(*inputs, {}, Duration32::Millis(200));

  // Since all positions are within `brush_epsilon`, there should only be a
  // single modeled input; all others should get pruned.
  EXPECT_THAT(modeler.GetModeledInputs(), SizeIs(1));
  EXPECT_EQ(modeler.GetState().real_input_count, 1);
}

void CanModelAnyStrokeInputBatch(const StrokeInputBatch& inputs) {
  StrokeInputModeler modeler;
  modeler.StartStroke(
      BrushFamily::SlidingWindowModel{
          .window_size = Duration32::Millis(10),
          .upsampling_period = Duration32::Millis(5),
      },
      /*brush_epsilon=*/1);
  modeler.ExtendStroke(inputs, {}, Duration32::Zero());
  if (inputs.IsEmpty()) {
    EXPECT_THAT(modeler.GetModeledInputs(), IsEmpty());
  } else {
    EXPECT_THAT(modeler.GetModeledInputs(), Not(IsEmpty()));
  }
  EXPECT_EQ(modeler.GetState().real_input_count,
            modeler.GetModeledInputs().size());
  EXPECT_LE(modeler.GetState().stable_input_count,
            modeler.GetState().real_input_count);
}
FUZZ_TEST(SlidingWindowInputModelerTest, CanModelAnyStrokeInputBatch)
    .WithDomains(ArbitraryStrokeInputBatch());

}  // namespace
}  // namespace ink::strokes_internal
