// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! Simple linear elasticity example.
//!
//! This example demonstrates basic usage of the RTX-FEA crate
//! for a simple 2D plane stress problem.

use rtx_fea::prelude::*;
use rtx_fea::{
    analysis::{AnalysisConfig, StaticLinearAnalysis},
    assembly::DofComponent,
    boundary::{BoundaryConditionSet, DirichletBC, NeumannBC},
    materials::{LinearElastic, MaterialDatabase},
    mesh::MaterialId,
    utils::StringUtils,
};

fn main() -> FeaResult<()> {
    println!("RTX-FEA Linear Elasticity Example");
    println!("=================================");

    // Create a simple rectangular domain
    let width = 2.0; // meters
    let height = 1.0; // meters
    let mesh = Mesh::generate_rectangle(width, height, 9, 5)?;

    println!(
        "Mesh: {} nodes, {} elements",
        mesh.num_nodes(),
        mesh.num_elements()
    );

    // Set up material (aluminum)
    let mut materials = MaterialDatabase::new();
    let aluminum = LinearElastic::new(70e9, 0.33); // E = 70 GPa, ν = 0.33
    materials.add_material(MaterialId(0), aluminum, Some("Aluminum".to_string()));

    // Set up boundary conditions
    let mut boundary_conditions = BoundaryConditionSet::new();

    // Fix left edge (x = 0)
    let mut fixed_nodes = Vec::new();
    for (&node_id, node) in &mesh.nodes {
        if node.position().x < 1e-6 {
            fixed_nodes.push(node_id);
        }
    }

    let fixed_bc = DirichletBC::fixed_support(fixed_nodes);
    boundary_conditions.add_condition(rtx_fea::boundary::BoundaryCondition::Dirichlet(fixed_bc));

    // Apply tension at right edge
    let mut loaded_nodes = Vec::new();
    for (&node_id, node) in &mesh.nodes {
        if (node.position().x - width).abs() < 1e-6 {
            loaded_nodes.push(node_id);
        }
    }

    let tension_bc = NeumannBC::fixed_force(
        loaded_nodes,
        vec![DofComponent::DisplacementX],
        1000.0, // 1000 N tension
    );
    boundary_conditions.add_condition(rtx_fea::boundary::BoundaryCondition::Neumann(tension_bc));

    // Configure and run analysis
    let config = AnalysisConfig::default();
    let mut analysis = StaticLinearAnalysis::new(mesh, materials, boundary_conditions, config);

    println!("Running analysis...");
    let results = analysis.run()?;

    println!("Analysis complete!");
    println!("Max displacement: {:.6} m", results.max_displacement());
    println!("Solver converged: {}", results.convergence.converged);
    println!(
        "Total time: {}",
        StringUtils::format_duration(results.timing.total_time)
    );

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_linear_elasticity_example() {
        let result = main();
        assert!(
            result.is_ok(),
            "Linear elasticity example should run without errors"
        );
    }
}