import std/[strutils, algorithm, sequtils]
import aoc

proc parseInput(input: string): seq[seq[int]] =
  for line in input.splitLines():
    result.add line.splitWhitespace().mapIt(parseInt(it))

proc checkPair(x, y: int): bool =
  if y < x: return false
  let dist = abs(y - x)
  if dist < 1 or dist > 3: return false
  return true

proc checkWindow(x, y, z: int): bool =
  if not [x, y, z].isSorted: return false
  let dists = [abs(z-y), abs(y-x)]
  if dists.filterIt(it < 1 or it > 3).len != 0: return false
  return true

proc isSafe(report: seq[int]): bool =
  var levels = report
  # make all lists ascending order
  if levels[0] > levels[1]:
    levels.reverse()
  for i in 0..<levels.len:
    if i == 0:
      if not checkPair(levels[i], levels[i + 1]):
        return false
    elif i == levels.len - 1:
      if not checkPair(levels[i - 1], levels[i]):
        return false
    else:
      if not checkWindow(levels[i - 1], levels[i], levels[i+1]):
        return false
  result = true


proc partOne*(input: string): int =
  parseInput(input).filterIt(it.isSafe()).len

proc leaveOneOutCheck(report: seq[int]): bool =
  let leaveOneOutReports =
    @[report[1..^1]] & (1..<report.len).mapIt(report[0 .. it - 1] & report[it + 1 .. ^1])

  for report in leaveOneOutReports:
    if isSafe report:
      return true

proc partTwo*(input: string): int =
  ## brutest of force
  for report in parseInput(input):
    if isSafe report:
      inc result
    else:
      if leaveOneOutCheck report:
        inc result

solve:
  "example.txt":
    partOne: 2
    partTwo: 4
  "input.txt":
    partOne: 624
    partTwo: 658