path: root/src/main.rs

use winit::{
    event::*,
    event_loop::EventLoop,
    keyboard::{KeyCode, PhysicalKey},
};
use wgpu::util::DeviceExt;

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct Vertex {
    position: [f32; 2],
    color: [f32; 3],
}

impl Vertex {
    fn desc() -> wgpu::VertexBufferLayout<'static> {
        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &[
                wgpu::VertexAttribute {
                    offset: 0,
                    shader_location: 0,
                    format: wgpu::VertexFormat::Float32x2,
                },
                wgpu::VertexAttribute {
                    offset: std::mem::size_of::<[f32; 2]>() as wgpu::BufferAddress,
                    shader_location: 1,
                    format: wgpu::VertexFormat::Float32x3,
                },
            ],
        }
    }
}

#[derive(Clone)]
struct Viewport {
    x: f32,
    y: f32,
    width: f32,
    height: f32,
}

struct GraphView {
    viewport: Viewport,
    lines: Vec<Vec<Vertex>>,
    show_grid: bool,
}

impl GraphView {
    fn new(viewport: Viewport) -> Self {
        Self {
            viewport,
            lines: Vec::new(),
            show_grid: true,
        }
    }

    fn generate_grid_lines(&self) -> Vec<Vertex> {
        let mut vertices = Vec::new();
        let grid_color = [0.3, 0.7, 0.9];

        // Vertical grid lines (10 divisions)
        for i in 0..=10 {
            let x = -1.0 + (i as f32 / 10.0) * 2.0;
            vertices.push(Vertex { position: [x, -1.0], color: grid_color });
            vertices.push(Vertex { position: [x, 1.0], color: grid_color });
        }

        // Horizontal grid lines (10 divisions)
        for i in 0..=10 {
            let y = -1.0 + (i as f32 / 10.0) * 2.0;
            vertices.push(Vertex { position: [-1.0, y], color: grid_color });
            vertices.push(Vertex { position: [1.0, y], color: grid_color });
        }

        vertices
    }

    fn generate_border(&self) -> Vec<Vertex> {
        let border_color = [0.6, 0.7, 0.7];
        vec![
            // Top border
            Vertex { position: [-1.0, 1.0], color: border_color },
            Vertex { position: [1.0, 1.0], color: border_color },
            // Right border
            Vertex { position: [1.0, 1.0], color: border_color },
            Vertex { position: [1.0, -1.0], color: border_color },
            // Bottom border
            Vertex { position: [1.0, -1.0], color: border_color },
            Vertex { position: [-1.0, -1.0], color: border_color },
            // Left border
            Vertex { position: [-1.0, -1.0], color: border_color },
            Vertex { position: [-1.0, 1.0], color: border_color },
        ]
    }
}

struct State {
    surface: wgpu::Surface<'static>,
    device: wgpu::Device,
    queue: wgpu::Queue,
    config: wgpu::SurfaceConfiguration,
    size: winit::dpi::PhysicalSize<u32>,
    window: std::sync::Arc<winit::window::Window>,
    line_pipeline: wgpu::RenderPipeline,
    line_list_pipeline: wgpu::RenderPipeline,
    vertex_buffer: wgpu::Buffer,
    time: f32,
    graphs: Vec<GraphView>,
}

