/// Represents one of the motions specified in the input, either up,
/// down, left, or right by a given distance (or number of steps).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Motion {
    Up(u8),
    Down(u8),
    Left(u8),
    Right(u8),
}

/// Module wrapping the parser for today's puzzle. Produces a `Vec<Motion>`.
mod parser {
    use super::*;
    use anyhow::{anyhow, Result};
    use nom::{
        branch::alt, bytes::complete::tag, character::complete::u8, combinator::map,
        multi::separated_list1, sequence::preceded, Finish, IResult,
    };

    /// Nom parser for "U 5" -> Motion::Up(5)
    fn up(s: &str) -> IResult<&str, Motion> {
        map(preceded(tag("U "), u8), Motion::Up)(s)
    }

    /// Nom parser for "D 5" -> Motion::Down(5)
    fn down(s: &str) -> IResult<&str, Motion> {
        map(preceded(tag("D "), u8), Motion::Down)(s)
    }

    /// Nom parser for "L 5" -> Motion::Left(5)
    fn left(s: &str) -> IResult<&str, Motion> {
        map(preceded(tag("L "), u8), Motion::Left)(s)
    }

    /// Nom parser for "R 5" -> Motion::Right(5)
    fn right(s: &str) -> IResult<&str, Motion> {
        map(preceded(tag("R "), u8), Motion::Right)(s)
    }

    /// Nom parser to take all the lines of the input and produce a vector
    /// of `Motion`s
    pub fn parse(s: &str) -> Result<Vec<Motion>> {
        let result = separated_list1(tag("\n"), alt((up, down, left, right)))(s);
        let (_, motions) = result
            .finish()
            .map_err(|e| anyhow!("Could not parse motions! \n{e}"))?;
        Ok(motions)
    }
}

const INPUT: &str = include_str!("../../input/09/input.txt");

/// Parse that input!
fn read() -> Vec<Motion> {
    parser::parse(INPUT).unwrap()
}

use std::collections::HashSet;
use std::ops::AddAssign;

/// Solve Day 9, Part 1
fn solve(input: &[Motion]) -> u32 {
    // Brand new `RopeSimulator`(tm)
    let mut simulator = RopeSimulator::new();

    // For each specified motion, update the simulator
    input.iter().for_each(|motion| simulator.move_head(motion));

    // Return the number of unique tail positions from the simulator
    simulator.hist.len() as u32
}

/// Represents the position of a knot in x/y space.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
struct Knot(i32, i32);

/// Check if the current knot is "too far" from another knot. This means
/// it's more than one unit away in either dimension.
impl Knot {
    fn too_far(&self, other: &Knot) -> bool {
        let Knot(x1, y1) = self;
        let Knot(x2, y2) = other;
        x1.abs_diff(*x2) > 1 || y1.abs_diff(*y2) > 1
    }
}

/// You know what's nice? Adding an offset to a knot to 'move' the position!
impl AddAssign<(i32, i32)> for Knot {
    fn add_assign(&mut self, rhs: (i32, i32)) {
        let Knot(x, y) = self;
        let (xd, yd) = rhs;
        *self = Knot(*x + xd, *y + yd);
    }
}


/// A struct to encapsulate the state of the rope, with a HashSet to keep up with
/// the unique positions of the tail.
struct RopeSimulator {
    head: Knot,
    tail: Knot,
    hist: HashSet<Knot>,
}

impl RopeSimulator {
    /// Create a new RopeSimulator with head and tail both at the origin (0, 0) and
    /// initializing the set of tail positions to contain the initial tail position.
    fn new() -> Self {
        let head = Knot::default();
        let tail = Knot::default();
        let hist = HashSet::from([Knot::default()]);
        RopeSimulator { head, tail, hist }
    }

    /// Move the tail of the rope towards the head according to the rules given by
    /// the puzzle instructions.
    fn move_tail(&mut self) {
        let Knot(hx, hy) = self.head;
        let Knot(tx, ty) = self.tail;

        // This `use` statement means we don't have to fully quality all the
        // different `Ordering` variants below. Makes it cleaner to look at.
        use std::cmp::Ordering::*;

        // Set the position of the tail based on where the head is relative to it.
        // For example, if the head is positioned diagonally up and to the left of
        // the tail, then we'll match on `(Less, Less)` and move the tail up and
        // to the left.
        self.tail = match (hx.cmp(&tx), hy.cmp(&ty)) {
            (Less, Less) => Knot(tx - 1, ty - 1),
            (Less, Equal) => Knot(tx - 1, ty),
            (Less, Greater) => Knot(tx - 1, ty + 1),
            (Equal, Less) => Knot(tx, ty - 1),
            (Equal, Equal) => unreachable!(),
            (Equal, Greater) => Knot(tx, ty + 1),
            (Greater, Less) => Knot(tx + 1, ty - 1),
            (Greater, Equal) => Knot(tx + 1, ty),
            (Greater, Greater) => Knot(tx + 1, ty + 1),
        };

        // Add the new tail position to the set of tracked tail positions.
        self.hist.insert(self.tail);
    }

