use lmm::causal::CausalGraph;
use lmm::compression::compute_mse;
use lmm::consciousness::Consciousness;
use lmm::discovery::SymbolicRegression;
use lmm::encode::{decode_from_parts, decode_message, encode_text};
use lmm::equation::Expression;
use lmm::field::Field;
use lmm::operator::NeuralOperator;
use lmm::physics::{HarmonicOscillator, LorenzSystem, SIRModel};
use lmm::predict::TextPredictor;
use lmm::simulation::Simulator;
use lmm::tensor::Tensor;
use lmm::traits::{Causal, Discoverable, Simulatable};
use lmm::world::WorldModel;
use std::collections::HashMap;

#[test]
fn test_tensor_math() {
    let t1 = Tensor::new(vec![2], vec![1.0, 2.0]).unwrap();
    let t2 = Tensor::new(vec![2], vec![3.0, 4.0]).unwrap();
    let sum = (&t1 + &t2).unwrap();
    assert_eq!(sum.data, vec![4.0, 6.0]);
    let diff = (&t2 - &t1).unwrap();
    assert_eq!(diff.data, vec![2.0, 2.0]);
}

#[test]
fn test_simulation_rk4() {
    let osc = HarmonicOscillator::new(1.0, 1.0, 0.0).unwrap();
    let sim = Simulator { step_size: 0.01 };
    let s1 = sim.rk4_step(&osc, osc.state()).unwrap();
    assert!((s1.data[0] - 0.99995).abs() < 1e-4);
}

#[test]
fn test_simulate_trajectory() {
    let osc = HarmonicOscillator::new(1.0, 1.0, 0.0).unwrap();
    let sim = Simulator { step_size: 0.01 };
    let traj = sim.simulate_trajectory(&osc, osc.state(), 100).unwrap();
    assert_eq!(traj.len(), 101);
}

#[test]
fn test_equation_evaluation() {
    let eq = Expression::Add(
        Box::new(Expression::Variable("x".into())),
        Box::new(Expression::Constant(2.0)),
    );
    let mut vars = HashMap::new();
    vars.insert("x".into(), 5.0);
    assert_eq!(eq.evaluate(&vars).unwrap(), 7.0);
}

#[test]
fn test_equation_symbolic_diff_and_simplify() {
    let eq = Expression::Mul(
        Box::new(Expression::Constant(3.0)),
        Box::new(Expression::Variable("x".into())),
    );
    let d = eq.symbolic_diff("x").simplify();
    let mut vars = HashMap::new();
    vars.insert("x".into(), 99.0);
    assert!((d.evaluate(&vars).unwrap() - 3.0).abs() < 1e-10);
}

#[test]
fn test_symbolic_regression_pipeline() {
    let inputs: Vec<Vec<f64>> = (0..12).map(|i| vec![i as f64]).collect();
    let targets: Vec<f64> = (0..12).map(|i| 3.0 * i as f64).collect();
    let sr = SymbolicRegression::new(3, 30).with_variables(vec!["x".into()]);
    let expr = sr.fit(&inputs, &targets).unwrap();
    let mse = compute_mse(&expr, &inputs, &targets);
    let mean_y = targets.iter().sum::<f64>() / targets.len() as f64;
    let baseline =
        targets.iter().map(|&y| (y - mean_y).powi(2)).sum::<f64>() / targets.len() as f64;
    assert!(
        mse < baseline,
        "GP must beat constant baseline: mse={mse:.2}, baseline={baseline:.2}"
    );
}

#[test]
fn test_discover_trait() {
    let data: Vec<Tensor> = (0..6).map(|i| Tensor::from_vec(vec![i as f64])).collect();
    let targets: Vec<f64> = (0..6).map(|i| i as f64).collect();
    let expr = SymbolicRegression::discover(&data, &targets).unwrap();
    assert!(expr.complexity() > 0);
}

#[test]
fn test_neural_operator_identity() {
    let mut kw = vec![0.0; 3];
    kw[1] = 1.0;
    let op = NeuralOperator { kernel_weights: kw };
    let field = Field::new(vec![3], Tensor::new(vec![3], vec![1.0, 2.0, 3.0]).unwrap()).unwrap();
    let out = op.transform(&field).unwrap();
    for (a, b) in out.values.data.iter().zip(field.values.data.iter()) {
        assert!((a - b).abs() < 1e-10);
    }
}

#[test]
fn test_world_model_step() {
    let mut wm = WorldModel::new(Tensor::new(vec![2], vec![0.0, 0.0]).unwrap());
    let action = Tensor::new(vec![2], vec![1.0, 2.0]).unwrap();
    let ns = wm.step(&action).unwrap();
    assert_eq!(ns.data, vec![1.0, 2.0]);
}

#[test]
fn test_world_model_prediction_error() {
    let mut wm = WorldModel::new(Tensor::new(vec![2], vec![0.0, 0.0]).unwrap());
    let predicted = Tensor::new(vec![2], vec![1.0, 1.0]).unwrap();
    let actual = Tensor::new(vec![2], vec![1.5, 0.5]).unwrap();
    let err = wm.record_error(&predicted, &actual).unwrap();
    assert!((err - 0.25).abs() < 1e-10);
    assert!((wm.mean_prediction_error() - 0.25).abs() < 1e-10);
}

#[test]
fn test_consciousness_tick() {
    let init = Tensor::zeros(vec![4]);
    let mut consc = Consciousness::new(init, 3, 0.01);
    let input = vec![255u8, 127, 0, 64];
    let ns = consc.tick(&input).unwrap();
    assert_eq!(ns.shape, vec![4]);
}

