// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.

package org.wpilib.simulation;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.util.List;
import org.junit.jupiter.api.Test;
import org.wpilib.math.controller.LTVUnicycleController;
import org.wpilib.math.controller.LinearPlantInversionFeedforward;
import org.wpilib.math.geometry.Pose2d;
import org.wpilib.math.geometry.Rotation2d;
import org.wpilib.math.kinematics.DifferentialDriveKinematics;
import org.wpilib.math.linalg.Matrix;
import org.wpilib.math.linalg.VecBuilder;
import org.wpilib.math.linalg.Vector;
import org.wpilib.math.numbers.N1;
import org.wpilib.math.numbers.N7;
import org.wpilib.math.system.DCMotor;
import org.wpilib.math.system.Models;
import org.wpilib.math.system.NumericalIntegration;
import org.wpilib.math.trajectory.TrajectoryConfig;
import org.wpilib.math.trajectory.TrajectoryGenerator;
import org.wpilib.math.trajectory.constraint.DifferentialDriveKinematicsConstraint;
import org.wpilib.math.util.Nat;
import org.wpilib.math.util.Units;

class DifferentialDrivetrainSimTest {
  @Test
  void testConvergence() {
    var motor = DCMotor.getNEO(2);
    var plant =
        Models.differentialDriveFromPhysicalConstants(
            motor, 50, Units.inchesToMeters(2), Units.inchesToMeters(12), 0.5, 1.0);

    var kinematics = new DifferentialDriveKinematics(Units.inchesToMeters(24));
    var sim =
        new DifferentialDrivetrainSim(
            plant,
            motor,
            1,
            kinematics.trackwidth,
            Units.inchesToMeters(2),
            VecBuilder.fill(0, 0, 0.0001, 0.1, 0.1, 0.005, 0.005));

    var feedforward = new LinearPlantInversionFeedforward<>(plant, 0.020);
    var feedback = new LTVUnicycleController(0.020);
    feedforward.reset(VecBuilder.fill(0, 0));

    // ground truth
    Matrix<N7, N1> groundTruthX = new Vector<>(Nat.N7());

    var traj =
        TrajectoryGenerator.generateTrajectory(
            Pose2d.kZero,
            List.of(),
            new Pose2d(2, 2, Rotation2d.kZero),
            new TrajectoryConfig(1, 1)
                .addConstraint(new DifferentialDriveKinematicsConstraint(kinematics, 1)));

    for (double t = 0; t < traj.getTotalTime(); t += 0.020) {
      var state = traj.sample(t);
      var feedbackOut = feedback.calculate(sim.getPose(), state);

      var wheelVelocities = kinematics.toWheelVelocities(feedbackOut);

      var voltages =
          feedforward.calculate(VecBuilder.fill(wheelVelocities.left, wheelVelocities.right));

      // Sim periodic code
      sim.setInputs(voltages.get(0, 0), voltages.get(1, 0));
      sim.update(0.020);

      // Update our ground truth
      groundTruthX = NumericalIntegration.rkdp(sim::getDynamics, groundTruthX, voltages, 0.020);
    }

    // 2 inch tolerance is OK since our ground truth is an approximation of the
    // ODE solution using RKDP anyway
    assertEquals(
        groundTruthX.get(DifferentialDrivetrainSim.State.X.value, 0),
        sim.getState(DifferentialDrivetrainSim.State.X),
        0.05);
    assertEquals(
        groundTruthX.get(DifferentialDrivetrainSim.State.Y.value, 0),
        sim.getState(DifferentialDrivetrainSim.State.Y),
        0.05);
    assertEquals(
        groundTruthX.get(DifferentialDrivetrainSim.State.HEADING.value, 0),
        sim.getState(DifferentialDrivetrainSim.State.HEADING),
        0.01);
  }

  @Test
  void testCurrent() {
    var motor = DCMotor.getNEO(2);
    var plant =
        Models.differentialDriveFromPhysicalConstants(
            motor, 50, Units.inchesToMeters(2), Units.inchesToMeters(12), 0.5, 1.0);
    var kinematics = new DifferentialDriveKinematics(Units.inchesToMeters(24));
    var sim =
        new DifferentialDrivetrainSim(
            plant, motor, 1, kinematics.trackwidth, Units.inchesToMeters(2), null);

    sim.setInputs(-12, -12);
    for (int i = 0; i < 10; i++) {
      sim.update(0.020);
    }
    assertTrue(sim.getCurrentDraw() > 0);

    sim.setInputs(12, 12);
    for (int i = 0; i < 20; i++) {
      sim.update(0.020);
    }
    assertTrue(sim.getCurrentDraw() > 0);

    sim.setInputs(-12, 12);
    for (int i = 0; i < 30; i++) {
      sim.update(0.020);
    }
    assertTrue(sim.getCurrentDraw() > 0);
  }

  @Test
  void testModelStability() {
    var motor = DCMotor.getNEO(2);
    var plant =
        Models.differentialDriveFromPhysicalConstants(
            motor, 50, Units.inchesToMeters(2), Units.inchesToMeters(12), 2.0, 5.0);

    var kinematics = new DifferentialDriveKinematics(Units.inchesToMeters(24));
    var sim =
        new DifferentialDrivetrainSim(
            plant,
            motor,
            5,
            kinematics.trackwidth,
            Units.inchesToMeters(2),
            VecBuilder.fill(0, 0, 0.0001, 0.1, 0.1, 0.005, 0.005));

    sim.setInputs(2, 4);

    // 10 seconds should be enough time to verify stability
    for (int i = 0; i < 500; i++) {
      sim.update(0.020);
    }

    assertTrue(Math.abs(sim.getPose().getTranslation().getNorm()) < 100);
  }
}
