use itertools::Itertools;

use fjadra::{
    force::{Link, ManyBody, SimulationBuilder},
    Node,
};
use rerun::{components::GraphType, Color, GraphEdges, GraphNodes};

const NUM_NODES: usize = 10;

fn main() -> anyhow::Result<()> {
    let rec = rerun::RecordingStreamBuilder::new("fjadra_lattice").spawn()?;

    let coordinates = (0..NUM_NODES).cartesian_product(0..NUM_NODES);

    let (nodes, colors): (Vec<_>, Vec<_>) = coordinates
        .clone()
        .enumerate()
        .map(|(i, (x, y))| {
            let r = ((x as f32 / (NUM_NODES - 1) as f32) * 255.0).round() as u8;
            let g = ((y as f32 / (NUM_NODES - 1) as f32) * 255.0).round() as u8;
            (i, Color::from_rgb(r, g, 0))
        })
        .unzip();

    let mut edges = Vec::new();
    for (x, y) in coordinates.clone() {
        if y > 0 {
            let source = (y - 1) * NUM_NODES + x;
            let target = y * NUM_NODES + x;
            edges.push((source, target));
        }
        if x > 0 {
            let source = y * NUM_NODES + (x - 1);
            let target = y * NUM_NODES + x;
            edges.push((source, target));
        }
    }

    let mut simulation = SimulationBuilder::default()
        .build(nodes.iter().map(|_| Node::default()))
        .add_force(
            "link",
            Link::new(edges.clone().into_iter())
                .strength(1.0)
                .distance(60.0)
                .iterations(10),
        )
        .add_force("charge", ManyBody::new().strength(-70.0));

    rec.log_static(
        "/lattice",
        &GraphEdges::new(
            edges
                .iter()
                .map(|(source, target)| (source.to_string(), target.to_string())),
        )
        .with_graph_type(GraphType::Directed),
    )?;

    for positions in simulation.iter() {
        rec.log(
            "/lattice",
            &GraphNodes::new(nodes.iter().map(|key| key.to_string()))
                .with_positions(positions.into_iter().map(|[x, y]| [x as f32, y as f32]))
                .with_colors(colors.clone())
                .with_labels(coordinates.clone().map(|(x, y)| format!("({x}, {y})"))),
        )?;
    }

    Ok(())
}