{-# OPTIONS_GHC -fexcess-precision #-}

-- Find the sum of all progressive perfect squares below one trillion

import qualified Data.Map as M
import List
import Data.Int
import Monad
import System


primes :: [Int64]
primes = 2:nextPrime 3 (tail $ inits primes)
    where
        nextPrime n ps =
            let next = head $ dropWhile (\n -> not $ all (\p -> n `mod` p /= 0) (head ps)) [n,n+2..]
            in  next:nextPrime (next + 2) (tail ps)

bins = M.fromListWith (++) (map setup (takeWhile (< 9999) . dropWhile (< 1000) $ primes))
    where
        setup n = (sort $ show n, [n])

ls = filter ((>= 4).length) $ map snd $ M.toList bins

pick :: Int64 -> [a] -> [[a]]
pick 0 _ = return []
pick n [] = mzero
pick n (h:t) = takeIt `mplus` skipIt
    where
        takeIt = do
            ls <- pick (n - 1) t
            return (h:ls)
        skipIt = pick n t

permute [] = return []
permute [e] = return [e]
permute l = msum (map (takeFrom l) l)
    where
        takeFrom l n = fmap (n:) $ permute (delete n l)

progressive :: Int64 -> Bool
progressive n = n `seq` maybeTrace $ go 1 (ceiling $ (sqrt (fromIntegral n) :: Double))
    where
        maybeTrace = id

        go m c
            | m `seq` c `seq` False = undefined

            | m >= c
            = False
            | infos m
            = True
            | otherwise
            = go (m + 1) c

        infos :: Int64 -> Bool
        infos d = d `seq`
            let (q, r) = n `quotRem` d
            in  progprop d q r

        progprop :: Int64 -> Int64 -> Int64 -> Bool
        progprop d q r = tuplesort d q r
            where
                tuplesort d q r
                    | d `seq` q `seq` r `seq` False = undefined
                    | q > d
                    = geom r d q
                    | q > r
                    = geom r q d
                    | otherwise
                    = geom q r d
                geom :: Int64 -> Int64 -> Int64 -> Bool
                geom a b c 
                    | a `seq` b `seq` c `seq` False = undefined
                    | a == 0
                    = False
                    | b == 0
                    = False
                    | otherwise
                    = b * b == a * c

tuplegeom = map (\n -> [1,n,n*n]) [1..]

perml l = concatMap permute l

genval [d,q,r] = q*d+r

traceList [] = []
traceList l@(h:_) = before:traceList after
    where (before, after) = splitAt 1000 l

main = do
    [minR, maxR] <- getArgs
    let (min, max) = (read minR, read maxR) :: (Int64, Int64)
        l = concatMap (concatMap (\n -> if progressive (n*n) then [n*n] else []))
                      (traceList [min..max])
    mapM_ (\n -> putStrLn $ "Found progressive: n^2 = " ++ show n) l

-- n = q * d + r
-- 
-- n = (a * b)*(a * b * b) + a
-- n = (a * b * b) * a + a * b
-- n = (a * b) * a + a * b * b

proggen a b = [a*a * b*b*b + a, a*a * b*b + a*b]
--proggen a b = [a*b*a*b*b + a, a*b*b*a + a * b, a*b*a + a*b*b]