#[test]
fn test_causal_end_to_end() {
    let mut g = CausalGraph::new();
    let x = g.add_node("x", Some(Expression::Constant(4.0)));
    let y = g.add_node(
        "y",
        Some(Expression::Mul(
            Box::new(Expression::Constant(3.0)),
            Box::new(Expression::Variable("x".into())),
        )),
    );
    g.add_edge(x, y, 1.0).unwrap();
    let vals = g.forward_pass().unwrap();
    assert!((vals[&y] - 12.0).abs() < 1e-10);
    g.intervene("x", 10.0).unwrap();
    let after = g.forward_pass().unwrap();
    assert!((after[&y] - 30.0).abs() < 1e-10);
}

#[test]
fn test_lorenz_produces_3d_trajectory() {
    let sys = LorenzSystem::canonical().unwrap();
    let sim = Simulator { step_size: 0.01 };
    let traj = sim.simulate_trajectory(&sys, sys.state(), 100).unwrap();
    assert_eq!(traj.first().unwrap().shape, vec![3]);
    assert_eq!(traj.last().unwrap().shape, vec![3]);
}

#[test]
fn test_sir_full_pipeline() {
    let sir = SIRModel::new(0.3, 0.1, 990.0, 10.0, 0.0).unwrap();
    let n0 = sir.total_population();
    let sim = Simulator { step_size: 0.5 };
    let traj = sim.simulate_trajectory(&sir, sir.state(), 200).unwrap();
    let nf: f64 = traj.last().unwrap().data.iter().sum();
    assert!(
        (nf - n0).abs() < 20.0,
        "SIR population drift too large: {}",
        (nf - n0).abs()
    );
    let peak_i = traj.iter().map(|s| s.data[1]).fold(0.0f64, f64::max);
    assert!(peak_i > 10.0, "Infection should peak above initial");
}

#[test]
fn test_encode_decode_roundtrip() {
    let text = "Hi LMM";
    let encoded = encode_text(text, 30, 3).unwrap();
    let decoded = decode_message(&encoded).unwrap();
    assert_eq!(
        decoded, text,
        "Round-trip must recover original text exactly"
    );
}

#[test]
fn test_decode_from_parts_known() {
    let eq = Expression::Constant(65.0);
    let decoded = decode_from_parts(&eq, 3, &[0, 1, 2]).unwrap();
    assert_eq!(decoded.as_bytes(), &[65u8, 66, 67]);
}

#[test]
fn test_field_gradient_1d() {
    let data: Vec<f64> = (0..8).map(|i| (i as f64).powi(2)).collect();
    let tensor = Tensor::new(vec![8], data).unwrap();
    let field = Field::new(vec![8], tensor).unwrap();
    let grad = field.compute_gradient().unwrap();
    assert!((grad.values.data[4] - 8.0).abs() < 1.0);
}

#[test]
fn test_field_laplacian_1d() {
    let data: Vec<f64> = (0..8).map(|i| (i as f64).powi(2)).collect();
    let tensor = Tensor::new(vec![8], data).unwrap();
    let field = Field::new(vec![8], tensor).unwrap();
    let lap = field.compute_laplacian().unwrap();
    for v in &lap.values.data[1..7] {
        assert!((v - 2.0).abs() < 1e-6, "Expected ≈2 got {v}");
    }
}

#[test]
fn test_causal_intervention_pipeline() {
    let mut graph = CausalGraph::new();
    let x_id = graph.add_node("x", Some(Expression::Constant(3.0)));
    let y_id = graph.add_node(
        "y",
        Some(Expression::Mul(
            Box::new(Expression::Constant(2.0)),
            Box::new(Expression::Variable("x".into())),
        )),
    );
    let z_id = graph.add_node(
        "z",
        Some(Expression::Add(
            Box::new(Expression::Variable("y".into())),
            Box::new(Expression::Constant(1.0)),
        )),
    );
    graph.add_edge(x_id, y_id, 1.0).unwrap();
    graph.add_edge(y_id, z_id, 1.0).unwrap();
    graph.nodes[x_id].observed_value = Some(3.0);
    let before = graph.forward_pass().unwrap();
    assert!((before[&y_id] - 6.0).abs() < 1e-10);
    assert!((before[&z_id] - 7.0).abs() < 1e-10);
    graph.intervene("x", 10.0).unwrap();
    let after = graph.forward_pass().unwrap();
    assert!((after[&y_id] - 20.0).abs() < 1e-10);
    assert!((after[&z_id] - 21.0).abs() < 1e-10);
}

#[test]
fn test_predict_produces_real_words() {
    let input = "Wise AI built the first LMM framework";
    let predictor = TextPredictor::new(10, 20, 3);
    let result = predictor.predict_continuation(input, 30).unwrap();
    let input_words: std::collections::BTreeSet<&str> = input.split_whitespace().collect();
    for word in result.continuation.split_whitespace() {
        assert!(
            input_words.contains(word),
            "Predicted '{word}' not in input vocabulary"
        );
    }
}

#[test]
fn test_predict_has_multiple_words() {
    let input = "Large Mathematical Models compress the world into equations";
    let predictor = TextPredictor::new(10, 20, 3);
    let result = predictor.predict_continuation(input, 40).unwrap();
    let word_count = result.continuation.split_whitespace().count();
    assert!(word_count >= 3, "Expected >=3 words, got {word_count}");
}

#[test]
fn test_predict_trajectory_equation_has_variable() {
    let input = "The Pharaohs encoded reality in mathematics";
    let predictor = TextPredictor::new(8, 20, 3);
    let result = predictor.predict_continuation(input, 20).unwrap();
    assert!(result.trajectory_equation.complexity() > 0);
    assert!(result.rhythm_equation.complexity() > 0);
}
