tilezz 0.2.0

//! `cyc_explore`: explore a cyclotomic ring's reachable lattice points and
//! render them.
//!
//! Enumerates the points of `Z[zeta_n]` reachable within a given number of
//! unit-vector rounds (`cyclotomic::explore::reachable_points`) and renders the
//! growing cloud as a PNG grid or an animated GIF via the `vis` point-cloud
//! backend. A thin CLI front-end; the reachability and rendering live in the
//! library.

use std::io::{Write, stdout};

use clap::Parser;

use tilezz::cyclotomic::Units;
use tilezz::cyclotomic::explore::reachable_points;
use tilezz::cyclotomic::geometry::point_mod_rect;
use tilezz::cyclotomic::*;
use tilezz::vis::plotutils::{P64, R64, points_bounds};
use tilezz::vis::pointcloud::{MarkerScale, animation_gif, grid_png};

/// Print a round's new-point count as live progress: the count, with a
/// ` + ` separator between rounds and none after the last.
fn print_round(k: usize, n: usize, count: usize) {
    print!("{}{}", count, if k == n { "" } else { " + " });
    stdout().flush().unwrap();
}

/// Close the progress line with the running total across all rounds.
fn finish_progress<ZZ>(rounds: &[Vec<ZZ>]) {
    let total: usize = rounds.iter().map(|v| v.len()).sum();
    println!("\n= {total}");
}

/// Free exploration: the raw reachable point cloud, no folding. Works for
/// every ring (the walk needs only `Units`).
fn explore_free<ZZ>(n: usize, num_threads: usize) -> Vec<Vec<ZZ>>
where
    ZZ: IsRing + Units + Send + Sync,
{
    let rounds = reachable_points(n, num_threads, |p| *p, |k, count| print_round(k, n, count));
    finish_progress(&rounds);
    rounds
}

/// Torus exploration: each step is folded back into the origin-centered
/// unit cell `[-1/2, 1/2)^2` (see [`point_mod_rect`]). Restricted to
/// `HasZZ4` rings: the cell folded into is the square period lattice
/// `Z + Z*i`, and the vertical shift `i*height` only exists when the
/// imaginary unit is in the ring (rings without `i` -- ZZ6/ZZ10/ZZ14 --
/// have no unit square to fold into).
fn explore_folded<ZZ>(n: usize, num_threads: usize) -> Vec<Vec<ZZ>>
where
    ZZ: HasZZ4 + Send + Sync,
{
    let anchor: ZZ = ZZ::zero();
    let rounds = reachable_points(
        n,
        num_threads,
        move |p| point_mod_rect(p, &anchor, (1, 1)),
        |k, count| print_round(k, n, count),
    );
    finish_progress(&rounds);
    rounds
}

fn prepare_render_free<ZZ>(num_rounds: usize, num_threads: usize) -> (Vec<Vec<P64>>, R64)
where
    ZZ: IsRing + Units + Send + Sync,
{
    let points: Vec<Vec<P64>> = explore_free::<ZZ>(num_rounds, num_threads)
        .iter()
        .map(|v| v.iter().map(|p| p.xy()).collect())
        .collect();
    // Free exploration bounds to the discovered point cloud.
    let bounds = points_bounds(points.iter()).unwrap_or(((-0.5, -0.5), (0.5, 0.5)));
    (points, bounds)
}

fn prepare_render_folded<ZZ>(num_rounds: usize, num_threads: usize) -> (Vec<Vec<P64>>, R64)
where
    ZZ: HasZZ4 + Send + Sync,
{
    let points: Vec<Vec<P64>> = explore_folded::<ZZ>(num_rounds, num_threads)
        .iter()
        .map(|v| v.iter().map(|p| p.xy()).collect())
        .collect();
    // The fundamental domain is exactly the origin-centered unit cell,
    // so draw that fixed frame rather than the point cloud's own extent.
    let bounds = ((-0.5, -0.5), (0.5, 0.5));
    (points, bounds)
}

