//! TRACES: FR-MRG-10 | FR-MRG-11
//! The merge: source frames warped into an output surface, chunk by chunk.
//!
//! The per-pixel half of a panorama (FR-MRG-10), on the GPU: the warp of a
//! source tile into an output chunk, the weighted accumulation across
//! frames, and the resolve to sixteen-bit samples. The geometry it is
//! given — rotations, focal length, projection — is `dr-pano`'s, solved on
//! proxies before any full-resolution pixel exists (panorama.md §5), and
//! that is what makes this simple: every output pixel's source coordinates
//! are a closed-form function, so a chunk can be produced from the source
//! tiles that project into it and nothing else.
//!
//! # The loop
//!
//! ```text
//! for each band of rows of the output:
//!     for each chunk across the band:
//!         zero the accumulator
//!         for each frame whose footprint meets the chunk:
//!             the source rectangle the chunk needs, from the geometry
//!             render it camera-linear through the pipeline (the tile)
//!             warp the tile into the chunk, accumulate       ← GPU
//!         resolve the chunk to u16                            ← GPU
//!         copy it into the band
//!     hand the band to the writer (one DNG strip)
//! ```
//!
//! No stage holds the composite (FR-MRG-11): the working set is one
//! chunk's accumulator, one tile, one band of u16 rows. The frame textures
//! are the caller's to provide and cache — `source` is asked for frame `k`
//! as it is needed, and a caller short of memory may demosaic on demand.
//!
//! # The blend
//!
//! With a seam map (`dr_pano::seam`), a frame's weight at a pixel is its
//! share of the map about that pixel — whole on its own side of a seam,
//! nothing on the other, and a ramp across a window `seam_blend` pixels
//! wide that follows the seam. Without one, or where the map has nothing
//! to say, the weight is the distance to the frame's edge over `feather`,
//! which hides exposure steps and does not hide parallax: the average draws
//! anything the frames disagree on twice.
//!
//! # What is not here yet
//!
//! A Laplacian pyramid, which would let the seam's blend be narrow for
//! detail and wide for exposure at once. Gain is a scalar per frame the
//! caller supplies.

use std::sync::Arc;

use dr_pano::bundle::Cameras;
use dr_pano::projection::{Bounds, Projection};
use dr_pano::seam::SeamMap;
use wgpu::util::DeviceExt;

use crate::readback::await_mapping;
use crate::{AdjustPass, DemosaicedImage, GpuContext, GpuError};

/// One frame's part in the merge.
pub struct MergeFrame {
    /// The frame's edit, for its lens corrections — the only part of an
    /// edit the camera-space tap uses (FR-MRG-2).
    pub graph: Arc<dr_pipeline::EditGraph>,
    /// Multiplies the frame's samples, to bring its exposure to the
    /// reference frame's. 1.0 for no correction.
    pub gain: f32,
}

/// The output the merge produces.
#[derive(Debug, Clone, PartialEq)]
pub struct MergeOutput {
    pub projection: Projection,
    /// The projection's scale in output pixels: the cylinder's radius, the
    /// plane's distance. The source focal length at full resolution gives
    /// output pixels the size of source pixels at the centre.
    pub scale: f64,
    /// The rectangle of the projection to produce, centred coordinates.
    pub bounds: Bounds,
    /// Pixels over which a frame's weight ramps up from its edge.
    pub feather: f32,
    /// Which frame each part of the output is taken from, laid out at the
    /// proxies' scale; `None` for the feathered average everywhere.
    pub seams: Option<Arc<SeamMap>>,
    /// The width, in output pixels, of the blend across a seam.
    pub seam_blend: f32,
    /// Chunk size: the unit of GPU work and of memory.
    pub chunk: (u32, u32),
    /// Multiplies a normalised sample (1.0 = white) to the sensor's scale.
    pub sample_scale: f32,
    /// The white balance the composite will be developed with — the
    /// inverse of its `AsShotNeutral` — so that a blown sample can be
    /// written as the camera value that balance calls grey.
    pub balance: [f32; 3],
}

impl MergeOutput {
    pub fn width(&self) -> u32 {
        self.bounds.width().ceil().max(1.0) as u32
    }
    pub fn height(&self) -> u32 {
        self.bounds.height().ceil().max(1.0) as u32
    }
}