impl State {
    async fn new(window: std::sync::Arc<winit::window::Window>) -> Self {
        let size = window.inner_size();

        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
            backends: wgpu::Backends::VULKAN,
            ..Default::default()
        });

        let surface = instance.create_surface(window.clone()).unwrap();

        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                compatible_surface: Some(&surface),
                force_fallback_adapter: false,
            })
            .await
            .unwrap();

        let (device, queue) = adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    required_features: wgpu::Features::empty(),
                    required_limits: wgpu::Limits::default(),
                    label: None,
                    memory_hints: Default::default(),
                },
                None,
            )
            .await
            .unwrap();

        let surface_caps = surface.get_capabilities(&adapter);
        let surface_format = surface_caps
            .formats
            .iter()
            .copied()
            .find(|f| f.is_srgb())
            .unwrap_or(surface_caps.formats[0]);

        let config = wgpu::SurfaceConfiguration {
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            format: surface_format,
            width: size.width,
            height: size.height,
            present_mode: surface_caps.present_modes[0],
            alpha_mode: surface_caps.alpha_modes[0],
            view_formats: vec![],
            desired_maximum_frame_latency: 2,
        };
        surface.configure(&device, &config);

        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("Waterfall Shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("shader.wgsl").into()),
        });

        let render_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Render Pipeline Layout"),
                bind_group_layouts: &[],
                push_constant_ranges: &[],
            });

        let line_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Line Strip Pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: "vs_main",
                buffers: &[Vertex::desc()],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: "fs_main",
                targets: &[Some(wgpu::ColorTargetState {
                    format: config.format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::LineStrip,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: None,
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            multiview: None,
            cache: None,
        });

        let line_list_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Line List Pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: "vs_main",
                buffers: &[Vertex::desc()],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: "fs_main",
                targets: &[Some(wgpu::ColorTargetState {
                    format: config.format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::LineList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: None,
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            multiview: None,
            cache: None,
        });

        // Create initial empty buffer (will be updated each frame)
        let vertex_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Vertex Buffer"),
            size: (std::mem::size_of::<Vertex>() * 100 * 100) as u64, // Larger buffer for multiple views
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        // Create 2 graph views side-by-side
        let graphs = vec![
            GraphView::new(Viewport { x: 0.0, y: 0.0, width: 0.5, height: 1.0 }),
            GraphView::new(Viewport { x: 0.5, y: 0.0, width: 0.5, height: 1.0 }),
        ];

        Self {
            surface,
            device,
            queue,
            config,
            size,
            window,
            line_pipeline,
            line_list_pipeline,
            vertex_buffer,
            time: 0.0,
            graphs,
        }
    }

    fn resize(&mut self, new_size: winit::dpi::PhysicalSize<u32>) {
        if new_size.width > 0 && new_size.height > 0 {
            self.size = new_size;
            self.config.width = new_size.width;
            self.config.height = new_size.height;
            self.surface.configure(&self.device, &self.config);
        }
    }

    fn update(&mut self) {
        self.time += 0.016; // ~60fps

        // Update each graph independently
        for (graph_idx, graph) in self.graphs.iter_mut().enumerate() {
            // Add new line every 10 frames
            if (self.time * 60.0) as u32 % 10 == 0 && graph.lines.len() < 50 {
                let mut line = Vec::new();
                let phase = self.time + (graph_idx as f32 * 2.0);
                let freq = 2.0 + (self.time * 0.5 + graph_idx as f32).sin() * 1.0;

                for i in 0..100 {
                    let x = (i as f32 / 100.0) * 2.0 - 1.0;
                    let y = ((i as f32) * 0.1 * freq + phase).sin() * 0.3;

                    // Different color per graph
                    let hue = (self.time * 0.1 + graph_idx as f32 * 0.5) % 1.0;
                    let color = [
                        (hue * 6.0).sin().abs(),
                        ((hue + 0.33) * 6.0).sin().abs(),
                        ((hue + 0.66) * 6.0).sin().abs(),
                    ];

                    line.push(Vertex {
                        position: [x, y],
                        color,
                    });
                }
                graph.lines.push(line);
            }

            // Scroll lines down
            for line in graph.lines.iter_mut() {
                for vertex in line.iter_mut() {
                    vertex.position[1] -= 0.01;
                }
            }

            // Remove lines that have scrolled off screen
            graph.lines.retain(|line| {
                line.first().map(|v| v.position[1] > -1.1).unwrap_or(false)
            });
        }
    }

    fn render(&mut self) -> Result<(), wgpu::SurfaceError> {
        let output = self.surface.get_current_texture()?;
        let view = output
            .texture
            .create_view(&wgpu::TextureViewDescriptor::default());

        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Render Encoder"),
            });

        // Collect all vertex data for all graphs
        struct DrawData {
            viewport: Viewport,
            border_offset: usize,
            border_count: usize,
            grid_offset: usize,
            grid_count: usize,
            show_grid: bool,
            lines_offset: usize,
            lines_count: usize,
            num_lines: usize,
        }

        let mut all_vertices = Vec::new();
        let mut draw_data = Vec::new();

        for graph in &self.graphs {
            let border_vertices = graph.generate_border();
            let border_offset = all_vertices.len();
            all_vertices.extend_from_slice(&border_vertices);
            let border_count = border_vertices.len();

            let grid_offset = all_vertices.len();
            let grid_vertices = graph.generate_grid_lines();
            all_vertices.extend_from_slice(&grid_vertices);
            let grid_count = grid_vertices.len();

            let lines_offset = all_vertices.len();
            let mut line_vertices = Vec::new();
            for line in &graph.lines {
                line_vertices.extend_from_slice(line);
            }
            all_vertices.extend_from_slice(&line_vertices);
            let lines_count = line_vertices.len();

            draw_data.push(DrawData {
                viewport: graph.viewport.clone(),
                border_offset,
                border_count,
                grid_offset,
                grid_count,
                show_grid: graph.show_grid,
                lines_offset,
                lines_count,
                num_lines: graph.lines.len(),
            });
        }

        // Write all vertices at once
        if !all_vertices.is_empty() {
            self.queue.write_buffer(
                &self.vertex_buffer,
                0,
                bytemuck::cast_slice(&all_vertices),
            );
        }

        {
            let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("Render Pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color {
                            r: 0.1,
                            g: 0.1,
                            b: 0.15,
                            a: 1.0,
                        }),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                occlusion_query_set: None,
                timestamp_writes: None,
            });

            // Render each graph view
            for data in &draw_data {
                // Set viewport for this graph
                render_pass.set_viewport(
                    data.viewport.x * self.size.width as f32,
                    data.viewport.y * self.size.height as f32,
                    data.viewport.width * self.size.width as f32,
                    data.viewport.height * self.size.height as f32,
                    0.0,
                    1.0,
                );

                render_pass.set_vertex_buffer(0, self.vertex_buffer.slice(..));

                // Draw border
                render_pass.set_pipeline(&self.line_list_pipeline);
                render_pass.draw(
                    data.border_offset as u32..(data.border_offset + data.border_count) as u32,
                    0..1,
                );

                // Draw grid if enabled
                if data.show_grid {
                    render_pass.set_pipeline(&self.line_list_pipeline);
                    render_pass.draw(
                        data.grid_offset as u32..(data.grid_offset + data.grid_count) as u32,
                        0..1,
                    );
                }

                // Draw waterfall lines
                if data.lines_count > 0 {
                    render_pass.set_pipeline(&self.line_pipeline);
                    let points_per_line = 100;
                    for i in 0..data.num_lines {
                        let start = (data.lines_offset + i * points_per_line) as u32;
                        let end = start + points_per_line as u32;
                        render_pass.draw(start..end, 0..1);
                    }
                }
            }
        }

        self.queue.submit(std::iter::once(encoder.finish()));
        output.present();

        Ok(())
    }
}