fn prepare_render_for(
    ring: u8,
    num_rounds: usize,
    unit_square: bool,
    num_threads: usize,
) -> (Vec<Vec<P64>>, R64) {
    if unit_square {
        // The unit-square fold builds `i*height`, so only rings with a
        // ZZ4 subring (divisible by 4) qualify.
        match ring {
            4 => prepare_render_folded::<ZZ4>(num_rounds, num_threads),
            8 => prepare_render_folded::<ZZ8>(num_rounds, num_threads),
            12 => prepare_render_folded::<ZZ12>(num_rounds, num_threads),
            16 => prepare_render_folded::<ZZ16>(num_rounds, num_threads),
            20 => prepare_render_folded::<ZZ20>(num_rounds, num_threads),
            24 => prepare_render_folded::<ZZ24>(num_rounds, num_threads),
            32 => prepare_render_folded::<ZZ32>(num_rounds, num_threads),
            60 => prepare_render_folded::<ZZ60>(num_rounds, num_threads),
            // Non-divisible-by-4 rings are already rejected in main(); this
            // arm only sees a ring divisible by 4 but not compiled here (e.g.
            // 28) -- which DOES have an imaginary unit, it is just unsupported.
            _ => panic!(
                "ZZ{ring} is not a compiled ring for the unit-square fold \
                 (supported: 4, 8, 12, 16, 20, 24, 32, 60)"
            ),
        }
    } else {
        // Free exploration is ring-generic.
        tilezz::dispatch_ring!(ring, prepare_render_free::<ZZ>(num_rounds, num_threads))
    }
}

// ------------------------------------------------------------------------

#[derive(Clone, Copy)]
enum OutputFormat {
    Png,
    Gif,
}

#[derive(Parser, Debug)]
#[command(version = tilezz::VERSION, about = "Explore cyclotomic rings and render the discovered points", long_about = None)]
struct Cli {
    #[arg(short, long)]
    ring: u8,

    #[arg(
        short,
        long,
        help = "Number of BFS exploration rounds (distance from the starting point(s))"
    )]
    num_rounds: usize,

    #[arg(
        short,
        long,
        help = "Run exploration modulo the origin-centered unit square [-1/2, 1/2)^2"
    )]
    unit_square: bool,

    #[arg(
        short = 'o',
        long,
        help = "Filename (with .png or .gif extension), if missing => dry run"
    )]
    filename: Option<String>,

    #[arg(short, long, default_value_t = 1000, help = "Image width (in px)")]
    width: u32,

    #[arg(short, long, default_value_t = 500, help = "GIF frame delay (in ms)")]
    delay: u32,

    #[arg(short = 'p', long, default_value_t = 4, help = "PNG plots per row")]
    row: usize,

    #[arg(
        short,
        long,
        help = "Number of threads (= # of available cores if unset)"
    )]
    threads: Option<usize>,
}

#[cfg(feature = "cli")]
fn main() {
    let cli = Cli::parse();
    if cli.unit_square && cli.ring % 4 != 0 {
        panic!(
            "ZZ{} has no imaginary unit; --unit-square fold requires a ring divisible by 4 \
             (free exploration works for any ring)",
            cli.ring
        );
    }

    let filename = cli.filename.unwrap_or_default();
    let output_format = if filename.is_empty() {
        None
    } else if filename.ends_with(".gif") {
        Some(OutputFormat::Gif)
    } else if filename.ends_with(".png") {
        Some(OutputFormat::Png)
    } else {
        panic!("Unknown image format!")
    };

    // -------- Compute --------

    let num_threads = cli.threads.unwrap_or_else(tilezz::util::available_workers);

    let (points, bounds) =
        prepare_render_for(cli.ring, cli.num_rounds, cli.unit_square, num_threads);

    // -------- Render --------

    let Some(output_format) = output_format else {
        return; // dry run -> computation with no rendering
    };

    let marker = if cli.unit_square {
        MarkerScale::CellFraction(0.04)
    } else {
        MarkerScale::Pixels(6.0)
    };
    let bytes = match output_format {
        OutputFormat::Png => grid_png(&points, bounds, cli.row, cli.width, marker),
        OutputFormat::Gif => animation_gif(&points, bounds, cli.width, cli.delay as u16, marker),
    };
    std::fs::write(&filename, bytes).expect("write image");
}