//! Recursive Newton-Euler algorithm (serial chains).

use nabled_core::scalar::NabledReal;
use nabled_kinematics::chain::ChainSpec;
use nabled_kinematics::fk::link_transforms_view;
use nabled_linalg::geometry::{Transform3, se3, so3};
use nabled_model::robot::RobotModel;
use ndarray::{Array1, Array2, ArrayView1};

use crate::DynamicsError;
use crate::config::DynamicsConfig;
use crate::spatial::{
    force_cross_product, joint_motion_subspace, motion_cross_product, spatial_gravity,
    spatial_inertia_6x6, spatial_inertia_apply,
};

fn adjoint_from_transform<T: NabledReal>(transform: &Transform3<T>) -> Array2<T> {
    let r = &transform.rotation.matrix;
    let p = &transform.translation;
    let px = so3::hat(&p.view());
    let mut ad = Array2::<T>::zeros((6, 6));
    for i in 0..3 {
        for j in 0..3 {
            ad[[i, j]] = r[[i, j]];
            ad[[i + 3, j]] = px[[i, j]];
            ad[[i + 3, j + 3]] = r[[i, j]];
        }
    }
    ad
}

fn default_inertial<T: NabledReal + Default>() -> nabled_model::link::InertialSpec<T> {
    nabled_model::link::InertialSpec {
        mass:    T::one(),
        com:     [T::default(), T::default(), T::default()],
        inertia: Array2::<T>::eye(3) * T::from_f64(0.01).unwrap_or(T::zero()),
    }
}

pub fn rnea<T: NabledReal + Default>(
    model: &RobotModel<T>,
    chain: &ChainSpec<T>,
    q: &ArrayView1<'_, T>,
    qd: &ArrayView1<'_, T>,
    qdd: &ArrayView1<'_, T>,
) -> Result<Array1<T>, DynamicsError> {
    rnea_with_config(model, chain, q, qd, qdd, &DynamicsConfig::default())
}

pub fn rnea_with_config<T: NabledReal + Default>(
    model: &RobotModel<T>,
    chain: &ChainSpec<T>,
    q: &ArrayView1<'_, T>,
    qd: &ArrayView1<'_, T>,
    qdd: &ArrayView1<'_, T>,
    config: &DynamicsConfig<T>,
) -> Result<Array1<T>, DynamicsError> {
    if model.dof() != chain.num_joints() || q.len() != chain.num_joints() {
        return Err(DynamicsError::DimensionMismatch);
    }
    if qd.len() != q.len() || qdd.len() != q.len() {
        return Err(DynamicsError::DimensionMismatch);
    }

    let transforms = link_transforms_view(chain, q)
        .map_err(|_| DynamicsError::InvalidInput("FK failed".to_string()))?;
    let num_joints = chain.num_joints();
    let indices = model.actuated_indices();
    if indices.len() != num_joints {
        return Err(DynamicsError::DimensionMismatch);
    }

    let mut link_vel = vec![Array1::<T>::zeros(6); num_joints + 1];
    let mut link_acc = vec![Array1::<T>::zeros(6); num_joints + 1];
    link_acc[0] = spatial_gravity(&config.gravity);

    let mut subspaces = Vec::with_capacity(num_joints);
    let mut inertias = Vec::with_capacity(num_joints);
    let mut adjoints = Vec::with_capacity(num_joints);

    for joint_idx in 0..num_joints {
        let body_idx = indices[joint_idx];
        let body = model.joint(body_idx).ok_or(DynamicsError::EmptyModel)?;
        let spec = body.inertial.clone().unwrap_or_else(default_inertial);
        inertias.push(spatial_inertia_6x6(&spec));
        subspaces.push(joint_motion_subspace(body.joint_type, body.axis));
        let parent_tf = &transforms[joint_idx];
        let child_tf = transforms.get(joint_idx + 1).ok_or(DynamicsError::EmptyModel)?;
        let relative = se3::compose(&se3::inverse(parent_tf), child_tf)
            .map_err(|_| DynamicsError::InvalidInput("transform compose failed".to_string()))?;
        adjoints.push(adjoint_from_transform(&relative));
    }

    for joint_idx in 0..num_joints {
        let subspace = &subspaces[joint_idx];
        let vj = subspace.mapv(|val| val * qd[joint_idx]);
        link_vel[joint_idx + 1] = adjoints[joint_idx].dot(&link_vel[joint_idx]) + &vj;
        let aj = subspace.mapv(|val| val * qdd[joint_idx]);
        link_acc[joint_idx + 1] = adjoints[joint_idx].dot(&link_acc[joint_idx])
            + aj
            + motion_cross_product(&link_vel[joint_idx + 1].view(), &vj.view());
    }

    let mut tau = Array1::<T>::zeros(num_joints);
    let mut f_child = Array1::<T>::zeros(6);
    for joint_idx in (0..num_joints).rev() {
        let iv = spatial_inertia_apply(&inertias[joint_idx], &link_vel[joint_idx + 1].view());
        let ia = spatial_inertia_apply(&inertias[joint_idx], &link_acc[joint_idx + 1].view());
        let local = ia + force_cross_product(&link_vel[joint_idx + 1].view(), &iv.view());
        let f_i = if joint_idx + 1 < num_joints {
            local + adjoints[joint_idx + 1].t().dot(&f_child)
        } else {
            local
        };
        tau[joint_idx] = subspaces[joint_idx].dot(&f_i);
        f_child = f_i;
    }

    Ok(tau)
}

