import Control.Monad
import Data.Bifunctor
import Data.List
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set

type C = (Int, Int)

readInput :: String -> (Map C Char, [C])
readInput s =
  let (room, _ : moves) = break null $ lines s
   in ( Map.fromList [((i, j), c) | (i, l) <- zip [0 ..] room, (j, c) <- zip [0 ..] l],
        map dir $ concat moves
      )
  where
    dir '^' = (-1, 0)
    dir 'v' = (1, 0)
    dir '<' = (0, -1)
    dir '>' = (0, 1)

moveInto :: Int -> Set C -> C -> C -> Set C -> Maybe (Set C)
moveInto boxWidth walls (di, dj) = go
  where
    go (i, j) boxes
      | (i, j) `Set.member` walls = Nothing
      | Just j' <- find (\j' -> (i, j') `Set.member` boxes) $ map (j -) [0 .. boxWidth - 1] =
          Set.insert (i + di, j' + dj)
            <$> foldM
              (flip go)
              (Set.delete (i, j') boxes)
              [(i + di, j' + z + dj) | z <- [0 .. boxWidth - 1]]
      | otherwise = Just boxes

runMoves :: (Map C Char, [C]) -> Int -> Int
runMoves (room, moves) scale = score $ snd $ foldl' move (start, boxes) moves
  where
    room' = Map.mapKeysMonotonic (second (* scale)) room
    Just start = fst <$> find ((== '@') . snd) (Map.assocs room')
    walls =
      let ps = Map.keysSet $ Map.filter (== '#') room'
       in Set.unions [Set.mapMonotonic (second (+ z)) ps | z <- [0 .. scale - 1]]
    boxes = Map.keysSet $ Map.filter (== 'O') room'
    move (pos@(i, j), boxes) dir@(di, dj) =
      let pos' = (i + di, j + dj)
       in maybe (pos, boxes) (pos',) $ moveInto scale walls dir pos' boxes
    score = sum . map (\(i, j) -> i * 100 + j) . Set.elems

main = do
  input <- readInput <$> readFile "input15"
  mapM_ (print . runMoves input) [1, 2]

type
  Vec2 = tuple[x,y: int]
  Box = array[2, Vec2]
  Dir = enum
    U = "^"
    R = ">"
    D = "v"
    L = "<"

proc convertPart2(grid: seq[string]): seq[string] =
  for y in 0..grid.high:
    result.add ""
    for x in 0..grid[0].high:
      result[^1] &= (
        if grid[y][x] == 'O': "[]"
        elif grid[y][x] == '#': "##"
        else: "..")

proc shiftLeft(grid: var seq[string], col: int, range: HSlice[int,int]) =
  for i in range.a ..< range.b:
    grid[col][i] = grid[col][i+1]
  grid[col][range.b] = '.'

proc shiftRight(grid: var seq[string], col: int, range: HSlice[int,int]) =
  for i in countDown(range.b, range.a+1):
    grid[col][i] = grid[col][i-1]
  grid[col][range.a] = '.'

proc box(pos: Vec2, grid: seq[string]): array[2, Vec2] =
  if grid[pos.y][pos.x] == '[':
    [pos, (pos.x+1, pos.y)]
  else:
    [(pos.x-1, pos.y), pos]

proc step(grid: var seq[string], bot: var Vec2, dir: Dir) =
  var (x, y) = bot
  case dir
  of U:
    while (dec y; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y-1][bot.x] == 'O': swap(grid[bot.y-1][bot.x], grid[y][x])
    dec bot.y
  of R:
    while (inc x; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y][bot.x+1] == 'O': swap(grid[bot.y][bot.x+1], grid[y][x])
    inc bot.x
  of L:
    while (dec x; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y][bot.x-1] == 'O': swap(grid[bot.y][bot.x-1], grid[y][x])
    dec bot.x
  of D:
    while (inc y; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y+1][bot.x] == 'O': swap(grid[bot.y+1][bot.x], grid[y][x])
    inc bot.y

proc canMoveVert(box: Box, grid: seq[string], boxes: var HashSet[Box], dy: int): bool =
  boxes.incl box
  var left, right = false
  let (lbox, rbox) = (box[0], box[1])
  let lbigBox = box((lbox.x, lbox.y+dy), grid)
  let rbigBox = box((rbox.x, lbox.y+dy), grid)

  if grid[lbox.y+dy][lbox.x] == '#' or
     grid[rbox.y+dy][rbox.x] == '#': return false
  elif grid[lbox.y+dy][lbox.x] == '.': left = true
  else:
    left = canMoveVert(box((lbox.x,lbox.y+dy), grid), grid, boxes, dy)

  if grid[rbox.y+dy][rbox.x] == '.': right = true
  elif lbigBox == rbigBox: right = left
  else:
    right = canMoveVert(box((rbox.x, rbox.y+dy), grid), grid, boxes, dy)

  left and right

proc moveBoxes(grid: var seq[string], boxes: var HashSet[Box], d: Vec2) =
  for box in boxes:
    grid[box[0].y][box[0].x] = '.'
    grid[box[1].y][box[1].x] = '.'
  for box in boxes:
    grid[box[0].y+d.y][box[0].x+d.x] = '['
    grid[box[1].y+d.y][box[1].x+d.x] = ']'
  boxes.clear()

proc step2(grid: var seq[string], bot: var Vec2, dir: Dir) =
  case dir
  of U:
    if grid[bot.y-1][bot.x] == '#': return
    if grid[bot.y-1][bot.x] == '.': dec bot.y
    else:
      var boxes: HashSet[Box]
      if canMoveVert(box((x:bot.x, y:bot.y-1), grid), grid, boxes, -1):
        grid.moveBoxes(boxes, (0, -1))
        dec bot.y
  of R:
    var (x, y) = bot
    while (inc x; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y][bot.x+1] == '[': grid.shiftRight(bot.y, bot.x+1..x)
    inc bot.x
  of L:
    var (x, y) = bot
    while (dec x; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y][bot.x-1] == ']': grid.shiftLeft(bot.y, x..bot.x-1)
    dec bot.x
  of D:
    if grid[bot.y+1][bot.x] == '#': return
    if grid[bot.y+1][bot.x] == '.': inc bot.y
    else:
      var boxes: HashSet[Box]
      if canMoveVert(box((x:bot.x, y:bot.y+1), grid), grid, boxes, 1):
        grid.moveBoxes(boxes, (0, 1))
        inc bot.y


proc solve(input: string): AOCSolution[int, int] =
  let chunks = input.split("\n\n")
  var grid = chunks[0].splitLines()
  let movements = chunks[1].splitLines().join().join()

  var robot: Vec2
  for y in 0..grid.high:
    for x in 0..grid[0].high:
      if grid[y][x] == '@':
        grid[y][x] = '.'
        robot = (x,y)

  block p1:
    var grid = grid
    var robot = robot
    for m in movements:
      let dir = parseEnum[Dir]($m)
      step(grid, robot, dir)
    for y in 0..grid.high:
      for x in 0..grid[0].high:
        if grid[y][x] == 'O':
          result.part1 += 100 * y + x

  block p2:
    var grid = grid.convertPart2()
    var robot = (robot.x*2, robot.y)
    for m in movements:
      let dir = parseEnum[Dir]($m)
      step2(grid, robot, dir)
      #grid.inspect(robot)

    for y in 0..grid.high:
      for x in 0..grid[0].high:
        if grid[y][x] == '[':
          result.part2 += 100 * y + x

import fs from "fs";

type Point = {x: number, y: number};

enum Direction {
    UP = '^',
    DOWN = 'v',
    LEFT = '<',
    RIGHT = '>'
}

const input = fs.readFileSync("./15/input.txt", "utf-8").split(/[\r\n]{4,}/);
const warehouse: string[][] = input[0]
    .split(/[\r\n]+/)
    .map(row => row.split(""));
const movements: Direction[] = input[1]
    .split("")
    .map(char => char.trim())
    .filter(Boolean)
    .map(char => char as Direction);

// Part 1
console.info("Part 1: " + solve(warehouse, movements));

// Part 2
const secondWarehouse = warehouse.map(row => {
    const newRow: string[] = [];
    for (const char of row) {
        if (char === '#') { newRow.push('#', '#'); }
        else if (char === 'O') { newRow.push('[', ']'); }
        else if (char === '.') { newRow.push('.', '.'); }
        else { newRow.push('@', '.'); }
    }
    return newRow;
});
console.info("Part 2: " + solve(secondWarehouse, movements));

function solve(warehouse: string[][], movements: Direction[]): number {
    let _warehouse = warehouse.map(row => [...row]); // Take a copy to avoid modifying the original
    const robotLocation: Point = findStartLocation(_warehouse);

    for (const move of movements) {
        // Under some very specific circumstances in part 2, tryMove returns false, but the grid has already been modified
        // "Fix" the issue ba taking a copy so we can easily revert all changes made
        // Slow AF of course but rest of this code isn't optimized either, so...
        const copy = _warehouse.map(row => [...row]);
    
        if (tryMove(robotLocation, move, _warehouse)) {
            if (move === Direction.UP) { robotLocation.y--; }
            else if (move === Direction.DOWN) { robotLocation.y++; }
            else if (move === Direction.LEFT) { robotLocation.x--; }
            else { robotLocation.x++; }
        } else {
            _warehouse = copy; // Revert changes
        }
    }

    // GPS
    let result = 0;
    for (let y = 0; y < _warehouse.length; y++) {
        for (let x = 0; x < _warehouse[y].length; x++) {
            if (_warehouse[y][x] === "O" || _warehouse[y][x] === "[") {
                result += 100 * y + x;
            }
        }
    }
    return result;
}

function tryMove(from: Point, direction: Direction, warehouse: string[][], movingPair = false): boolean {
    const moveWhat = warehouse[from.y][from.x];
    if (moveWhat === "#") {
        return false;
    }

    let to: Point;
    switch (direction) {
        case Direction.UP: to = {x: from.x, y: from.y - 1}; break;
        case Direction.DOWN: to = {x: from.x, y: from.y + 1}; break;
        case Direction.LEFT: to = {x: from.x - 1, y: from.y}; break;
        case Direction.RIGHT: to = {x: from.x + 1, y: from.y}; break;
    }

    const allowMove = warehouse[to.y][to.x] === "."
        || (direction === Direction.UP && tryMove({x: from.x, y: from.y - 1}, direction, warehouse))
        || (direction === Direction.DOWN && tryMove({x: from.x, y: from.y + 1}, direction, warehouse))
        || (direction === Direction.LEFT && tryMove({x: from.x - 1, y: from.y}, direction, warehouse))
        || (direction === Direction.RIGHT && tryMove({x: from.x + 1, y: from.y}, direction, warehouse));

    if (allowMove) {
        // Part 1 logic handles horizontal movement of larger boxes just fine. Needs special handling for vertical movement
        if (!movingPair && (direction === Direction.UP || direction === Direction.DOWN)) {
            if (moveWhat === "[" && !tryMove({x: from.x + 1, y: from.y}, direction, warehouse, true)) {
                return false;
            }
            if (moveWhat === "]" && !tryMove({x: from.x - 1, y: from.y}, direction, warehouse, true)) {
                return false;
            }
        }

        // Make the move
        warehouse[to.y][to.x] = moveWhat;
        warehouse[from.y][from.x] = ".";
        return true;
    }
    return false;
}

function findStartLocation(warehouse: string[][]): Point {
    for (let y = 0; y < warehouse.length; y++) {
        for (let x = 0; x < warehouse[y].length; x++) {
            if (warehouse[y][x] === "@") {
                return {x,y};
            }
        }
    }

    throw new Error("Could not find start location!");
}

use euclid::{default::*, vec2};

// Common type for both parts. In part 1 all boxes are BoxL.
#[derive(Clone, Copy)]
enum Spot {
    Empty,
    BoxL,
    BoxR,
    Wall,
}

impl From<u8> for Spot {
    fn from(value: u8) -> Self {
        match value {
            b'.' | b'@' => Spot::Empty,
            b'O' => Spot::BoxL,
            b'#' => Spot::Wall,
            other => panic!("Invalid spot: {other}"),
        }
    }
}

fn parse(input: &str) -> (Vec<Vec<Spot>>, Point2D<i32>, Vec<Vector2D<i32>>) {
    let (field_s, moves_s) = input.split_once("\n\n").unwrap();
    let mut field = Vec::new();
    let mut robot = None;
    for (y, l) in field_s.lines().enumerate() {
        let mut row = Vec::new();
        for (x, b) in l.bytes().enumerate() {
            row.push(Spot::from(b));
            if b == b'@' {
                robot = Some(Point2D::new(x, y).to_i32())
            }
        }
        field.push(row);
    }

    let moves = moves_s
        .bytes()
        .filter(|b| *b != b'\n')
        .map(|b| match b {
            b'^' => vec2(0, -1),
            b'>' => vec2(1, 0),
            b'v' => vec2(0, 1),
            b'<' => vec2(-1, 0),
            other => panic!("Invalid move: {other}"),
        })
        .collect();
    (field, robot.unwrap(), moves)
}

fn gps(field: &[Vec<Spot>]) -> u32 {
    let mut sum = 0;
    for (y, row) in field.iter().enumerate() {
        for (x, s) in row.iter().enumerate() {
            if let Spot::BoxL = s {
                sum += x + 100 * y;
            }
        }
    }
    sum as u32
}

fn part1(input: String) {
    let (mut field, mut robot, moves) = parse(&input);
    for m in moves {
        let next = robot + m;
        match field[next.y as usize][next.x as usize] {
            Spot::Empty => robot = next, // Move into space
            Spot::BoxL => {
                let mut search = next + m;
                let can_move = loop {
                    match field[search.y as usize][search.x as usize] {
                        Spot::BoxL => {}
                        Spot::Wall => break false,
                        Spot::Empty => break true,
                        Spot::BoxR => unreachable!(),
                    }
                    search += m;
                };
                if can_move {
                    robot = next;
                    field[next.y as usize][next.x as usize] = Spot::Empty;
                    field[search.y as usize][search.x as usize] = Spot::BoxL;
                }
            }
            Spot::Wall => {} // Cannot move
            Spot::BoxR => unreachable!(),
        }
    }
    println!("{}", gps(&field));
}

// Transform part 1 field to wider part 2 field
fn widen(field: &[Vec<Spot>]) -> Vec<Vec<Spot>> {
    field
        .iter()
        .map(|row| {
            row.iter()
                .flat_map(|s| match s {
                    Spot::Empty => [Spot::Empty; 2],
                    Spot::Wall => [Spot::Wall; 2],
                    Spot::BoxL => [Spot::BoxL, Spot::BoxR],
                    Spot::BoxR => unreachable!(),
                })
                .collect()
        })
        .collect()
}

// Recursively find out whether or not the robot can move in direction `dir` from `start`.
fn can_move_rec(field: &[Vec<Spot>], start: Point2D<i32>, dir: Vector2D<i32>) -> bool {
    let next = start + dir;
    match field[next.y as usize][next.x as usize] {
        Spot::Empty => true,
        Spot::BoxL => can_move_rec(field, next, dir) && can_move_rec(field, next + vec2(1, 0), dir),
        Spot::BoxR => can_move_rec(field, next - vec2(1, 0), dir) && can_move_rec(field, next, dir),
        Spot::Wall => false,
    }
}

// Recursively execute a move for vertical directions into boxes.
fn do_move(field: &mut [Vec<Spot>], start: Point2D<i32>, dir: Vector2D<i32>) {
    let next = start + dir;
    match field[next.y as usize][next.x as usize] {
        Spot::Empty | Spot::Wall => {}
        Spot::BoxL => {
            do_move(field, next, dir);
            do_move(field, next + vec2(1, 0), dir);
            let move_to = next + dir;
            field[next.y as usize][next.x as usize] = Spot::Empty;
            field[next.y as usize][next.x as usize + 1] = Spot::Empty;
            field[move_to.y as usize][move_to.x as usize] = Spot::BoxL;
            field[move_to.y as usize][move_to.x as usize + 1] = Spot::BoxR;
        }
        Spot::BoxR => {
            do_move(field, next - vec2(1, 0), dir);
            do_move(field, next, dir);
            let move_to = next + dir;
            field[next.y as usize][next.x as usize - 1] = Spot::Empty;
            field[next.y as usize][next.x as usize] = Spot::Empty;
            field[move_to.y as usize][move_to.x as usize - 1] = Spot::BoxL;
            field[move_to.y as usize][move_to.x as usize] = Spot::BoxR;
        }
    }
}

fn part2(input: String) {
    let (field1, robot1, moves) = parse(&input);
    let mut field = widen(&field1);
    let mut robot = Point2D::new(robot1.x * 2, robot1.y);
    for m in moves {
        let next = robot + m;
        match field[next.y as usize][next.x as usize] {
            Spot::Empty => robot = next, // Move into space
            Spot::BoxL | Spot::BoxR if m.y == 0 => {
                let mut search = next + m;
                let can_move = loop {
                    match field[search.y as usize][search.x as usize] {
                        Spot::BoxL | Spot::BoxR => {}
                        Spot::Wall => break false,
                        Spot::Empty => break true,
                    }
                    search += m;
                };
                if can_move {
                    robot = next;
                    // Shift boxes by array remove/insert
                    field[next.y as usize].remove(search.x as usize);
                    field[next.y as usize].insert(next.x as usize, Spot::Empty);
                }
            }
            Spot::BoxL | Spot::BoxR => {
                if can_move_rec(&field, robot, m) {
                    do_move(&mut field, robot, m);
                    robot = next;
                }
            }
            Spot::Wall => {} // Cannot move
        }
    }
    println!("{}", gps(&field));
}

util::aoc_main!();