// Integration tests for physics solvers
// Following strict TDD - RED phase first

#[cfg(test)]
mod tests {
    use crate::solver::{Solver, SolverConfig};
    use rtx_tensor::{Tensor, Device, DType};
    use rtx_autograd::variable::Variable;

    #[test]
    fn test_harmonic_oscillator_solver() {
        let device = Device::cpu();

        // Simple harmonic oscillator: d²x/dt² = -ω²x
        // Solution: x(t) = A*cos(ωt) + B*sin(ωt)

        let omega = 2.0; // Angular frequency
        let dt = 0.01; // Time step
        let t_final = 10.0;
        let steps = (t_final / dt) as usize;

        // Initial conditions: x(0) = 1, dx/dt(0) = 0
        let x0 = Tensor::scalar(1.0, DType::F32, &device).unwrap();
        let v0 = Tensor::scalar(0.0, DType::F32, &device).unwrap();

        // Create solver configuration
        let config = SolverConfig {
            method: "rk4".to_string(),
            tolerance: 1e-6,
            max_iterations: 1000,
            adaptive_timestep: false,
        };

        let solver = Solver::new(config);

        // Define the system dynamics
        let dynamics = |x: &Tensor, _t: f32| -> Tensor {
            x.mul_scalar(-omega * omega).unwrap()
        };

        // Solve the ODE
        let mut x = x0.clone();
        let mut v = v0.clone();
        let mut trajectory = vec![x.to_vec().unwrap()[0]];

        for i in 0..steps {
            let t = i as f32 * dt;

            // Update using velocity
            x = x.add(&v.mul_scalar(dt).unwrap()).unwrap();

            // Update velocity using acceleration
            let accel = dynamics(&x, t);
            v = v.add(&accel.mul_scalar(dt).unwrap()).unwrap();

            trajectory.push(x.to_vec().unwrap()[0]);
        }

        // Verify solution is periodic
        let period = 2.0 * std::f32::consts::PI / omega;
        let cycles = t_final / period;

        // Check that we completed full cycles
        assert!(cycles > 1.0);

        // Final position should be close to initial (after full periods)
        let final_x = trajectory.last().unwrap();
        let initial_x = trajectory.first().unwrap();

        // Allow for numerical error accumulation
        assert!((final_x - initial_x).abs() < 0.1);
    }

    #[test]
    fn test_heat_equation_solver() {
        let device = Device::cpu();

        // 1D heat equation: ∂u/∂t = α * ∂²u/∂x²
        // Domain: [0, 1], initial condition: gaussian pulse

        let nx = 50; // Spatial points
        let dx = 1.0 / (nx as f32 - 1.0);
        let alpha = 0.01; // Thermal diffusivity
        let dt = 0.5 * dx * dx / alpha; // CFL condition
        let t_final = 0.1;
        let nt = (t_final / dt) as usize;

        // Initial condition: Gaussian
        let mut x_vals = Vec::new();
        let mut u_vals = Vec::new();
        for i in 0..nx {
            let x = i as f32 * dx;
            x_vals.push(x);
            let u = (-50.0 * (x - 0.5) * (x - 0.5)).exp();
            u_vals.push(u);
        }

        let u = Tensor::from_vec(u_vals.clone(), &[nx], &device).unwrap();

        // Time evolution
        let mut current_u = u;
        for _ in 0..nt {
            // Compute second derivative using finite differences
            let u_vec = current_u.to_vec().unwrap();
            let mut laplacian = vec![0.0; nx];

            for i in 1..(nx - 1) {
                laplacian[i] = (u_vec[i + 1] - 2.0 * u_vec[i] + u_vec[i - 1]) / (dx * dx);
            }

            // Boundary conditions (Dirichlet: u = 0 at boundaries)
            laplacian[0] = 0.0;
            laplacian[nx - 1] = 0.0;

            let lap_tensor = Tensor::from_vec(laplacian, &[nx], &device).unwrap();

            // Update using forward Euler
            current_u = current_u.add(&lap_tensor.mul_scalar(alpha * dt).unwrap()).unwrap();
        }

        // Verify heat has diffused (peak should be lower)
        let final_vals = current_u.to_vec().unwrap();
        let initial_max = u_vals.iter().fold(0.0f32, |a, &b| a.max(b));
        let final_max = final_vals.iter().fold(0.0f32, |a, &b| a.max(b));

        assert!(final_max < initial_max);
        assert!(final_max > 0.0); // Should still have some heat
    }