pub fn rnea_view<T: NabledReal + Default>(
    model: &RobotModel<T>,
    chain: &ChainSpec<T>,
    q: &ArrayView1<'_, T>,
    qd: &ArrayView1<'_, T>,
    qdd: &ArrayView1<'_, T>,
) -> Result<Array1<T>, DynamicsError> {
    rnea(model, chain, q, qd, qdd)
}

pub fn rnea_into<T: NabledReal + Default>(
    model: &RobotModel<T>,
    chain: &ChainSpec<T>,
    q: &ArrayView1<'_, T>,
    qd: &ArrayView1<'_, T>,
    qdd: &ArrayView1<'_, T>,
    output: &mut Array1<T>,
) -> Result<(), DynamicsError> {
    let tau = rnea(model, chain, q, qd, qdd)?;
    if output.len() != tau.len() {
        return Err(DynamicsError::DimensionMismatch);
    }
    output.assign(&tau);
    Ok(())
}

#[cfg(test)]
mod tests {
    use approx::assert_relative_eq;
    use nabled_kinematics::chain::{ChainSpec, DhConvention, JointType};
    use nabled_model::fixture::Planar2rFixture;
    use nabled_model::joint::{JointAxis, JointType as ModelJointType};
    use nabled_model::robot::{BodySpec, DhParams};
    use ndarray::arr1;

    use super::*;

    fn load_fixture() -> Planar2rFixture {
        let path = concat!(
            env!("CARGO_MANIFEST_DIR"),
            "/../nabled/tests/fixtures/physical_ai/2r_planar.json"
        );
        Planar2rFixture::from_file(path).expect("fixture")
    }

    #[test]
    fn pendulum_gravity_torque() {
        let mut model = RobotModel::<f64>::new();
        let body = BodySpec {
            link:         nabled_model::link::LinkSpec { name: "link1".to_string() },
            parent_link:  "base".to_string(),
            joint_type:   ModelJointType::Revolute,
            axis:         JointAxis::Z,
            limits:       None,
            inertial:     Some(nabled_model::link::InertialSpec {
                mass:    1.0,
                com:     [0.5, 0.0, 0.0],
                inertia: Array2::<f64>::zeros((3, 3)),
            }),
            joint_origin: nabled_model::origin::identity_transform(),
            dh_params:    Some(DhParams {
                a:            0.0,
                alpha:        0.0,
                d:            0.0,
                theta_offset: 0.0,
            }),
        };
        let _ = model.add_body(None, body);
        let chain = ChainSpec::from_dh(
            DhConvention::Standard,
            vec![JointType::Revolute],
            arr1(&[0.0]),
            arr1(&[0.0]),
            arr1(&[0.0]),
            arr1(&[0.0]),
        )
        .unwrap();
        let q = arr1(&[0.0_f64]);
        let zeros = arr1(&[0.0]);
        let config = DynamicsConfig { gravity: [0.0, -9.81, 0.0], ..DynamicsConfig::default() };
        let tau =
            rnea_with_config(&model, &chain, &q.view(), &zeros.view(), &zeros.view(), &config)
                .unwrap();
        assert!(tau[0].abs() > 0.1, "expected gravity torque, got {}", tau[0]);
    }

