// Copyright 2025 International Digital Economy Academy
//
// 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.

///|
const REAL_MAX : @core.Real = 1.0e30F

///|
pub(all) enum JointAxis {
  LinX
  LinY
  LinZ
  AngX
  AngY
  AngZ
}

///|
pub struct JointAxesMask {
  bits : Int
}

///|
pub fn JointAxesMask::empty() -> JointAxesMask {
  { bits: 0 }
}

///|
pub fn JointAxesMask::JointAxesMask(bits : Int) -> JointAxesMask {
  { bits, }
}

///|
pub fn JointAxesMask::bits(self : JointAxesMask) -> Int {
  self.bits
}

///|
pub fn JointAxesMask::contains(
  self : JointAxesMask,
  other : JointAxesMask,
) -> Bool {
  (self.bits & other.bits) == other.bits
}

///|
pub fn JointAxesMask::or(
  self : JointAxesMask,
  other : JointAxesMask,
) -> JointAxesMask {
  { bits: self.bits | other.bits }
}

///|
pub fn JointAxesMask::intersection(
  self : JointAxesMask,
  other : JointAxesMask,
) -> JointAxesMask {
  { bits: self.bits & other.bits }
}

///|
pub fn JointAxesMask::lin_x() -> JointAxesMask {
  { bits: 1 << 0 }
}

///|
pub fn JointAxesMask::lin_y() -> JointAxesMask {
  { bits: 1 << 1 }
}

///|
pub fn JointAxesMask::lin_z() -> JointAxesMask {
  { bits: 1 << 2 }
}

///|
pub fn JointAxesMask::ang_x() -> JointAxesMask {
  { bits: 1 << 3 }
}

///|
pub fn JointAxesMask::ang_y() -> JointAxesMask {
  { bits: 1 << 4 }
}

///|
pub fn JointAxesMask::ang_z() -> JointAxesMask {
  { bits: 1 << 5 }
}

///|
pub fn JointAxesMask::lin_axes() -> JointAxesMask {
  JointAxesMask::lin_x().or(JointAxesMask::lin_y())
}

///|
pub fn JointAxesMask::ang_axes() -> JointAxesMask {
  JointAxesMask::ang_z()
}

///|
pub fn JointAxesMask::lin_axes3() -> JointAxesMask {
  JointAxesMask::lin_axes().or(JointAxesMask::lin_z())
}

///|
pub fn JointAxesMask::ang_axes3() -> JointAxesMask {
  JointAxesMask::ang_x().or(JointAxesMask::ang_y()).or(JointAxesMask::ang_z())
}

///|
pub fn JointAxesMask::all() -> JointAxesMask {
  JointAxesMask::lin_axes().or(JointAxesMask::ang_axes())
}

///|
pub fn JointAxesMask::all3() -> JointAxesMask {
  JointAxesMask::lin_axes3().or(JointAxesMask::ang_axes3())
}

///|
pub fn JointAxesMask::locked_revolute_axes() -> JointAxesMask {
  // 2D hinge: lock both linear axes; keep angular free.
  JointAxesMask::lin_axes()
}

///|
pub fn JointAxesMask::locked_revolute_axes3() -> JointAxesMask {
  JointAxesMask::lin_axes3()
  .or(JointAxesMask::ang_y())
  .or(JointAxesMask::ang_z())
}

///|
pub fn JointAxesMask::locked_prismatic_axes() -> JointAxesMask {
  JointAxesMask::lin_y().or(JointAxesMask::ang_z())
}

///|
pub fn JointAxesMask::locked_prismatic_axes3() -> JointAxesMask {
  JointAxesMask::lin_y()
  .or(JointAxesMask::lin_z())
  .or(JointAxesMask::ang_x())
  .or(JointAxesMask::ang_y())
  .or(JointAxesMask::ang_z())
}

///|
pub fn JointAxesMask::locked_pin_slot_axes() -> JointAxesMask {
  JointAxesMask(1 << 1)
}

///|
pub fn JointAxesMask::locked_fixed_axes() -> JointAxesMask {
  JointAxesMask::lin_axes().or(JointAxesMask::ang_z())
}

///|
pub fn JointAxesMask::locked_fixed_axes3() -> JointAxesMask {
  JointAxesMask::all3()
}

///|
pub fn JointAxesMask::locked_spherical_axes() -> JointAxesMask {
  JointAxesMask::lin_axes3()
}

///|
pub fn JointAxesMask::free_revolute_axes() -> JointAxesMask {
  JointAxesMask::ang_z()
}

///|
pub fn JointAxesMask::free_revolute_axes3() -> JointAxesMask {
  JointAxesMask::ang_x()
}

///|
pub fn JointAxesMask::free_prismatic_axes() -> JointAxesMask {
  JointAxesMask::lin_x()
}

///|
pub fn JointAxesMask::free_prismatic_axes3() -> JointAxesMask {
  JointAxesMask::lin_x()
}

///|
pub fn JointAxesMask::free_fixed_axes() -> JointAxesMask {
  JointAxesMask::empty()
}

///|
pub fn JointAxesMask::free_spherical_axes() -> JointAxesMask {
  JointAxesMask::ang_axes3()
}

///|
pub struct JointLimits {
  min : @core.Real
  max : @core.Real
  mut impulse : @core.Real
}

///|
pub fn JointLimits::default() -> JointLimits {
  { min: -REAL_MAX, max: REAL_MAX, impulse: 0.0F }
}

///|
pub fn JointLimits::JointLimits(
  min : @core.Real,
  max : @core.Real,
) -> JointLimits {
  { min, max, impulse: 0.0F }
}

///|
pub fn JointLimits::min(self : JointLimits) -> @core.Real {
  self.min
}

///|
pub fn JointLimits::max(self : JointLimits) -> @core.Real {
  self.max
}