    #[test]
    fn test_wave_equation_solver() {
        let device = Device::cpu();

        // 1D wave equation: ∂²u/∂t² = c² * ∂²u/∂x²
        let nx = 100;
        let dx = 1.0 / (nx as f32 - 1.0);
        let c = 1.0; // Wave speed
        let dt = 0.5 * dx / c; // CFL condition
        let t_final = 2.0;
        let nt = (t_final / dt) as usize;

        // Initial conditions: Gaussian pulse, zero velocity
        let mut u_vals = Vec::new();
        for i in 0..nx {
            let x = i as f32 * dx;
            let u = (-100.0 * (x - 0.5) * (x - 0.5)).exp();
            u_vals.push(u);
        }

        let u_prev = Tensor::from_vec(u_vals.clone(), &[nx], &device).unwrap();
        let mut u_curr = u_prev.clone();
        let mut u_next;

        // Time evolution using finite differences
        for _ in 0..nt {
            let u_vec = u_curr.to_vec().unwrap();
            let u_prev_vec = u_prev.to_vec().unwrap();
            let mut new_u = vec![0.0; nx];

            for i in 1..(nx - 1) {
                let d2u_dx2 = (u_vec[i + 1] - 2.0 * u_vec[i] + u_vec[i - 1]) / (dx * dx);
                new_u[i] = 2.0 * u_vec[i] - u_prev_vec[i] + c * c * dt * dt * d2u_dx2;
            }

            // Boundary conditions (fixed ends)
            new_u[0] = 0.0;
            new_u[nx - 1] = 0.0;

            u_next = Tensor::from_vec(new_u, &[nx], &device).unwrap();

            // Update for next iteration
            u_prev = u_curr;
            u_curr = u_next;
        }

        // Wave should have propagated
        let final_vals = u_curr.to_vec().unwrap();

        // Check energy is approximately conserved (with some numerical dissipation)
        let initial_energy: f32 = u_vals.iter().map(|x| x * x).sum();
        let final_energy: f32 = final_vals.iter().map(|x| x * x).sum();

        // Allow up to 20% energy loss due to numerical dissipation
        assert!(final_energy > 0.8 * initial_energy);
    }

    #[test]
    fn test_gradient_flow_optimization() {
        let device = Device::cpu();

        // Minimize f(x, y) = x² + 2y² using gradient descent
        let mut x = Variable::new(Tensor::scalar(3.0, DType::F32, &device).unwrap(), true);
        let mut y = Variable::new(Tensor::scalar(4.0, DType::F32, &device).unwrap(), true);

        let learning_rate = 0.1;
        let iterations = 100;

        for _ in 0..iterations {
            // Compute f = x² + 2y²
            let x_sq = x.multiply(&x).unwrap();
            let y_sq = y.multiply(&y).unwrap();
            let two_y_sq = y_sq.multiply_scalar(2.0).unwrap();
            let f = x_sq.add(&two_y_sq).unwrap();

            // Compute gradients
            f.backward(None).unwrap();

            // Get gradients
            let x_grad = x.grad().unwrap();
            let y_grad = y.grad().unwrap();

            // Update parameters
            let x_tensor = x.tensor();
            let y_tensor = y.tensor();

            let new_x_tensor = x_tensor.sub(&x_grad.mul_scalar(learning_rate).unwrap()).unwrap();
            let new_y_tensor = y_tensor.sub(&y_grad.mul_scalar(learning_rate).unwrap()).unwrap();

            // Create new variables
            x = Variable::new(new_x_tensor, true);
            y = Variable::new(new_y_tensor, true);
        }

        // Should converge to (0, 0)
        let final_x = x.tensor().to_vec().unwrap()[0];
        let final_y = y.tensor().to_vec().unwrap()[0];

        assert!(final_x.abs() < 0.01);
        assert!(final_y.abs() < 0.01);
    }

    #[test]
    fn test_nonlinear_system_solver() {
        let device = Device::cpu();

        // Solve nonlinear system:
        // x² + y² = 1 (circle)
        // y = x² (parabola)
        // Solutions: approximately (0.786, 0.618) and (-0.786, 0.618)

        let mut x = Variable::new(Tensor::scalar(0.5, DType::F32, &device).unwrap(), true);
        let mut y = Variable::new(Tensor::scalar(0.5, DType::F32, &device).unwrap(), true);

        let iterations = 50;
        let lr = 0.01;

        for _ in 0..iterations {
            // Compute residuals
            let x_sq = x.multiply(&x).unwrap();
            let y_sq = y.multiply(&y).unwrap();

            // r1 = x² + y² - 1
            let sum_sq = x_sq.add(&y_sq).unwrap();
            let one = Variable::new(Tensor::scalar(1.0, DType::F32, &device).unwrap(), false);
            let r1 = sum_sq.add(&one.multiply_scalar(-1.0).unwrap()).unwrap();

            // r2 = y - x²
            let r2 = y.add(&x_sq.multiply_scalar(-1.0).unwrap()).unwrap();

            // Loss = r1² + r2²
            let r1_sq = r1.multiply(&r1).unwrap();
            let r2_sq = r2.multiply(&r2).unwrap();
            let loss = r1_sq.add(&r2_sq).unwrap();

            // Backward pass
            loss.backward(None).unwrap();

            // Update
            let x_grad = x.grad().unwrap();
            let y_grad = y.grad().unwrap();

            let new_x = x.tensor().sub(&x_grad.mul_scalar(lr).unwrap()).unwrap();
            let new_y = y.tensor().sub(&y_grad.mul_scalar(lr).unwrap()).unwrap();

            x = Variable::new(new_x, true);
            y = Variable::new(new_y, true);
        }

        // Check solution satisfies constraints
        let x_val = x.tensor().to_vec().unwrap()[0];
        let y_val = y.tensor().to_vec().unwrap()[0];

        // Check x² + y² ≈ 1
        let circle_error = (x_val * x_val + y_val * y_val - 1.0).abs();
        assert!(circle_error < 0.1);

        // Check y ≈ x²
        let parabola_error = (y_val - x_val * x_val).abs();
        assert!(parabola_error < 0.1);
    }
}