#! /usr/bin/env stack
-- stack --resolver lts-12.20 --install-ghc script --optimize
{-# LANGUAGE ViewPatterns #-}
module Main where

import Control.Applicative
import Data.Function (on)

import qualified Data.Array as A
import qualified Data.List  as L

type Point = (Int, Int)
type Size  = Int
type Total = Int

gridSerialNum = 6548
main = do
   let ((p1X, p1Y), (p1Sz, p1Total)) = solution [3]
   let ((p2X, p2Y), (p2Sz, p2Total)) = solution [1..300]
   putStrLn $ "Part1: " ++ show (p1X, p1Y) ++ " Total: " ++ show p1Total
   putStrLn $ "Part2: " ++ show (p2X, p2Y, p2Sz) ++ " Total: " ++ show p2Total

fnArray bounds f = A.listArray bounds [ f ix | ix <- A.range bounds ]

powerLevels = fnArray ((1,1),(300,300)) $ \(x,y) ->
   let rackID = x + 10
       hundreds n = (n `div` 100) `mod` 10
   in hundreds ((rackID * y + gridSerialNum) * rackID) - 5

totals =
   let lookup sz pt = snd (totals sz A.! pt)
       quarter sz (x,y) = sum [ lookup sz (x', y') | x' <- [x,x+sz]
                                                   , y' <- [y,y+sz] ]
       totalsArr = fnArray (1,300) $ \sz ->
         fnArray ((1,1),(301-sz,301-sz)) $ \(x,y) ->
            if sz == 1 then (sz, powerLevels A.! (x, y)) else
            if even sz then (sz, quarter (sz `div` 2) (x, y)) else
            (sz, lookup (sz-1) (x, y)
               + sum [ powerLevels A.! (x', y+sz-1) | x' <- [x..x+sz-2] ]
               + sum [ powerLevels A.! (x+sz-1, y') | y' <- [y..y+sz-1] ])
   in \sz -> totalsArr A.! sz

solution :: [Size] -> (Point, (Size, Total))
solution = L.maximumBy (compare `on` snd . snd)
         . L.concatMap (A.assocs . totals)