/// A band of finished rows: `rows × width × 3` RGB `u16`, plus a coverage
/// mask (`true` where any frame reached the pixel).
pub struct Band<'a> {
    pub first_row: u32,
    pub rows: u32,
    pub rgb: &'a [u16],
    pub covered: &'a [bool],
}

#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct WarpParams {
    chunk_origin: [f32; 2],
    chunk_size: [u32; 2],
    projection: u32,
    proj_scale: f32,
    focal: f32,
    gain: f32,
    r0: [f32; 4],
    r1: [f32; 4],
    r2: [f32; 4],
    frame_size: [f32; 2],
    tile_origin: [f32; 2],
    tile_size: [u32; 2],
    feather: f32,
    clip_onset: f32,
    balance: [f32; 4],
    seam_origin: [f32; 2],
    seam_size: [u32; 2],
    seam_px: f32,
    seam_radius: f32,
    frame_index: u32,
    seam_on: u32,
}

#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct ResolveParams {
    chunk_size: [u32; 2],
    scale: f32,
    _pad: f32,
}

/// The two pipelines and the chunk buffers.
pub struct MergePass {
    ctx: GpuContext,
    warp: wgpu::ComputePipeline,
    warp_layout: wgpu::BindGroupLayout,
    resolve: wgpu::ComputePipeline,
    resolve_layout: wgpu::BindGroupLayout,
    /// Accumulator and packed output for the current chunk size.
    buffers: Option<(wgpu::Buffer, wgpu::Buffer, wgpu::Buffer, (u32, u32))>,
}

