use crfs::train::Trainer;
use crfs::{Attribute, Model};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("CRF Training and Tagging Example");
    println!("=================================\n");

    // Create training data
    let xseq = vec![
        vec![Attribute::new("walk", 1.0), Attribute::new("shop", 0.5)],
        vec![Attribute::new("walk", 1.0)],
        vec![Attribute::new("walk", 1.0), Attribute::new("clean", 0.5)],
        vec![Attribute::new("shop", 0.5), Attribute::new("clean", 0.5)],
        vec![Attribute::new("walk", 0.5), Attribute::new("clean", 1.0)],
        vec![Attribute::new("clean", 1.0), Attribute::new("shop", 0.1)],
        vec![Attribute::new("walk", 1.0), Attribute::new("shop", 0.5)],
        vec![],
        vec![Attribute::new("clean", 1.0)],
    ];
    let yseq = vec![
        "sunny", "sunny", "sunny", "rainy", "rainy", "rainy", "sunny", "sunny", "rainy",
    ];

    println!("Training data:");
    println!("  Sequence length: {}", xseq.len());
    println!("  Labels: {:?}\n", yseq);

    // Create and configure trainer
    println!("Creating trainer...");
    let mut trainer = Trainer::lbfgs();
    trainer.verbose(true).append(&xseq, &yseq)?;

    // Set parameters
    println!("Setting parameters:");
    trainer.params_mut().set_c1(0.0)?;
    trainer.params_mut().set_c2(1.0)?;
    trainer.params_mut().set_max_iterations(100)?;
    println!("  L1 regularization (c1): {}", trainer.params().c1());
    println!("  L2 regularization (c2): {}", trainer.params().c2());
    println!("  Max iterations: {}\n", trainer.params().max_iterations());

    // Train
    let temp_file = tempfile::NamedTempFile::new()?;
    println!("Training model...\n");
    trainer.train(temp_file.path())?;

    println!("\n=================================");
    println!("Training completed successfully!");
    println!("=================================\n");

    // Load model and tag
    println!("Loading trained model...");
    let model_data = std::fs::read(temp_file.path())?;
    let model = Model::new(&model_data)?;

    println!("Creating tagger...");
    let tagger = model.tagger()?;

    // Test on training data
    println!("\nTesting on training data:");
    let test_seq = vec![
        vec![Attribute::new("walk", 1.0)],
        vec![Attribute::new("shop", 1.0)],
        vec![Attribute::new("clean", 1.0)],
    ];

    let result = tagger.tag(&test_seq)?;
    println!("  Input: walk -> shop -> clean");
    println!("  Predicted labels: {:?}", result);

    Ok(())
}
