//! What the vertex shader's move-chain walk costs, against how deep the chain
//! is. Every active widget owns a slot, so the depth a primitive resolves
//! through is its depth in the widget tree.
//!
//!     cargo test --release --test chain_cost -- --ignored --nocapture
//!
//! Timed on the GPU with timestamp queries rather than by the clock: wall time
//! here varied by 2x between runs of one unchanged binary. The pass is
//! submitted and waited on, so this is the GPU's cost and not the recording
//! loop's -- which is what `draw_cost.rs` measures instead.
//!
//! The instances are two pixels wide so that vertex work dominates; a chain
//! walk that does not show up against small quads will not show up against
//! anything.
//!
//! The instance is leaked deliberately, for the reason `draw_cost.rs` gives.

use iris::prelude::*;
use iris_core::{
    MaskIdx, MoveIdx, PrimitiveInst, RectPrimitive, UiData, UiRegion, UiRenderNode, UiRenderState,
    UiScalar, UiSpan,
};
use wgpu::{Color as GpuColor, *};

const SIZE: u32 = 1024;
const INSTANCES: usize = 200_000;
const FRAMES: u32 = 20;
/// Reported as the best of this many batches, since the mean moves by more
/// than the thing being measured.
const BATCHES: u32 = 8;

fn gpu() -> Option<(Device, Queue, f32)> {
    let all = Instance::new(InstanceDescriptor::new_without_display_handle());
    let instance = match pollster::block_on(all.request_adapter(&RequestAdapterOptions::default()))
    {
        Ok(_) => all,
        Err(_) => Instance::new(InstanceDescriptor {
            backends: Backends::GL,
            ..InstanceDescriptor::new_without_display_handle()
        }),
    };
    let instance: &'static Instance = Box::leak(Box::new(instance));
    let adapter =
        pollster::block_on(instance.request_adapter(&RequestAdapterOptions::default())).ok()?;
    if !adapter.features().contains(Features::TIMESTAMP_QUERY) {
        println!("no timestamp queries on {:?}", adapter.get_info().name);
        return None;
    }
    println!("adapter: {:?}", adapter.get_info().name);
    let (device, queue) = pollster::block_on(adapter.request_device(&DeviceDescriptor {
        required_features: Features::TIMESTAMP_QUERY,
        ..Default::default()
    }))
    .ok()?;
    let period = queue.get_timestamp_period();
    Some((device, queue, period))
}

fn config(format: TextureFormat) -> SurfaceConfiguration {
    SurfaceConfiguration {
        usage: TextureUsages::RENDER_ATTACHMENT,
        format,
        color_space: SurfaceColorSpace::Auto,
        width: SIZE,
        height: SIZE,
        present_mode: PresentMode::Fifo,
        desired_maximum_frame_latency: 2,
        alpha_mode: CompositeAlphaMode::Auto,
        view_formats: vec![],
    }
}

/// A chain `depth` slots long, and instances that all resolve through its end.
fn fill(ui: &mut UiData, render: &mut UiRenderState, depth: usize) {
    let kind = ui.primitives.kind::<RectPrimitive>();
    let id = ui.widgets.add_strong(Rect::new(UiColor::WHITE)).id();

    let mut slot = MoveIdx::NONE;
    for _ in 0..depth {
        slot = render.moves.push(slot, UiRegion::FULL);
    }

    let px = |v: f32| UiScalar { rel: 0.0, abs: v };
    for i in 0..INSTANCES {
        let x = (i % (SIZE as usize / 2)) as f32 * 2.0;
        let y = (i / (SIZE as usize / 2)) as f32;
        render.layers.write(
            0,
            PrimitiveInst {
                kind,
                id,
                primitive: RectPrimitive::color(UiColor::WHITE),
                region: UiRegion::new(
                    UiSpan::new(px(x), px(x + 2.0)),
                    UiSpan::new(px(y), px(y + 1.0)),
                ),
                mask_idx: MaskIdx::NONE,
                move_idx: slot,
            },
        );
    }
}

