use burn_central::{
    experiment::ExperimentRun,
    macros::register,
    runtime::{Args, ArtifactLoader, Model, MultiDevice},
};
use burn::prelude::*;

#[register(training, name = "mnist")]
pub fn training<B: AutodiffBackend>(
    experiment: &ExperimentRun,
    config: Args<YourExperimentConfig>,
    MultiDevice(devices): MultiDevice<B>,
    loader: ArtifactLoader<ModelArtifact<B>>,
) -> Result<Model<ModelArtifact<B::InnerBackend>>, String> {
    // Log your configuration
    experiment.log_config("Training Config", &training_config)
        .expect("Logging config failed");

    // Your training logic here...
    let model = train::<B>(experiment, artifact_dir, &training_config, devices[0].clone())?;

    Ok(Model(ModelArtifact {
        model_record: model.into_record(),
        config: training_config,
    }))
}

use burn_central::experiment::integration::training::{
    ExperimentCheckpointRecorder,
    ExperimentMetricLogger,
    experiment_interrupter,
};
use burn_central::experiment::integration::tracing::try_init_tracing_subscriber;
use burn::train::{LearnerBuilder, metric::{AccuracyMetric, LossMetric}};

let _ = try_init_tracing_subscriber();

let learner = LearnerBuilder::new(artifact_dir)
    .metric_train_numeric(AccuracyMetric::new())
    .metric_valid_numeric(AccuracyMetric::new())
    .metric_train_numeric(LossMetric::new())
    .metric_valid_numeric(LossMetric::new())
    // Experiment metric logging
    .with_metric_logger(ExperimentMetricLogger::new(experiment))
    // Experiment checkpoint saving
    .with_file_checkpointer(ExperimentCheckpointRecorder::new(experiment))
    // Experiment interruption handling
    .with_interrupter(experiment_interrupter(experiment))
    .num_epochs(config.num_epochs)
    .summary()
    .build(
        model.init::<B>(&device),
        optimizer.init(),
        learning_rate,
        LearningStrategy::SingleDevice(device),
    );