    #[test]
    fn planar2r_matches_fixture_when_present() {
        let fixture = load_fixture();
        if fixture.cases.iter().all(|c| c.tau.is_none()) {
            return;
        }
        let model = fixture.to_robot_model().unwrap();
        let chain = fixture.to_chain_spec().unwrap();
        let gravity: [f64; 3] = fixture.gravity.unwrap_or([0.0, -9.81, 0.0]);
        let config = DynamicsConfig { gravity, ..DynamicsConfig::default() };
        for case in &fixture.cases {
            if let (Some(qd), Some(qdd), Some(tau_ref)) = (&case.qd, &case.qdd, &case.tau) {
                let q = arr1(&case.q);
                let qd = arr1(qd);
                let qdd = arr1(qdd);
                let tau =
                    rnea_with_config(&model, &chain, &q.view(), &qd.view(), &qdd.view(), &config)
                        .unwrap();
                for (computed, expected) in tau.iter().zip(tau_ref.iter()) {
                    assert_relative_eq!(computed, expected, epsilon = 1e-6);
                }
            }
        }
    }

    #[test]
    fn rnea_rejects_dimension_mismatch() {
        let fixture = load_fixture();
        let model = fixture.to_robot_model().unwrap();
        let chain = fixture.to_chain_spec().unwrap();
        let q = arr1(&[0.0_f64, 0.0]);
        let qd = arr1(&[0.0_f64]);
        let qdd = arr1(&[0.0_f64, 0.0]);
        assert!(matches!(
            rnea(&model, &chain, &q.view(), &qd.view(), &qdd.view()),
            Err(DynamicsError::DimensionMismatch)
        ));
    }

    #[test]
    fn prismatic_joint_gravity_force() {
        let mut model = RobotModel::<f64>::new();
        let body = BodySpec {
            link:         nabled_model::link::LinkSpec { name: "slide".to_string() },
            parent_link:  "base".to_string(),
            joint_type:   ModelJointType::Prismatic,
            axis:         JointAxis::Z,
            limits:       None,
            inertial:     Some(nabled_model::link::InertialSpec {
                mass:    1.0,
                com:     [0.0, 0.0, 0.0],
                inertia: Array2::<f64>::zeros((3, 3)),
            }),
            joint_origin: nabled_model::origin::identity_transform(),
            dh_params:    Some(DhParams {
                a:            0.0,
                alpha:        0.0,
                d:            0.0,
                theta_offset: 0.0,
            }),
        };
        let _ = model.add_body(None, body);
        let chain = ChainSpec::from_dh(
            DhConvention::Standard,
            vec![JointType::Prismatic],
            arr1(&[0.0]),
            arr1(&[0.0]),
            arr1(&[0.0]),
            arr1(&[0.0]),
        )
        .unwrap();
        let q = arr1(&[0.0_f64]);
        let zeros = arr1(&[0.0]);
        let config = DynamicsConfig { gravity: [0.0, 0.0, -9.81], ..DynamicsConfig::default() };
        let tau =
            rnea_with_config(&model, &chain, &q.view(), &zeros.view(), &zeros.view(), &config)
                .unwrap();
        assert!(tau[0].abs() > 0.1, "expected gravity force on prismatic joint, got {}", tau[0]);
    }

    #[test]
    fn rnea_into_writes_output_buffer() {
        let fixture = load_fixture();
        let model = fixture.to_robot_model().unwrap();
        let chain = fixture.to_chain_spec().unwrap();
        let q = arr1(&[0.3_f64, 0.5]);
        let zeros = arr1(&[0.0, 0.0]);
        let mut output = arr1(&[0.0, 0.0]);
        rnea_into(&model, &chain, &q.view(), &zeros.view(), &zeros.view(), &mut output).unwrap();
        let tau = rnea_view(&model, &chain, &q.view(), &zeros.view(), &zeros.view()).unwrap();
        assert_relative_eq!(output, tau, epsilon = 1e-12);
    }

    #[test]
    fn rnea_into_rejects_output_dimension_mismatch() {
        let fixture = load_fixture();
        let model = fixture.to_robot_model().unwrap();
        let chain = fixture.to_chain_spec().unwrap();
        let q = arr1(&[0.3_f64, 0.5]);
        let zeros = arr1(&[0.0, 0.0]);
        let mut output = arr1(&[0.0]);
        assert!(matches!(
            rnea_into(&model, &chain, &q.view(), &zeros.view(), &zeros.view(), &mut output),
            Err(DynamicsError::DimensionMismatch)
        ));
    }
}
