//! Example demonstrating how to append frames to an existing TACO trajectory file
//!
//! This example shows:
//! - Creating an initial trajectory file
//! - Reading it back to verify
//! - Appending additional frames
//! - Verifying the complete trajectory
//!
//! Run with: cargo run --example append_trajectory

use ndarray::Array2;
use taco_format::{
    AtomMetadata, CompressionSettings, Frame, FrameData, Reader, SimulationMetadata, Writer,
};

fn main() -> taco_format::Result<()> {
    println!("=== TACO Append Trajectory Example ===\n");

    // Step 1: Create an initial trajectory file with some frames
    println!("Step 1: Creating initial trajectory file with 5 frames...");
    let file_path = "example_trajectory.taco";

    let sim_metadata = SimulationMetadata {
        name: Some("Water Molecule MD".to_string()),
        temperature: Some(300.0),
        pressure: Some(1.0),
        ensemble: Some("NVT".to_string()),
        ..Default::default()
    };

    let atom_metadata = AtomMetadata {
        names: Some(vec!["O".to_string(), "H".to_string(), "H".to_string()]),
        elements: Some(vec!["O".to_string(), "H".to_string(), "H".to_string()]),
        masses: Some(vec![15.999, 1.008, 1.008]),
        ..Default::default()
    };

    let compression_settings = CompressionSettings::default();

    // Create writer and write initial frames
    let mut writer = Writer::create(
        file_path,
        3, // 3 atoms (H2O)
        0.001,
        sim_metadata.clone(),
        atom_metadata.clone(),
        compression_settings.clone(),
    )?;

    // Simulate a simple trajectory where atoms move slightly each frame
    let mut positions = Array2::<f32>::zeros((3, 3));

    // Initial positions for water molecule
    positions[[0, 0]] = 0.0; // O atom at origin
    positions[[0, 1]] = 0.0;
    positions[[0, 2]] = 0.0;
    positions[[1, 0]] = 0.96; // H atom 1
    positions[[1, 1]] = 0.0;
    positions[[1, 2]] = 0.0;
    positions[[2, 0]] = -0.24; // H atom 2
    positions[[2, 1]] = 0.93;
    positions[[2, 2]] = 0.0;

    for i in 0..5 {
        // Add small displacements to simulate thermal motion
        for j in 0..3 {
            for k in 0..3 {
                positions[[j, k]] += 0.001 * (i as f32);
            }
        }

        let frame_data = FrameData::new(positions.clone());
        let frame = Frame::new(i, i as f64 * 0.001, frame_data)
            .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0])
            .with_thermo(300.0, 1.0);

        writer.write_frame(frame)?;
    }

    writer.finish()?;
    println!("✓ Initial trajectory created with 5 frames\n");

    // Step 2: Read back and verify initial trajectory
    println!("Step 2: Reading back initial trajectory...");
    {
        let mut reader = Reader::open(file_path)?;
        println!("  Number of frames: {}", reader.num_frames());
        println!("  Number of atoms: {}", reader.num_atoms());
        println!(
            "  Simulation: {:?}",
            reader.header().simulation_metadata.name
        );

        // Read a sample frame
        let frame = reader.read_frame(0)?;
        println!(
            "  Frame 0 position of O atom: [{:.3}, {:.3}, {:.3}]",
            frame.data.positions.as_ref().unwrap()[[0, 0]],
            frame.data.positions.as_ref().unwrap()[[0, 1]],
            frame.data.positions.as_ref().unwrap()[[0, 2]]
        );
    }
    println!("✓ Initial trajectory verified\n");

    // Step 3: Append additional frames
    println!("Step 3: Appending 3 more frames to the trajectory...");
    {
        let mut writer = Writer::append(file_path)?;

        for i in 5..8 {
            // Continue the trajectory
            for j in 0..3 {
                for k in 0..3 {
                    positions[[j, k]] += 0.001 * (i as f32);
                }
            }

            let frame_data = FrameData::new(positions.clone());
            let frame = Frame::new(i, i as f64 * 0.001, frame_data)
                .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0])
                .with_thermo(300.0, 1.0);

            writer.write_frame(frame)?;
        }

        writer.finish()?;
    }
    println!("✓ Successfully appended 3 frames\n");

    // Step 4: Verify the complete trajectory
    println!("Step 4: Verifying complete trajectory...");
    {
        let mut reader = Reader::open(file_path)?;
        println!("  Total number of frames: {}", reader.num_frames());

        assert_eq!(
            reader.num_frames(),
            8,
            "Expected 8 frames (5 initial + 3 appended)"
        );

        // Read and display info about frames across the boundary
        for i in [4, 5, 6] {
            let frame = reader.read_frame(i)?;
            println!(
                "  Frame {} time: {:.4} ps, is_full_frame: {}",
                i, frame.time, frame.is_full_frame
            );
        }
    }
    println!("✓ Complete trajectory verified\n");

    // Step 5: Demonstrate multiple appends
    println!("Step 5: Demonstrating multiple sequential appends...");
    for append_round in 1..=3 {
        let mut writer = Writer::append(file_path)?;
        let start_frame = 8 + (append_round - 1) * 2;

        for i in start_frame..(start_frame + 2) {
            for j in 0..3 {
                for k in 0..3 {
                    positions[[j, k]] += 0.001 * (i as f32);
                }
            }

            let frame_data = FrameData::new(positions.clone());
            let frame = Frame::new(i, i as f64 * 0.001, frame_data)
                .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0])
                .with_thermo(300.0, 1.0);

            writer.write_frame(frame)?;
        }

        writer.finish()?;
        println!("  Append round {}: Added 2 frames", append_round);
    }

    let reader = Reader::open(file_path)?;
    println!(
        "✓ After {} append operations, total frames: {}\n",
        3,
        reader.num_frames()
    );

    assert_eq!(
        reader.num_frames(),
        14,
        "Expected 14 frames total (8 + 3*2)"
    );

    // Step 6: Show random access still works
    println!("Step 6: Demonstrating random access across appended frames...");
    {
        let mut reader = Reader::open(file_path)?;

        // Access frames in non-sequential order
        for &frame_idx in &[0, 7, 3, 13, 5, 10] {
            let frame = reader.read_frame(frame_idx)?;
            println!(
                "  Random access frame {}: time={:.4} ps, full_frame={}",
                frame_idx, frame.time, frame.is_full_frame
            );
        }
    }
    println!("✓ Random access working correctly\n");

    // Cleanup
    std::fs::remove_file(file_path)?;
    println!("=== Example completed successfully! ===");
    println!("\nKey takeaways:");
    println!("  • Append operation maintains delta encoding chain");
    println!("  • Multiple appends can be performed sequentially");
    println!("  • All metadata is preserved across appends");
    println!("  • Random access works seamlessly across original and appended frames");
    println!("  • Full frame intervals are respected based on total frame count");

    Ok(())
}
