//
//  File.swift
//
//
//  Created by Max Nuding on 05.12.21.
//

import Foundation
import Runner

struct Coord: Hashable {
    let row: Int
    let col: Int
    
    func getLeft() -> Coord? {
        guard col > 0 else {
            return nil
        }
        return Coord(row: row, col: col - 1)
    }
    
    func getRight() -> Coord { Coord(row: row, col: col + 1) }
    
    func getAbove() -> Coord? {
        guard row > 0 else {
            return nil
        }
        return Coord(row: row - 1, col: col)
    }
    
    func getBelow() -> Coord { Coord(row: row + 1, col: col) }
}

struct Field:Sequence, IteratorProtocol {
    typealias Element = Coord
    let numbers: [[Int]]
    let numCols: Int
    let numRows: Int
    
    var neighbors = [Coord:[Coord]]()
    
    var current: Coord? = Coord(row: 0, col: 0)
    
    init(numbers: [[Int]]) {
        self.numbers = numbers
        numCols = numbers.first?.count ?? 0
        numRows = numbers.count
    }
    
    mutating func getNeighborsCoordsFor(coord: Coord) -> [Coord] {
        let n = neighbors[coord] ?? [coord.getAbove(), coord.getRight(), coord.getBelow(), coord.getLeft()]
            .compactMap { $0 }
            .filter { $0.row < numRows && $0.col < numCols }
        neighbors[coord] = n
        return n
    }
    
    mutating func isLowPointAt(coord: Coord) -> Bool {
        let value = self[coord]
        return getNeighborsCoordsFor(coord: coord)
            .map { numbers[$0.row][$0.col] }
            .allSatisfy { $0 > value }
    }
    
    private func getNextCoord() -> Coord? {
        guard let current = current else {
            return nil
        }
        if current.col == numCols - 1 && current.row == numRows - 1 {
            return nil
        }
        if current.col == numCols - 1 {
            return Coord(row: current.row + 1, col: 0)
        }
        return Coord(row: current.row, col: current.col + 1)
    }
    
    mutating func next() -> Coord? {
        defer {
            current = getNextCoord()
        }
        return current
    }
    
    subscript(index: Coord) -> Int {
        get {
            numbers[index.row][index.col]
        }/*
        set(newValue) {
            // Perform a suitable setting action here.
        }*/
    }
}

class  Day09: Runnable {
    let inputPath: String
    var lowestNeighbors = [Coord: Coord]()
    var lowPoints = Set<Coord>()
    var field = Field(numbers: [[Int]]())
    
    required init(inputPath: String) {
        self.inputPath = inputPath
    }
    
    public func run() {
        let input = try! String(contentsOfFile: inputPath)
        let entries = input
            .trimmingCharacters(in: .newlines)
            .components(separatedBy: .newlines)
            .map { line in line.map { $0.wholeNumberValue! } }
        field = Field(numbers: entries)
        lowPoints = Set(field.filter { field.isLowPointAt(coord: $0) })
        
        //Part 1
        let sum = lowPoints.map { field[$0] }.reduce(0, +) + lowPoints.count
        print(sum)
        
        // Part 2:
        var basinSize = [Coord: Int]()
        for coord in field {
            guard field[coord] != 9 else {
                continue
            }
            let bl = basinLocation(for: coord)
            basinSize[bl, default: 0] += 1
        }
        let result = basinSize.values.sorted(by: >)[0..<3].reduce(1, *)
        print(result)
    }
    
    private func basinLocation(for coord: Coord) -> Coord {
        guard !lowPoints.contains(coord) else {
            return coord
        }
        let closestNeighbor = lowestNeighbors[coord] ?? field.getNeighborsCoordsFor(coord: coord).min(by: { field[$0] < field[$1] })!
        lowestNeighbors[coord] = closestNeighbor
        return basinLocation(for: closestNeighbor)
    }
}