{-# LANGUAGE ScopedTypeVariables, LambdaCase #-}
{-

Expressing numbers as a sum of 4 square numbers (Lagrange's
four-square theorem).

Copyright 2018 Ken Takusagawa

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or (at
your option) any later version.

This program is distributed in the hope that it will be useful, but
WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

-}
module Main where {
import System.Environment(getArgs);
import Control.Exception(assert);
import Debug.Trace(trace);
import Data.Function((&));
import Control.Category((>>>));
import Prelude hiding((.),(>>));
--import System.IO(hSetBuffering,stdout,BufferMode(LineBuffering));
import Data.List;
--import Control.Monad;
--import Data.Maybe;
--import qualified Data.Map as Map;
--import Data.Map(Map);
--import qualified Data.Set as Set;
--import Data.Set(Set);
import Control.Arrow(second);

-- to avoid the redundancy warning
trace_placeholder :: ();
trace_placeholder = (trace,assert) & (id >>> error) "trace_placeholder";

choose_n_of_max :: Integer -> Int -> [[Integer]];
choose_n_of_max amax num = case compare num 0 of {
LT -> error "negative choose_n_of_max";
EQ -> return [];
GT -> do {
        x <- [0..amax];
        y <- choose_n_of_max x (pred num);
        return (x:y);
}};

sumsq :: [Integer] -> Integer;
sumsq = map (\x -> x*x) >>> sum;

get_sumsq :: [Integer] -> (Integer,[Integer]);
get_sumsq x = (sumsq x,x);

ordgroup :: [(Integer,a)] -> [[(Integer,a)]];
ordgroup = sortOn fst >>> groupBy (\x y -> (fst x) == (fst y));

-- this is a mapreduce operation, though we do not do it

clean_heads :: [(a,b)] -> (a,[b]);
clean_heads l = (head l & fst, map snd l);

main :: IO();
main = getArgs >>= \case{
[n] -> do_s (read n) & mapM_ print;
["latex",n] -> do_s (read n) & map latex_out & latex_join & putStrLn;
["latexmin",n] -> do_s (read n) & tail & map (second minimize_numterms) & map latex_out & latex_join & putStrLn;
_ -> undefined
};

do_s :: Integer -> [(Integer,[[Integer]])];
do_s n = choose_n_of_max n 4 & map get_sumsq & filter (\x -> fst x <= n*n) & ordgroup & map clean_heads ;

latex_join :: [String] -> String;
latex_join = intersperse " \\\\\n" >>> concat;

latex_out :: (Integer, [[Integer]]) -> String;
latex_out (s,xs) = (show s) ++ (map a xs & latex_join) where {
a :: [Integer] -> String;
a l = " &=& " ++ (map b l & intersperse " + " & concat);
b :: Integer -> String;
b x = (show x) ++ "^2";
};

-- 90 can be expressed in 9 ways, 130 and 138 can be expressed in 10

-- consider going up to 112=4*4*7 to see Legendre's three-square theorem in action

num_nonzero :: [Integer] -> Int;
num_nonzero = filter (\x -> x /= 0) >>> length;

minimize_numterms :: [[Integer]] -> [[Integer]];
minimize_numterms ls = let {
 smallest :: Int;
 smallest = map num_nonzero ls & minimum
} in filter (\l -> smallest == num_nonzero l) ls & map (filter (\l -> l/=0)) ;
} --end