use burn::backend::wgpu::{CubeBackend, WgpuRuntime};
use burn::prelude::*;
use burn::tensor::Distribution;
use burn_fusion::Fusion;
use burn_tracing_backend::{Profiler, start_tracing, finish_tracing, write_trace, marker};

// Profiler sits inside the Fusion layer to intercept fusion dispatch
type B = Fusion<Profiler<CubeBackend<WgpuRuntime, f32, i32, u32>>>;

fn main() {
    let device = Default::default();

    // Start recording
    start_tracing();

    // Your operations are traced automatically
    let a: Tensor<B, 2> = Tensor::random([8, 8], Distribution::Uniform(0.0, 1.0), &device);
    let b: Tensor<B, 2> = Tensor::random([8, 8], Distribution::Uniform(0.0, 1.0), &device);
    let c = a.clone().add(b.clone()).mul(a.clone()).exp();

    // Explicit GPU sync — shows as a sync point on the timeline
    B::sync(&device).unwrap();

    // Implicit GPU sync — into_data() forces flush, also shown on the timeline
    let _data = c.to_data();

    // Collect events and write HTML + JS visualization
    let events = finish_tracing();
    write_trace(&events, "trace_data.js").unwrap();
    // Open trace_explorer.html in a browser
}

use burn_tracing_backend::marker;

// Instant marker — appears as a vertical line on the timeline
marker("checkpoint").debug("Epoch 5 complete, loss=0.032").emit();

// Timed span — measures duration from creation to drop
{
    let _span = marker("forward_pass").debug("Layer 1 → 2 → output").span();
    // ... operations run here ...
} // span duration recorded on drop

// Attach tensor data to markers
marker("output")
    .tensor::<B>(&result.clone().into_primitive().tensor())
    .debug(format!("shape: {:?}", result.dims()))
    .emit();

// Add tensor data to spans after they start
let mut span = marker("inference").span();
// ... compute ...
span.add_tensor::<B>(&output.clone().into_primitive().tensor());
span.set_debug("batch complete");
drop(span);

// Always-on tensor capture (works without `trace-data` feature)
// Second arg: Some(n) = first n values, None = all values
marker("weights")
    .tensor_with_data::<B>(&w.clone().into_primitive().tensor(), Some(128))
    .emit();

use burn_tracing_backend::set_data_capture_limit;

// Capture first 128 values per tensor (default: 64, 0 = disabled)
set_data_capture_limit(128);

use burn::backend::wgpu::{CubeBackend, WgpuRuntime};
use burn::prelude::*;
use burn::tensor::Distribution;
use burn_fusion::Fusion;
use burn_tracing_backend::{Profiler, start_tracing, finish_tracing, write_trace, marker, OpCategory};

type B = Fusion<Profiler<CubeBackend<WgpuRuntime, f32, i32, u32>>>;

fn main() {
    let device = Default::default();
    start_tracing();

    marker("begin").debug("Fusion tracing demo").emit();

    // --- Elementwise fusion chain ---
    {
        let _span = marker("elementwise").debug("add → mul → exp chain").span();
        let a: Tensor<B, 2> = Tensor::random([4, 16], Distribution::Normal(0.0, 1.0), &device);
        let b: Tensor<B, 2> = Tensor::random([4, 16], Distribution::Normal(0.0, 1.0), &device);
        let _c = a.clone().add(b).mul(a).exp().log().abs();
        B::sync(&device).unwrap(); // explicit sync
    }

    // --- Matmul fusion ---
    {
        let _span = marker("matmul").debug("linear layer: matmul + bias + relu").span();
        let x: Tensor<B, 2> = Tensor::random([4, 16], Distribution::Normal(0.0, 1.0), &device);
        let w: Tensor<B, 2> = Tensor::random([16, 8], Distribution::Normal(0.0, 0.1), &device);
        let bias: Tensor<B, 1> = Tensor::zeros([8], &device);
        let _out = x.matmul(w).add(bias.unsqueeze()).clamp_min(0.0);
        B::sync(&device).unwrap();
    }

    // --- Reductions ---
    {
        let _span = marker("reductions").span();
        let t: Tensor<B, 2> = Tensor::random([8, 8], Distribution::Uniform(0.0, 1.0), &device);
        let _mean = t.clone().mean();
        let _sum = t.clone().sum();
        B::sync(&device).unwrap();
    }

    // --- Implicit sync via into_data ---
    let result: Tensor<B, 2> = Tensor::ones([2, 2], &device);
    let _data = result.to_data(); // implicit GPU sync shown on timeline

    marker("end").emit();

    let events = finish_tracing();

    // Print summary
    let total_us: f64 = events.iter().map(|e| e.duration_us).sum();
    let fusions = events.iter().filter(|e| e.category == OpCategory::Fusion).count();
    let syncs = events.iter().filter(|e| e.is_sync == Some(true)).count();
    println!("Events: {}, Fusions: {}, Syncs: {}, Total: {:.0}us", events.len(), fusions, syncs, total_us);

    for e in &events {
        if let Some(kind) = &e.fusion_kind {
            println!("  {:?} — {} ops fused, {:.1}us", kind, e.num_fused_ops.unwrap_or(0), e.duration_us);
        }
    }

    // Write visualization
    write_trace(&events, "trace_data.js").unwrap();
    println!("Open trace_explorer.html in a browser to explore the trace.");
}