fn main() {
    env_logger::init();
    let event_loop = EventLoop::new().unwrap();
    let window = std::sync::Arc::new(
        event_loop
            .create_window(
                winit::window::Window::default_attributes()
                    .with_title("TimePlot - Waterfall Display")
                    .with_inner_size(winit::dpi::LogicalSize::new(1280, 720))
            )
            .unwrap(),
    );

    let mut state = pollster::block_on(State::new(window.clone()));

    event_loop
        .run(move |event, control_flow| match event {
            Event::WindowEvent {
                ref event,
                window_id,
            } if window_id == state.window.id() => match event {
                WindowEvent::CloseRequested => control_flow.exit(),
                WindowEvent::Resized(physical_size) => {
                    state.resize(*physical_size);
                }
                WindowEvent::KeyboardInput {
                    event: key_event,
                    ..
                } => {
                    if key_event.state == winit::event::ElementState::Pressed {
                        if let PhysicalKey::Code(keycode) = key_event.physical_key {
                            match keycode {
                                KeyCode::KeyG => {
                                    println!("Grid toggle pressed");
                                    // Toggle grid for all graphs
                                    for graph in &mut state.graphs {
                                        graph.show_grid = !graph.show_grid;
                                        println!("Grid now: {}", graph.show_grid);
                                    }
                                }
                                KeyCode::Escape => control_flow.exit(),
                                _ => {}
                            }
                        }
                    }
                }
                WindowEvent::RedrawRequested => {
                    state.update();
                    match state.render() {
                        Ok(_) => {}
                        Err(wgpu::SurfaceError::Lost) => state.resize(state.size),
                        Err(wgpu::SurfaceError::OutOfMemory) => control_flow.exit(),
                        Err(e) => eprintln!("{:?}", e),
                    }
                }
                _ => {}
            },
            Event::AboutToWait => {
                state.window.request_redraw();
            }
            _ => {}
        })
        .unwrap();
}