module Problem142 where

import Data.Array
import qualified Data.IntMap as Map
import Data.List

lim :: Int
lim = 950

main :: IO ()
-- x+y=a1^2 | x-y=a2^2 | y-z=a3^2 | y+z=a4^2 | x-z=a5^2 | x+z=a6^2
-- we say that y is related to x if x+y and x-y both are perfect squares and x>y
-- now we form list of related numbers and a map. we find entries (x,y) (y,z) (x,z)
main =
    print
        . head
        $ [ x + y + z
          | x <- Map.keys relatedMap
          , length (relatedMap Map.! x) > 1
          , y <- sort $ relatedMap Map.! x
          , y `Map.member` relatedMap
          , z <- relatedMap Map.! y
          , z `elem` (relatedMap Map.! x)
          ]
  where
    ps =
        array (1, lim) ((1, 1) : [ (i, ps ! (i - 1) + 2 * i - 1) | i <- [2 .. lim] ])
    relatedMap = Map.fromListWith
        (++)
        [ (x, [y])
        | i <- [1 .. lim]
        , j <- [i + 2, i + 4 .. lim]
        , let x = (ps ! j + ps ! i) `div` 2
        , let y = (ps ! j - ps ! i) `div` 2
        ]