/// Nanoseconds the pass took on the GPU, best of `BATCHES`.
fn pass_cost(device: &Device, queue: &Queue, period: f32, depth: usize) -> f64 {
    let format = TextureFormat::Bgra8Unorm;
    let mut node = UiRenderNode::new(device, &config(format));
    let mut ui = UiData::default();
    let mut render = UiRenderState::new();
    fill(&mut ui, &mut render, depth);
    node.update(device, queue, &mut ui, &mut render);

    let target = device.create_texture(&TextureDescriptor {
        label: Some("chain cost"),
        size: Extent3d {
            width: SIZE,
            height: SIZE,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: TextureDimension::D2,
        format,
        usage: TextureUsages::RENDER_ATTACHMENT,
        view_formats: &[],
    });
    let view = target.create_view(&TextureViewDescriptor::default());

    let queries = device.create_query_set(&QuerySetDescriptor {
        label: Some("chain cost"),
        ty: QueryType::Timestamp,
        count: 2,
    });
    let resolved = device.create_buffer(&BufferDescriptor {
        label: Some("resolved"),
        size: 16,
        usage: BufferUsages::QUERY_RESOLVE | BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    let readback = device.create_buffer(&BufferDescriptor {
        label: Some("readback"),
        size: 16,
        usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });

    let frame = || {
        let mut encoder = device.create_command_encoder(&CommandEncoderDescriptor::default());
        {
            let pass = &mut encoder.begin_render_pass(&RenderPassDescriptor {
                label: None,
                color_attachments: &[Some(RenderPassColorAttachment {
                    view: &view,
                    resolve_target: None,
                    ops: Operations {
                        load: LoadOp::Clear(GpuColor::BLACK),
                        store: StoreOp::Store,
                    },
                    depth_slice: None,
                })],
                depth_stencil_attachment: None,
                timestamp_writes: Some(RenderPassTimestampWrites {
                    query_set: &queries,
                    beginning_of_pass_write_index: Some(0),
                    end_of_pass_write_index: Some(1),
                }),
                occlusion_query_set: None,
                multiview_mask: None,
            });
            node.draw(pass);
        }
        encoder.resolve_query_set(&queries, 0..2, &resolved, 0);
        encoder.copy_buffer_to_buffer(&resolved, 0, &readback, 0, 16);
        queue.submit(Some(encoder.finish()));

        let slice = readback.slice(..);
        slice.map_async(MapMode::Read, |_| {});
        let _ = device.poll(PollType::Wait {
            submission_index: None,
            timeout: None,
        });
        let ns = {
            let view = slice.get_mapped_range().expect("timestamps did not map");
            let stamps: [u64; 2] = [
                u64::from_le_bytes(view[..8].try_into().unwrap()),
                u64::from_le_bytes(view[8..16].try_into().unwrap()),
            ];
            (stamps[1].saturating_sub(stamps[0])) as f64 * period as f64
        };
        readback.unmap();
        ns
    };

    frame();
    let mut best = f64::MAX;
    for _ in 0..BATCHES {
        let mut total = 0.0;
        for _ in 0..FRAMES {
            total += frame();
        }
        best = best.min(total / FRAMES as f64);
    }
    best
}

#[test]
#[ignore = "measurement, not a check"]
fn chain_cost_by_depth() {
    let Some((device, queue, period)) = gpu() else {
        println!("no gpu with timestamps; nothing measured");
        return;
    };
    println!("{INSTANCES} instances, {SIZE}x{SIZE}, best of {BATCHES} batches");
    let mut base = None;
    for depth in [1, 2, 4, 8, 16, 32, 64] {
        let ns = pass_cost(&device, &queue, period, depth);
        let base = *base.get_or_insert(ns);
        println!(
            "depth {depth:>3}: {:>9.1} us   {:+6.1}% against depth 1",
            ns / 1000.0,
            (ns - base) / base * 100.0
        );
    }
}