//! Generate geometry JSON from explosion for Three.js rendering.
//!
//! Run with: cargo run --example render_braiding
//!
//! Outputs fixtures/half_braid_geometry.json with vertices, wires, surfaces, volumes.

use std::collections::{HashMap, HashSet, VecDeque};
use std::fs;

use zigzag_engine::diagram::{Diagram, DiagramN};
use zigzag_engine::explosion::{HeightLabel, Point, Poset};
use zigzag_engine::import::load_homotopy_diagram_n;

use serde::Serialize;

#[derive(Serialize)]
struct Geometry {
    metadata: Metadata,
    vertices: Vec<Vertex>,
    wires: Vec<Wire>,
    surfaces: Vec<Surface>,
    volumes: Vec<Volume>,
}

#[derive(Serialize)]
struct Metadata {
    source: String,
    dimension: usize,
    total_points: usize,
    total_covers: usize,
}

#[derive(Serialize)]
struct Vertex {
    id: usize,
    label: String,
    point: String,
    coords: [f64; 3],
}

#[derive(Serialize)]
struct Wire {
    id: usize,
    label: String,
    point: String,
    coords: [f64; 3],        // The wire's own coordinate (midpoint/waypoint)
    endpoints: [usize; 2],   // Vertex IDs
    endpoint_coords: [[f64; 3]; 2],  // For convenience in renderer
}

#[derive(Serialize)]
struct Surface {
    id: usize,
    label: String,
    point: String,
    coords: [f64; 3],
    boundary_wires: Vec<usize>,  // Wire IDs on boundary
}

#[derive(Serialize)]
struct Volume {
    id: usize,
    label: String,
    point: String,
    coords: [f64; 3],
}

/// Format a point as a string like "s0,s1,r0"
fn format_point(p: &Point) -> String {
    p.0.iter()
        .map(|h| match h {
            HeightLabel::Regular(j) => format!("r{}", j),
            HeightLabel::Singular(j) => format!("s{}", j),
        })
        .collect::<Vec<_>>()
        .join(",")
}

/// Compute linear coordinates from height labels
fn linear_coords(p: &Point) -> [f64; 3] {
    let coords: Vec<f64> = p.0.iter().map(|h| h.to_linear_index() as f64).collect();
    // Pad to 3D if needed
    [
        coords.get(0).copied().unwrap_or(0.0),
        coords.get(1).copied().unwrap_or(0.0),
        coords.get(2).copied().unwrap_or(0.0),
    ]
}

/// Count singular labels in a point
fn singular_count(p: &Point) -> usize {
    p.0.iter().filter(|h| h.is_singular()).count()
}

/// Compute geometric dimension from singular count
fn geom_dim(p: &Point, n: usize) -> usize {
    n - singular_count(p)
}

