Rust

use std::{
    cmp::Reverse,
    collections::{HashMap, HashSet},
    f64::consts::PI,
};

use itertools::Itertools;
use libm::atan2;
use priority_queue::PriorityQueue;

fn l2(i: usize, j: usize) -> usize {
    i * i + j * j
}

fn erupt(data: &[Vec<char>], vi: usize, vj: usize, r: usize) -> usize {
    let r2 = r * r;
    data.iter()
        .enumerate()
        .flat_map(|(i, l)| {
            l.iter()
                .enumerate()
                .filter(move |&(j, v)| *v != '@' && l2(vi.abs_diff(i), vj.abs_diff(j)) <= r2)
        })
        .map(|(_, v)| (*v as u8 - b'0') as usize)
        .sum::<usize>()
}

fn find(data: &[Vec<char>], c: char) -> (usize, usize) {
    data.iter()
        .enumerate()
        .flat_map(|(i, l)| {
            l.iter()
                .enumerate()
                .filter(|&(_, v)| *v == c)
                .map(move |(j, _)| (i, j))
        })
        .exactly_one()
        .unwrap()
}

pub fn solve_part_1(input: &str) -> String {
    let data = input
        .lines()
        .map(|l| l.chars().collect::<Vec<_>>())
        .collect::<Vec<_>>();
    let (vi, vj) = find(&data, '@');
    erupt(&data, vi, vj, 10).to_string()
}

pub fn solve_part_2(input: &str) -> String {
    let data = input
        .lines()
        .map(|l| l.chars().collect::<Vec<_>>())
        .collect::<Vec<_>>();
    let (vi, vj) = find(&data, '@');
    let h = data.len();
    let w = data[0].len();
    let max_r = vi.max(vj).max(h - vi - 1).max(w - vj - 1);
    let mut last_eruption = 0;
    let mut max_eruption = 0;
    let mut max_eruption_r = 0;
    for r in 1..=max_r {
        let eruption = erupt(&data, vi, vj, r);
        let de = eruption - last_eruption;
        if de > max_eruption {
            max_eruption = de;
            max_eruption_r = r;
        }
        last_eruption = eruption;
    }
    (max_eruption_r * max_eruption).to_string()
}

pub fn solve_part_3(input: &str) -> String {
    let data = input
        .lines()
        .map(|l| l.chars().collect::<Vec<_>>())
        .collect::<Vec<_>>();
    let width = data[0].len();
    let height = data.len();
    let (vi, vj) = find(&data, '@');
    let (si, sj) = find(&data, 'S');

    let azimuth = |i: usize, j: usize| atan2(i as f64 - vi as f64, j as f64 - vj as f64);

    let small_rot = |az1: f64, az2: f64| {
        let d = az1 - az2;
        if d > PI {
            d - 2. * PI
        } else if d < -PI {
            d + 2. * PI
        } else {
            d
        }
    };

    let solve = |radius: usize| {
        let r2 = radius * radius;
        let time_limit = ((radius + 1) * 30) as i64;
        let mut queue = PriorityQueue::new();
        let mut rotations = HashMap::new();
        rotations.insert((si, sj, false), 0f64);
        let mut visited = HashSet::new();
        queue.push((si, sj, false), Reverse(0));
        while let Some(((i, j, rotated), Reverse(time))) = queue.pop() {
            if time >= time_limit {
                break;
            }
            visited.insert((i, j, rotated));
            let az = azimuth(i, j);
            let rotation = rotations[&(i, j, rotated)];
            for (di, dj) in [(-1, 0), (1, 0), (0, -1), (0, 1)] {
                let (ni, nj) = (i.wrapping_add_signed(di), j.wrapping_add_signed(dj));
                if ni >= height || nj >= width {
                    continue;
                }
                if l2(ni.abs_diff(vi), nj.abs_diff(vj)) <= r2 {
                    continue;
                }
                let is_rotated = if let Some(previous_rotation) = rotations.get(&(ni, nj, false)) {
                    let rotation = rotation + small_rot(azimuth(ni, nj), az);
                    (rotation - previous_rotation).abs() > 6.
                } else {
                    false
                };
                if (ni, nj, is_rotated) == (si, sj, true) {
                    return Some(time);
                }
                if visited.contains(&(ni, nj, is_rotated)) {
                    continue;
                }
                let new_time: i64 = time + (data[ni][nj] as i8 - '0' as i8) as i64;
                let should_update =
                    match queue.push_increase((ni, nj, is_rotated), Reverse(new_time)) {
                        None => true,
                        Some(Reverse(t)) => t > new_time,
                    };
                if should_update {
                    rotations.insert(
                        (ni, nj, is_rotated),
                        rotation + small_rot(azimuth(ni, nj), az),
                    );
                };
            }
        }
        None
    };
    let (radius, time) = (1..(width.min(height) / 2))
        .map(|radius| (radius, solve(radius)))
        .filter(|(_, s)| s.is_some())
        .min()
        .unwrap();
    (radius as i64 * time.unwrap()).to_string()
}