///|
pub fn JointLimits::impulse(self : JointLimits) -> @core.Real {
  self.impulse
}

///|
pub fn JointLimits::set_min(
  self : JointLimits,
  min : @core.Real,
) -> JointLimits {
  { min, max: self.max, impulse: self.impulse }
}

///|
pub fn JointLimits::set_max(
  self : JointLimits,
  max : @core.Real,
) -> JointLimits {
  { min: self.min, max, impulse: self.impulse }
}

///|
pub fn JointLimits::clear_impulse(self : JointLimits) -> JointLimits {
  self.impulse = 0.0F
  self
}

///|
pub fn JointLimits::set_impulse(
  self : JointLimits,
  impulse : @core.Real,
) -> JointLimits {
  self.impulse = impulse
  self
}

///|
pub struct JointLimits3 {
  lin_x : JointLimits
  lin_y : JointLimits
  lin_z : JointLimits
  ang_x : JointLimits
  ang_y : JointLimits
  ang_z : JointLimits
}

///|
pub fn JointLimits3::default() -> JointLimits3 {
  {
    lin_x: JointLimits::default(),
    lin_y: JointLimits::default(),
    lin_z: JointLimits::default(),
    ang_x: JointLimits::default(),
    ang_y: JointLimits::default(),
    ang_z: JointLimits::default(),
  }
}

///|
pub(all) enum MotorModel {
  AccelerationBased
  ForceBased
}

///|
pub fn MotorModel::combine_coefficients(
  self : MotorModel,
  dt : @core.Real,
  stiffness : @core.Real,
  damping : @core.Real,
) -> (@core.Real, @core.Real, @core.Real) {
  fn inv(x : @core.Real) -> @core.Real {
    if x == 0.0F {
      0.0F
    } else {
      1.0F / x
    }
  }

  let denom = dt * stiffness + damping
  let erp_inv_dt = stiffness * inv(denom)
  let pgs_denom = dt * dt * stiffness + dt * damping
  match self {
    AccelerationBased => (erp_inv_dt, inv(pgs_denom), 0.0F)
    ForceBased => (erp_inv_dt, 0.0F, inv(pgs_denom))
  }
}

///|
pub struct JointMotor {
  target_vel : @core.Real
  target_pos : @core.Real
  stiffness : @core.Real
  damping : @core.Real
  max_force : @core.Real
  mut impulse : @core.Real
  model : MotorModel
}

///|
pub fn JointMotor::default() -> JointMotor {
  {
    target_pos: 0.0F,
    target_vel: 0.0F,
    stiffness: 0.0F,
    damping: 0.0F,
    max_force: REAL_MAX,
    impulse: 0.0F,
    model: AccelerationBased,
  }
}

///|
pub fn JointMotor::set_target_pos(
  self : JointMotor,
  target_pos : @core.Real,
) -> JointMotor {
  {
    target_vel: self.target_vel,
    target_pos,
    stiffness: self.stiffness,
    damping: self.damping,
    max_force: self.max_force,
    impulse: self.impulse,
    model: self.model,
  }
}

///|
pub fn JointMotor::set_target_vel(
  self : JointMotor,
  target_vel : @core.Real,
) -> JointMotor {
  {
    target_vel,
    target_pos: self.target_pos,
    stiffness: self.stiffness,
    damping: self.damping,
    max_force: self.max_force,
    impulse: self.impulse,
    model: self.model,
  }
}

///|
pub fn JointMotor::set_stiffness(
  self : JointMotor,
  stiffness : @core.Real,
) -> JointMotor {
  {
    target_vel: self.target_vel,
    target_pos: self.target_pos,
    stiffness,
    damping: self.damping,
    max_force: self.max_force,
    impulse: self.impulse,
    model: self.model,
  }
}

///|
pub fn JointMotor::set_damping(
  self : JointMotor,
  damping : @core.Real,
) -> JointMotor {
  {
    target_vel: self.target_vel,
    target_pos: self.target_pos,
    stiffness: self.stiffness,
    damping,
    max_force: self.max_force,
    impulse: self.impulse,
    model: self.model,
  }
}

///|
pub fn JointMotor::set_max_force(
  self : JointMotor,
  max_force : @core.Real,
) -> JointMotor {
  {
    target_vel: self.target_vel,
    target_pos: self.target_pos,
    stiffness: self.stiffness,
    damping: self.damping,
    max_force,
    impulse: self.impulse,
    model: self.model,
  }
}

///|
pub fn JointMotor::set_model(
  self : JointMotor,
  model : MotorModel,
) -> JointMotor {
  {
    target_vel: self.target_vel,
    target_pos: self.target_pos,
    stiffness: self.stiffness,
    damping: self.damping,
    max_force: self.max_force,
    impulse: self.impulse,
    model,
  }
}

///|
pub fn JointMotor::clear_impulse(self : JointMotor) -> JointMotor {
  self.impulse = 0.0F
  self
}

///|
pub fn JointMotor::impulse(self : JointMotor) -> @core.Real {
  self.impulse
}

///|
pub fn JointMotor::set_impulse(
  self : JointMotor,
  impulse : @core.Real,
) -> JointMotor {
  self.impulse = impulse
  self
}

///|
pub struct JointMotors3 {
  lin_x : JointMotor
  lin_y : JointMotor
  lin_z : JointMotor
  ang_x : JointMotor
  ang_y : JointMotor
  ang_z : JointMotor
}

///|
pub fn JointMotors3::default() -> JointMotors3 {
  {
    lin_x: JointMotor::default(),
    lin_y: JointMotor::default(),
    lin_z: JointMotor::default(),
    ang_x: JointMotor::default(),
    ang_y: JointMotor::default(),
    ang_z: JointMotor::default(),
  }
}