impl MergePass {
    pub fn new(ctx: &GpuContext) -> Result<Self, GpuError> {
        let module = ctx
            .device
            .create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("merge"),
                source: wgpu::ShaderSource::Wgsl(include_str!("shaders/merge.wgsl").into()),
            });
        let uniform = |binding| wgpu::BindGroupLayoutEntry {
            binding,
            visibility: wgpu::ShaderStages::COMPUTE,
            ty: wgpu::BindingType::Buffer {
                ty: wgpu::BufferBindingType::Uniform,
                has_dynamic_offset: false,
                min_binding_size: None,
            },
            count: None,
        };
        let storage = |binding, read_only| wgpu::BindGroupLayoutEntry {
            binding,
            visibility: wgpu::ShaderStages::COMPUTE,
            ty: wgpu::BindingType::Buffer {
                ty: wgpu::BufferBindingType::Storage { read_only },
                has_dynamic_offset: false,
                min_binding_size: None,
            },
            count: None,
        };
        let warp_layout = ctx
            .device
            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("merge-warp-bgl"),
                entries: &[
                    uniform(0),
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Texture {
                            // Unfilterable: rgba32float, loaded by hand.
                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
                            view_dimension: wgpu::TextureViewDimension::D2,
                            multisampled: false,
                        },
                        count: None,
                    },
                    storage(2, false),
                    wgpu::BindGroupLayoutEntry {
                        binding: 3,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Texture {
                            sample_type: wgpu::TextureSampleType::Uint,
                            view_dimension: wgpu::TextureViewDimension::D2,
                            multisampled: false,
                        },
                        count: None,
                    },
                ],
            });
        let resolve_layout =
            ctx.device
                .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                    label: Some("merge-resolve-bgl"),
                    entries: &[uniform(0), storage(1, true), storage(2, false)],
                });
        let pipeline = |name: &str, layout: &wgpu::BindGroupLayout| {
            let pl = ctx
                .device
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                    label: Some(name),
                    bind_group_layouts: &[Some(layout)],
                    immediate_size: 0,
                });
            ctx.device
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                    label: Some(name),
                    layout: Some(&pl),
                    module: &module,
                    entry_point: Some(name),
                    compilation_options: Default::default(),
                    cache: None,
                })
        };
        Ok(MergePass {
            ctx: ctx.clone(),
            warp: pipeline("warp", &warp_layout),
            warp_layout,
            resolve: pipeline("resolve", &resolve_layout),
            resolve_layout,
            buffers: None,
        })
    }

    /// Allocate the chunk buffers for this size if the last ones differ.
    fn ensure_buffers(&mut self, chunk: (u32, u32)) {
        if self.buffers.as_ref().is_none_or(|b| b.3 != chunk) {
            let n = u64::from(chunk.0) * u64::from(chunk.1);
            let acc = self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("merge-acc"),
                size: n * 16,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            let out = self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("merge-out"),
                size: n * 8,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            let read = self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("merge-read"),
                size: n * 8,
                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
                mapped_at_creation: false,
            });
            self.buffers = Some((acc, out, read, chunk));
        }
    }

    fn chunk_buffers(&self) -> (&wgpu::Buffer, &wgpu::Buffer, &wgpu::Buffer) {
        let b = self.buffers.as_ref().expect("ensured by the caller");
        (&b.0, &b.1, &b.2)
    }

    /// Produce the whole output, band by band, handing each finished band
    /// to `sink`.
    ///
    /// `cameras` are in **full-resolution source pixels** (`frame_size`),
    /// with frame `k` corresponding to `frames[k]` and `source(k)`. `source`
    /// supplies the demosaiced frame on demand and may cache as it sees fit.
    #[allow(clippy::too_many_arguments)]
    pub fn merge<S, F>(
        &mut self,
        adjust: &mut AdjustPass,
        frames: &[MergeFrame],
        cameras: &Cameras,
        frame_size: (u32, u32),
        output: &MergeOutput,
        mut source: S,
        mut sink: F,
        mut cancelled: impl FnMut() -> bool,
    ) -> Result<(), GpuError>
    where
        S: FnMut(usize) -> Result<Arc<DemosaicedImage>, GpuError>,
        F: FnMut(Band<'_>) -> Result<(), GpuError>,
    {
        let (out_w, out_h) = (output.width(), output.height());
        let (cw, ch) = (output.chunk.0.max(8), output.chunk.1.max(8));
        let (fw, fh) = (frame_size.0 as f64, frame_size.1 as f64);

        let mut band_rgb = vec![0u16; (out_w * ch * 3) as usize];
        let mut band_cov = vec![false; (out_w * ch) as usize];
        let mut chunk_px: Vec<u32> = Vec::new();

        // The seam map, once for the whole output, and where it sits in
        // this output's coordinates. A one-texel stand-in when there is
        // none, because the binding is not optional.
        let (seam_tex, seam_origin, seam_px, seam_radius, seam_size) = match &output.seams {
            Some(m) => {
                let ((ou, ov), px) = m.at_scale(output.scale);
                let radius = m.blend_radius(output.scale, f64::from(output.seam_blend));
                (
                    self.label_texture(m.width as u32, m.height as u32, &m.labels),
                    [ou as f32, ov as f32],
                    px as f32,
                    radius as f32,
                    [m.width as u32, m.height as u32],
                )
            }
            None => (
                self.label_texture(1, 1, &[dr_pano::seam::NONE]),
                [0.0; 2],
                1.0,
                1.0,
                [1, 1],
            ),
        };
        let seam_view = seam_tex.create_view(&Default::default());
        let seam_on = u32::from(output.seams.is_some());

        let mut y = 0u32;
        while y < out_h {
            let rows = ch.min(out_h - y);
            band_rgb.iter_mut().for_each(|v| *v = 0);
            band_cov.iter_mut().for_each(|v| *v = false);

            let mut x = 0u32;
            while x < out_w {
                if cancelled() {
                    return Err(GpuError::Readback("merge cancelled".into()));
                }
                let cols = cw.min(out_w - x);
                let origin = (
                    output.bounds.min_u + f64::from(x),
                    output.bounds.min_v + f64::from(y),
                );
                self.zero_accumulator((cols, rows));

                for (k, frame) in frames.iter().enumerate() {
                    let Some(rect) = source_rect(
                        output.projection,
                        output.scale,
                        cameras,
                        k,
                        origin,
                        (cols, rows),
                        (fw, fh),
                    ) else {
                        continue;
                    };
                    let image = source(k)?;
                    // The tile: that rectangle of the frame, camera-linear,
                    // at 1:1.
                    let view = dr_pipeline::CropRect {
                        x: (rect.0 as f32) / fw as f32,
                        y: (rect.1 as f32) / fh as f32,
                        width: (rect.2 as f32) / fw as f32,
                        height: (rect.3 as f32) / fh as f32,
                    };
                    let shader = frame.graph.compose_camera_linear(view);
                    let tile = adjust.render_camera_linear(&image, &shader, rect.2, rect.3)?;
                    let r = cameras.rotations[k].transpose();
                    let params = WarpParams {
                        chunk_origin: [origin.0 as f32, origin.1 as f32],
                        chunk_size: [cols, rows],
                        projection: match output.projection {
                            Projection::Perspective => 0,
                            Projection::Cylindrical => 1,
                            Projection::Spherical => 2,
                        },
                        proj_scale: output.scale as f32,
                        focal: cameras.focal as f32,
                        gain: frame.gain,
                        r0: [r.0[0][0] as f32, r.0[0][1] as f32, r.0[0][2] as f32, 0.0],
                        r1: [r.0[1][0] as f32, r.0[1][1] as f32, r.0[1][2] as f32, 0.0],
                        r2: [r.0[2][0] as f32, r.0[2][1] as f32, r.0[2][2] as f32, 0.0],
                        frame_size: [fw as f32, fh as f32],
                        tile_origin: [rect.0 as f32, rect.1 as f32],
                        tile_size: [rect.2, rect.3],
                        feather: output.feather,
                        clip_onset: dr_pipeline::CLIP_ONSET,
                        balance: [
                            output.balance[0].max(1e-3),
                            output.balance[1].max(1e-3),
                            output.balance[2].max(1e-3),
                            0.0,
                        ],
                        seam_origin,
                        seam_size,
                        seam_px,
                        seam_radius,
                        frame_index: k as u32,
                        seam_on,
                    };
                    self.accumulate(&params, tile, &seam_view);
                }

                self.resolve_chunk((cols, rows), output.sample_scale, &mut chunk_px)?;
                // Into the band.
                for row in 0..rows as usize {
                    for col in 0..cols as usize {
                        let px = chunk_px[(row * cols as usize + col) * 2..][..2].to_vec();
                        let i = row * out_w as usize + (x as usize + col);
                        band_rgb[i * 3] = (px[0] & 0xFFFF) as u16;
                        band_rgb[i * 3 + 1] = (px[0] >> 16) as u16;
                        band_rgb[i * 3 + 2] = (px[1] & 0xFFFF) as u16;
                        band_cov[i] = (px[1] >> 16) != 0;
                    }
                }
                x += cols;
            }

            sink(Band {
                first_row: y,
                rows,
                rgb: &band_rgb[..(out_w * rows * 3) as usize],
                covered: &band_cov[..(out_w * rows) as usize],
            })?;
            y += rows;
        }
        Ok(())
    }

    fn zero_accumulator(&mut self, chunk: (u32, u32)) {
        self.ensure_buffers(chunk);
        let (acc, _, _) = self.chunk_buffers();
        let n = u64::from(chunk.0) * u64::from(chunk.1) * 16;
        let mut enc = self.ctx.device.create_command_encoder(&Default::default());
        enc.clear_buffer(acc, 0, Some(n));
        self.ctx.queue.submit(Some(enc.finish()));
    }

    /// The seam map's labels as an `r8uint` texture.
    fn label_texture(&self, width: u32, height: u32, labels: &[u8]) -> wgpu::Texture {
        let size = wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        };
        let tex = self.ctx.device.create_texture(&wgpu::TextureDescriptor {
            label: Some("merge-seams"),
            size,
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::R8Uint,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });
        self.ctx.queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture: &tex,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            labels,
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(width),
                rows_per_image: Some(height),
            },
            size,
        );
        tex
    }

    fn accumulate(&mut self, params: &WarpParams, tile: &wgpu::Texture, seams: &wgpu::TextureView) {
        let chunk = (params.chunk_size[0], params.chunk_size[1]);
        let uniforms = self
            .ctx
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("merge-warp-params"),
                contents: bytemuck::bytes_of(params),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let view = tile.create_view(&Default::default());
        self.ensure_buffers(chunk);
        let (acc, _, _) = self.chunk_buffers();
        let bind = self
            .ctx
            .device
            .create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("merge-warp-bg"),
                layout: &self.warp_layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: uniforms.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: wgpu::BindingResource::TextureView(&view),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: acc.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 3,
                        resource: wgpu::BindingResource::TextureView(seams),
                    },
                ],
            });
        let mut enc = self.ctx.device.create_command_encoder(&Default::default());
        {
            let mut pass = enc.begin_compute_pass(&Default::default());
            pass.set_pipeline(&self.warp);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(chunk.0.div_ceil(8), chunk.1.div_ceil(8), 1);
        }
        self.ctx.queue.submit(Some(enc.finish()));
    }

    fn resolve_chunk(
        &mut self,
        chunk: (u32, u32),
        scale: f32,
        out: &mut Vec<u32>,
    ) -> Result<(), GpuError> {
        let params = ResolveParams {
            chunk_size: [chunk.0, chunk.1],
            scale,
            _pad: 0.0,
        };
        let uniforms = self
            .ctx
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("merge-resolve-params"),
                contents: bytemuck::bytes_of(&params),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let n = u64::from(chunk.0) * u64::from(chunk.1);
        self.ensure_buffers(chunk);
        let (acc, packed, read) = self.chunk_buffers();
        let bind = self
            .ctx
            .device
            .create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("merge-resolve-bg"),
                layout: &self.resolve_layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: uniforms.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: acc.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: packed.as_entire_binding(),
                    },
                ],
            });
        let mut enc = self.ctx.device.create_command_encoder(&Default::default());
        {
            let mut pass = enc.begin_compute_pass(&Default::default());
            pass.set_pipeline(&self.resolve);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(chunk.0.div_ceil(8), chunk.1.div_ceil(8), 1);
        }
        enc.copy_buffer_to_buffer(packed, 0, read, 0, n * 8);
        self.ctx.queue.submit(Some(enc.finish()));

        let slice = read.slice(..n * 8);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |r| {
            let _ = tx.send(r);
        });
        await_mapping(&self.ctx, &rx)?;
        {
            let data = slice.get_mapped_range();
            out.clear();
            out.extend_from_slice(bytemuck::cast_slice::<u8, u32>(&data));
        }
        read.unmap();
        Ok(())
    }
}