fn main() {
    // Load diagram
    let json = fs::read_to_string("fixtures/half_braid.json")
        .expect("Failed to read fixtures/half_braid.json");
    let diagram_n = load_homotopy_diagram_n(&json)
        .expect("Failed to parse half_braid.json");
    let diagram = Diagram::DiagramN(diagram_n);

    let n = diagram.dimension();
    let pts = diagram.full_points();

    eprintln!("Loaded half_braid.json: dim={}, {} points, {} covers",
        n, pts.len(), pts.covers().len());

    // Group points by geometric dimension
    let mut by_geom_dim: HashMap<usize, Vec<usize>> = HashMap::new();
    for (idx, point) in pts.elements().iter().enumerate() {
        let gd = geom_dim(point, n);
        by_geom_dim.entry(gd).or_default().push(idx);
    }

    // Build adjacency for reachability
    let mut successors: Vec<Vec<usize>> = vec![vec![]; pts.len()];
    let mut predecessors: Vec<Vec<usize>> = vec![vec![]; pts.len()];
    for &(lower, upper) in pts.covers() {
        successors[lower].push(upper);
        predecessors[upper].push(lower);
    }

    // Helper: find all transitively reachable points (for wire→vertex connections)
    let reachable_from = |start: usize, adj: &[Vec<usize>]| -> HashSet<usize> {
        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();
        queue.push_back(start);
        visited.insert(start);
        while let Some(curr) = queue.pop_front() {
            for &next in &adj[curr] {
                if visited.insert(next) {
                    queue.push_back(next);
                }
            }
        }
        visited
    };

    // Get vertex indices
    let vertex_indices = by_geom_dim.get(&0).map(|v| v.as_slice()).unwrap_or(&[]);
    let vertex_set: HashSet<usize> = vertex_indices.iter().copied().collect();

    // Get wire indices
    let wire_indices = by_geom_dim.get(&1).map(|v| v.as_slice()).unwrap_or(&[]);

    // Build vertices
    let mut vertices: Vec<Vertex> = Vec::new();
    for (i, &idx) in vertex_indices.iter().enumerate() {
        let point = &pts.elements()[idx];
        vertices.push(Vertex {
            id: idx,
            label: format!("vertex_{}", i),
            point: format_point(point),
            coords: linear_coords(point),
        });
    }

    // Build wires with endpoint connections
    let mut wires: Vec<Wire> = Vec::new();
    for (i, &idx) in wire_indices.iter().enumerate() {
        let point = &pts.elements()[idx];

        // Find connected vertices via TRANSITIVE reachability
        // A wire (strand) spans between vertices even if the poset path has intermediate points
        let reachable_up = reachable_from(idx, &successors);
        let reachable_down = reachable_from(idx, &predecessors);

        let mut connected: Vec<usize> = reachable_up
            .union(&reachable_down)
            .filter(|v| vertex_set.contains(v))
            .copied()
            .collect();
        connected.sort();
        connected.dedup();

        // Default to first two vertices if we don't have exactly 2
        let endpoints = if connected.len() >= 2 {
            [connected[0], connected[1]]
        } else if connected.len() == 1 {
            // Wire connects to only one vertex - use same vertex twice
            // (This represents a loop or boundary wire)
            [connected[0], connected[0]]
        } else {
            // No connected vertices found - use first two vertices as fallback
            [vertex_indices[0], vertex_indices.get(1).copied().unwrap_or(vertex_indices[0])]
        };

        let endpoint_coords = [
            linear_coords(&pts.elements()[endpoints[0]]),
            linear_coords(&pts.elements()[endpoints[1]]),
        ];

        wires.push(Wire {
            id: idx,
            label: format!("wire_{}", i),
            point: format_point(point),
            coords: linear_coords(point),
            endpoints,
            endpoint_coords,
        });
    }

    // Build surfaces (geom_dim=2)
    let surface_indices = by_geom_dim.get(&2).map(|v| v.as_slice()).unwrap_or(&[]);
    let wire_set: HashSet<usize> = wire_indices.iter().copied().collect();

    let mut surfaces: Vec<Surface> = Vec::new();
    for (i, &idx) in surface_indices.iter().enumerate() {
        let point = &pts.elements()[idx];

        // Find boundary wires via DIRECT covering relations only
        // Surface (geom_dim=2) connects to wires (geom_dim=1) via single covers
        let mut boundary_wires: Vec<usize> = successors[idx]
            .iter()
            .chain(predecessors[idx].iter())
            .filter(|v| wire_set.contains(v))
            .copied()
            .collect();
        boundary_wires.sort();

        surfaces.push(Surface {
            id: idx,
            label: format!("surface_{}", i),
            point: format_point(point),
            coords: linear_coords(point),
            boundary_wires,
        });
    }

    // Build volumes (geom_dim=3)
    let volume_indices = by_geom_dim.get(&3).map(|v| v.as_slice()).unwrap_or(&[]);

    let mut volumes: Vec<Volume> = Vec::new();
    for (i, &idx) in volume_indices.iter().enumerate() {
        let point = &pts.elements()[idx];
        volumes.push(Volume {
            id: idx,
            label: format!("volume_{}", i),
            point: format_point(point),
            coords: linear_coords(point),
        });
    }

    // Build output
    let geometry = Geometry {
        metadata: Metadata {
            source: "half_braid.json".to_string(),
            dimension: n,
            total_points: pts.len(),
            total_covers: pts.covers().len(),
        },
        vertices,
        wires,
        surfaces,
        volumes,
    };

    // Output JSON
    let json_output = serde_json::to_string_pretty(&geometry).expect("Failed to serialize");

    // Write to file
    fs::write("fixtures/half_braid_geometry.json", &json_output)
        .expect("Failed to write fixtures/half_braid_geometry.json");

    eprintln!("\nWrote fixtures/half_braid_geometry.json");
    eprintln!("  {} vertices", geometry.vertices.len());
    eprintln!("  {} wires", geometry.wires.len());
    eprintln!("  {} surfaces", geometry.surfaces.len());
    eprintln!("  {} volumes", geometry.volumes.len());

    // Also print to stdout for piping
    println!("{}", json_output);
}