/*
 * MIT License
 *
 * Copyright (c) PhotonVision
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

#include "subsystems/SwerveDriveSim.h"

#include <wpi/math/system/Discretization.hpp>
#include <wpi/system/RobotController.hpp>

template <typename T>
int sgn(T val) {
  return (T(0) < val) - (val < T(0));
}

SwerveDriveSim::SwerveDriveSim(
    const wpi::math::SimpleMotorFeedforward<wpi::units::meters>& driveFF,
    const wpi::math::DCMotor& driveMotor, double driveGearing,
    wpi::units::meter_t driveWheelRadius,
    const wpi::math::SimpleMotorFeedforward<wpi::units::radians>& steerFF,
    const wpi::math::DCMotor& steerMotor, double steerGearing,
    const wpi::math::SwerveDriveKinematics<numModules>& kinematics)
    : SwerveDriveSim(
          wpi::math::LinearSystem<2, 1, 2>{
              (Eigen::MatrixXd(2, 2) << 0.0, 1.0, 0.0,
               -driveFF.GetKv().to<double>() / driveFF.GetKa().to<double>())
                  .finished(),
              Eigen::Matrix<double, 2, 1>{0.0,
                                          1.0 / driveFF.GetKa().to<double>()},
              (Eigen::MatrixXd(2, 2) << 1.0, 0.0, 0.0, 1.0).finished(),
              Eigen::Matrix<double, 2, 1>{0.0, 0.0}},
          driveFF.GetKs(), driveMotor, driveGearing, driveWheelRadius,
          wpi::math::LinearSystem<2, 1, 2>{
              (Eigen::MatrixXd(2, 2) << 0.0, 1.0, 0.0,
               -steerFF.GetKv().to<double>() / steerFF.GetKa().to<double>())
                  .finished(),
              Eigen::Matrix<double, 2, 1>{0.0,
                                          1.0 / steerFF.GetKa().to<double>()},
              (Eigen::MatrixXd(2, 2) << 1.0, 0.0, 0.0, 1.0).finished(),
              Eigen::Matrix<double, 2, 1>{0.0, 0.0}},
          steerFF.GetKs(), steerMotor, steerGearing, kinematics) {}

SwerveDriveSim::SwerveDriveSim(
    const wpi::math::LinearSystem<2, 1, 2>& drivePlant,
    wpi::units::volt_t driveKs, const wpi::math::DCMotor& driveMotor,
    double driveGearing, wpi::units::meter_t driveWheelRadius,
    const wpi::math::LinearSystem<2, 1, 2>& steerPlant,
    wpi::units::volt_t steerKs, const wpi::math::DCMotor& steerMotor,
    double steerGearing,
    const wpi::math::SwerveDriveKinematics<numModules>& kinematics)
    : drivePlant(drivePlant),
      driveKs(driveKs),
      driveMotor(driveMotor),
      driveGearing(driveGearing),
      driveWheelRadius(driveWheelRadius),
      steerPlant(steerPlant),
      steerKs(steerKs),
      steerMotor(steerMotor),
      steerGearing(steerGearing),
      kinematics(kinematics) {}

void SwerveDriveSim::SetDriveInputs(
    const std::array<wpi::units::volt_t, numModules>& inputs) {
  wpi::units::volt_t battVoltage = wpi::RobotController::GetBatteryVoltage();
  for (int i = 0; i < driveInputs.size(); i++) {
    wpi::units::volt_t input = inputs[i];
    driveInputs[i] = std::clamp(input, -battVoltage, battVoltage);
  }
}

void SwerveDriveSim::SetSteerInputs(
    const std::array<wpi::units::volt_t, numModules>& inputs) {
  wpi::units::volt_t battVoltage = wpi::RobotController::GetBatteryVoltage();
  for (int i = 0; i < steerInputs.size(); i++) {
    wpi::units::volt_t input = inputs[i];
    steerInputs[i] = std::clamp(input, -battVoltage, battVoltage);
  }
}

Eigen::Matrix<double, 2, 1> SwerveDriveSim::CalculateX(
    const Eigen::Matrix<double, 2, 2>& discA,
    const Eigen::Matrix<double, 2, 1>& discB,
    const Eigen::Matrix<double, 2, 1>& x, wpi::units::volt_t input,
    wpi::units::volt_t kS) {
  auto Ax = discA * x;
  double nextStateVel = Ax(1, 0);
  double inputToStop = nextStateVel / -discB(1, 0);
  double ksSystemEffect =
      std::clamp(inputToStop, -kS.to<double>(), kS.to<double>());

  nextStateVel += discB(1, 0) * ksSystemEffect;
  inputToStop = nextStateVel / -discB(1, 0);
  double signToStop = sgn(inputToStop);
  double inputSign = sgn(input.to<double>());
  double ksInputEffect = 0;

  if (std::abs(ksSystemEffect) < kS.to<double>()) {
    double absInput = std::abs(input.to<double>());
    ksInputEffect =
        -std::clamp(kS.to<double>() * inputSign, -absInput, absInput);
  } else if ((input.to<double>() * signToStop) > (inputToStop * signToStop)) {
    double absInput = std::abs(input.to<double>() - inputToStop);
    ksInputEffect =
        -std::clamp(kS.to<double>() * inputSign, -absInput, absInput);
  }

  auto sF = Eigen::Matrix<double, 1, 1>{input.to<double>() + ksSystemEffect +
                                        ksInputEffect};
  auto Bu = discB * sF;
  auto retVal = Ax + Bu;
  return retVal;
}

void SwerveDriveSim::Update(wpi::units::second_t dt) {
  Eigen::Matrix<double, 2, 2> driveDiscA;
  Eigen::Matrix<double, 2, 1> driveDiscB;
  wpi::math::DiscretizeABAB<2, 1>(drivePlant.A(), drivePlant.B(), dt,
                                  &driveDiscA, &driveDiscB);

  Eigen::Matrix<double, 2, 2> steerDiscA;
  Eigen::Matrix<double, 2, 1> steerDiscB;
  wpi::math::DiscretizeABAB<2, 1>(steerPlant.A(), steerPlant.B(), dt,
                                  &steerDiscA, &steerDiscB);

  std::array<wpi::math::SwerveModulePosition, 4> moduleDeltas;

  for (int i = 0; i < numModules; i++) {
    double prevDriveStatePos = driveStates[i](0, 0);
    driveStates[i] = CalculateX(driveDiscA, driveDiscB, driveStates[i],
                                driveInputs[i], driveKs);
    double currentDriveStatePos = driveStates[i](0, 0);
    steerStates[i] = CalculateX(steerDiscA, steerDiscB, steerStates[i],
                                steerInputs[i], steerKs);
    double currentSteerStatePos = steerStates[i](0, 0);
    moduleDeltas[i] = wpi::math::SwerveModulePosition{
        wpi::units::meter_t{currentDriveStatePos - prevDriveStatePos},
        wpi::math::Rotation2d{wpi::units::radian_t{currentSteerStatePos}}};
  }

  wpi::math::Twist2d twist = kinematics.ToTwist2d(moduleDeltas);
  pose = pose.Exp(twist);
  omega = twist.dtheta / dt;
}

void SwerveDriveSim::Reset(const wpi::math::Pose2d& pose, bool preserveMotion) {
  this->pose = pose;
  if (!preserveMotion) {
    for (int i = 0; i < numModules; i++) {
      driveStates[i] = Eigen::Matrix<double, 2, 1>{0, 0};
      steerStates[i] = Eigen::Matrix<double, 2, 1>{0, 0};
    }
    omega = 0_rad_per_s;
  }
}

void SwerveDriveSim::Reset(const wpi::math::Pose2d& pose,
                           const std::array<Eigen::Matrix<double, 2, 1>,
                                            numModules>& moduleDriveStates,
                           const std::array<Eigen::Matrix<double, 2, 1>,
                                            numModules>& moduleSteerStates) {
  this->pose = pose;
  driveStates = moduleDriveStates;
  steerStates = moduleSteerStates;
  omega = kinematics.ToChassisSpeeds(GetModuleStates()).omega;
}

wpi::math::Pose2d SwerveDriveSim::GetPose() const { return pose; }

std::array<wpi::math::SwerveModulePosition, numModules>
SwerveDriveSim::GetModulePositions() const {
  std::array<wpi::math::SwerveModulePosition, numModules> positions;
  for (int i = 0; i < numModules; i++) {
    positions[i] = wpi::math::SwerveModulePosition{
        wpi::units::meter_t{driveStates[i](0, 0)},
        wpi::math::Rotation2d{wpi::units::radian_t{steerStates[i](0, 0)}}};
  }
  return positions;
}

std::array<wpi::math::SwerveModulePosition, numModules>
SwerveDriveSim::GetNoisyModulePositions(wpi::units::meter_t driveStdDev,
                                        wpi::units::radian_t steerStdDev) {
  std::array<wpi::math::SwerveModulePosition, numModules> positions;
  for (int i = 0; i < numModules; i++) {
    positions[i] = wpi::math::SwerveModulePosition{
        wpi::units::meter_t{driveStates[i](0, 0)} +
            randDist(generator) * driveStdDev,
        wpi::math::Rotation2d{wpi::units::radian_t{steerStates[i](0, 0)} +
                              randDist(generator) * steerStdDev}};
  }
  return positions;
}

std::array<wpi::math::SwerveModuleState, numModules>
SwerveDriveSim::GetModuleStates() {
  std::array<wpi::math::SwerveModuleState, numModules> states;
  for (int i = 0; i < numModules; i++) {
    states[i] = wpi::math::SwerveModuleState{
        wpi::units::meters_per_second_t{driveStates[i](1, 0)},
        wpi::math::Rotation2d{wpi::units::radian_t{steerStates[i](0, 0)}}};
  }
  return states;
}

std::array<Eigen::Matrix<double, 2, 1>, numModules>
SwerveDriveSim::GetDriveStates() const {
  return driveStates;
}

std::array<Eigen::Matrix<double, 2, 1>, numModules>
SwerveDriveSim::GetSteerStates() const {
  return steerStates;
}

wpi::units::radians_per_second_t SwerveDriveSim::GetOmega() const {
  return omega;
}

wpi::units::ampere_t SwerveDriveSim::GetCurrentDraw(
    const wpi::math::DCMotor& motor, wpi::units::radians_per_second_t velocity,
    wpi::units::volt_t inputVolts, wpi::units::volt_t batteryVolts) const {
  wpi::units::volt_t effVolts = inputVolts - velocity / motor.Kv;
  if (inputVolts >= 0_V) {
    effVolts = std::clamp(effVolts, 0_V, inputVolts);
  } else {
    effVolts = std::clamp(effVolts, inputVolts, 0_V);
  }
  auto retVal = (inputVolts / batteryVolts) * (effVolts / motor.R);
  return retVal;
}

std::array<wpi::units::ampere_t, numModules>
SwerveDriveSim::GetDriveCurrentDraw() const {
  std::array<wpi::units::ampere_t, numModules> currents;
  for (int i = 0; i < numModules; i++) {
    wpi::units::radians_per_second_t speed =
        wpi::units::radians_per_second_t{driveStates[i](1, 0)} * driveGearing /
        driveWheelRadius.to<double>();
    currents[i] = GetCurrentDraw(driveMotor, speed, driveInputs[i],
                                 wpi::RobotController::GetBatteryVoltage());
  }
  return currents;
}

std::array<wpi::units::ampere_t, numModules>
SwerveDriveSim::GetSteerCurrentDraw() const {
  std::array<wpi::units::ampere_t, numModules> currents;
  for (int i = 0; i < numModules; i++) {
    wpi::units::radians_per_second_t speed =
        wpi::units::radians_per_second_t{steerStates[i](1, 0) * steerGearing};
    // TODO: If uncommented we get huge current values.. Not sure how to fix
    // atm. :(
    currents[i] = 20_A;
    // currents[i] = GetCurrentDraw(steerMotor, speed, steerInputs[i],
    // wpi::RobotController::GetBatteryVoltage());
  }
  return currents;
}

wpi::units::ampere_t SwerveDriveSim::GetTotalCurrentDraw() const {
  wpi::units::ampere_t total{0};
  for (const auto& val : GetDriveCurrentDraw()) {
    total += val;
  }
  for (const auto& val : GetSteerCurrentDraw()) {
    total += val;
  }
  return total;
}