/// The rectangle of frame `k` (x, y, w, h in source pixels) a chunk reads,
/// or `None` if the chunk sees nothing of the frame.
///
/// Walks the chunk's border, projects each point into the frame, and takes
/// the bounding box with a two-pixel margin for the bilinear fetch. The
/// border rather than the corners because under a cylinder or sphere the
/// extreme of a footprint is not at a corner.
fn source_rect(
    projection: Projection,
    scale: f64,
    cameras: &Cameras,
    k: usize,
    origin: (f64, f64),
    size: (u32, u32),
    frame: (f64, f64),
) -> Option<(u32, u32, u32, u32)> {
    let (w, h) = (f64::from(size.0), f64::from(size.1));
    let steps = 16;
    let mut min = (f64::MAX, f64::MAX);
    let mut max = (f64::MIN, f64::MIN);
    let mut any = false;
    let mut visit = |u: f64, v: f64| {
        let d = projection.to_direction(scale, u, v);
        if let Some((x, y)) = cameras.project(k, d) {
            let (x, y) = (x + frame.0 / 2.0, y + frame.1 / 2.0);
            min = (min.0.min(x), min.1.min(y));
            max = (max.0.max(x), max.1.max(y));
            any = true;
        }
    };
    for s in 0..=steps {
        let t = f64::from(s) / f64::from(steps);
        visit(origin.0 + w * t, origin.1);
        visit(origin.0 + w * t, origin.1 + h);
        visit(origin.0, origin.1 + h * t);
        visit(origin.0 + w, origin.1 + h * t);
    }
    // The interior too, coarsely: a chunk can contain a frame entirely.
    for i in 1..4 {
        for j in 1..4 {
            visit(
                origin.0 + w * f64::from(i) / 4.0,
                origin.1 + h * f64::from(j) / 4.0,
            );
        }
    }
    if !any {
        return None;
    }
    let x0 = (min.0.floor() - 2.0).max(0.0);
    let y0 = (min.1.floor() - 2.0).max(0.0);
    let x1 = (max.0.ceil() + 2.0).min(frame.0);
    let y1 = (max.1.ceil() + 2.0).min(frame.1);
    if x1 <= x0 || y1 <= y0 {
        return None;
    }
    Some((x0 as u32, y0 as u32, (x1 - x0) as u32, (y1 - y0) as u32))
}