    fn move_head(&mut self, motion: &Motion) {
        // Generate a specification for moving the head. We get the number of
        // steps from the `Motion`, and the offset indicates how the `Knot`
        // of the head is changed on each step.
        let (steps, offset) = match motion {
            Motion::Up(steps) => (steps, (0, -1)),
            Motion::Down(steps) => (steps, (0, 1)),
            Motion::Left(steps) => (steps, (-1, 0)),
            Motion::Right(steps) => (steps, (1, 0)),
        };

        // Now we move the head one step at a time, adjusting the head by
        // the offset each time. Any time the head is too far from the
        // tail, we adjust the tail and record its new position.
        for _ in 0..*steps {
            self.head += offset;
            if self.head.too_far(&self.tail) {
                self.move_tail();
            }
        }
    }
}

use itertools::Itertools;
use std::collections::HashSet;

/// Same solve function, more or less. Just uses the updated rope simulator.
pub fn solve(input: &[Motion]) -> u32 {
    // Brand new `RopeSimulator`(tm)
    let mut simulator: RopeSimulator<10> = RopeSimulator::new();

    // For each specified motion, update the simulator
    input.iter().for_each(|motion| simulator.move_head(motion));

    // Return the number of unique tail positions from the simulator
    simulator.hist.len() as u32
}

// New and improved! I realize we only _needed_ to support a rope with 10
// knots, but at this point, why not make it *const generic*?
struct RopeSimulator<const N: usize> {
    knots: [Knot; N],
    hist: HashSet<Knot>,
}

impl<const N: usize> RopeSimulator<N> {
    // New RopeSimulator, this time we're keeping all the knots in an array
    // of length N (10 for our case). Note that the order of these knots
    // matters, since `knots[0]` will be the head and `knots[N - 1]` will be
    // the tail.
    fn new() -> Self {
        let knots = [Knot::default(); N];
        let hist = HashSet::from([Knot::default()]);
        RopeSimulator { knots, hist }
    }

    // This time, instead of hard-coding the head and the tail, we pass in
    // the index of the `leader` knot and the `follower` knot. For our
    // implementation, follower == leader + 1;
    fn follow(&mut self, leader: usize, follower: usize) {
        let Knot(hx, hy) = self.knots[leader];
        let Knot(tx, ty) = self.knots[follower];

        // The logic here is exactly the same as for the first part
        use std::cmp::Ordering::*;
        self.knots[follower] = match (hx.cmp(&tx), hy.cmp(&ty)) {
            (Less, Less) => Knot(tx - 1, ty - 1),
            (Less, Equal) => Knot(tx - 1, ty),
            (Less, Greater) => Knot(tx - 1, ty + 1),
            (Equal, Less) => Knot(tx, ty - 1),
            (Equal, Equal) => unreachable!(),
            (Equal, Greater) => Knot(tx, ty + 1),
            (Greater, Less) => Knot(tx + 1, ty - 1),
            (Greater, Equal) => Knot(tx + 1, ty),
            (Greater, Greater) => Knot(tx + 1, ty + 1),
        };

        // Now we need to check to be sure the `follower` is in the tail
        // slot before we record its position.
        if follower == N - 1 {
            self.hist.insert(self.knots[N - 1]);
        }
    }

    fn move_head(&mut self, motion: &Motion) {
        // Generate a specification for moving the head. We get the number of
        // steps from the `Motion`, and the offset indicates how the `Knot`
        // of the head is changed on each step.
        let (reps, offset) = match motion {
            Motion::Up(reps) => (reps, (0, -1)),
            Motion::Down(reps) => (reps, (0, 1)),
            Motion::Left(reps) => (reps, (-1, 0)),
            Motion::Right(reps) => (reps, (1, 0)),
        };

        // For each step in the motion, move the first knot in the `knots` array
        // (that's the head), then move down the array of knots, updating each
        // knot in sequence based on the position of the previous knot.
        for _ in 0..*reps {
            self.knots[0] += offset;

            // Note the `tuple_windows()` method from the `itertools` crate. Handy
            // crate, that `itertools`.
            for (leader, follower) in (0..N).tuple_windows() {
                // If the first knot is too far away from the knot behind it, move
                // the follower.
                if self.knots[leader].too_far(&self.knots[follower]) {
                    self.follow(leader, follower);
                }
            }
        }
    }
}