module Exercise05 where

import Test.QuickCheck (quickCheck, (==>))
import Data.List
import Data.Bits
-- ~ import Data.List.NonEmpty (groupAllWith)
-- ~ import Debug.Trace

-- May or may not be useful: computes the logarithm base 2 (rounded down) of the given number.
-- You don't have to move this into the WETT tags if you want to use it.
log2 :: (Integral a, Num b) => a -> b
log2 = let go acc n = if n <= 1 then acc else go (acc + 1) (n `div` 2) in go 0

{-WETT-}

-- ~ decompose :: [Integer] -> [Integer]
-- ~ decompose ds = decomposeWithBases (pO2sr $ fromIntegral $ length ds) ds

decomposeWithBases :: [Integer] -> [Integer] -> [Integer]
decomposeWithBases [d] _ = decomposeDim d
decomposeWithBases (d:ds) (base:more_bases) = zipWith (*) old bas
    where
        one = decomposeDim d
        old = decomposeWithBases more_bases ds
        bas = backAdd base one

-- ~ -- new_dim -> reversed_pre -> acc -> new
-- ~ bla :: Integer -> [Integer] -> [Integer] -> [Integer]
-- ~ bla d -> []


decompose :: [Integer] -> [Integer]
decompose = fst . decomposeWithDepth

decomposeWithDepth :: [Integer] -> ([Integer], Integer)
decomposeWithDepth [d] = (decomposeDim d, 1)
decomposeWithDepth (d:ds) = (
        removeTrailing0s $
        map (sum . map fst) $
        -- ~ map fst $
        -- ~ groupAllWith snd $
        groupBy (\(_,a) (_,b) -> a == b) $
        sortOn snd $
        add0s $
        -- ~ [if sn < so then hfun cn sn (co * (so `div` sn)^depth) else hfun co so (cn * (sn `div` so)) | (cn, sn) <- (zip new $ pO2s 1), (co, so) <- (zip old $ pO2s 1)]
        -- ~ [hfun2 cn sn co so | (cn, sn) <- (zip new $ pO2s 1), cn /= 0, (co, so) <- (zip old $ pO2s 1), co /= 0]
        [hfun2 cn sn co so | (cn, sn) <- (zip new $ [1..]), cn /= 0, (co, so) <- (zip old $ [1..]), co /= 0]
        , depth + 1)
    where
        new = decomposeDim d
        (old, depth) = decomposeWithDepth ds
        -- count_small -> size_small -> count_long -> size_long -> (num, size)
        hfun :: Integer -> Integer -> Integer -> (Integer, Integer)
        -- ~ hfun cs ss cl sl = trace (show cs ++ "," ++ show cl ++ "," ++ show ss ++ "," ++ show sl ++ "," ++ show (cs * cl * small_per_large, ss)) (cs * cl * small_per_large, ss)
        hfun cs ss cl = (cs * cl, ss)
        -- ~ hfun co so (cn * (sn `div` so)) = (co * (cn * (sn `div` so)), so)
        -- ~ hfun cn sn (co * (so `div` sn)^depth) = (cn * (co * (so `div` sn)^depth), sn)
            -- ~ where small_per_large = (sl `div` ss)
        -- ~ hfun2 cn sn co so = if sn < so then ((co * (so `div` sn)^depth), sn) else (co * ((sn `div` so)), so)
        hfun2 cn sn co so = if sn < so then hfun20 else hfun21
            where
                -- ~ hfun20 = ((co * (power (so `div` sn) depth)), sn)
                -- ~ hfun20 = (co * (power 2 ((so - sn) * depth)), sn)
                hfun20 = tup -- so and sn grow logramithlically
                    where
                        tup = (mult2, sn)
                        mult2 = shiftL co $ alksjf
                        mult = co * basdfkbjh
                        basdfkbjh = ($!) identity asdasa
                            where identity a = a
                        alksjf = ($!) alksjff ()
                            where alksjff () = fromIntegral ((so - sn) * depth)
                        asdasa = ($!) asdasaf alksjf
                            where asdasaf alksjf = (shiftL 1 $ alksjf)
                -- ~ hfun21 = (co * ((sn `div` so)), so)
                hfun21 = (co * (shiftL 1 $ fromIntegral  (sn - so)), so)
        add0s :: [(Integer, Integer)] -> [(Integer, Integer)]
        -- ~ add0s xs = [(0, s) | s <- (takeWhile (<max_s) $ pO2s 1)] ++ xs
        add0s xs = [(0, s) | s <- (takeWhile (<max_s) $ [1..])] ++ xs
            where
                max_s = maximum $ (1:map snd xs)
        removeTrailing0s :: [Integer] -> [Integer]
        removeTrailing0s [] = []
        removeTrailing0s (x:xs)
            | x == 0 && null more = []
            | otherwise = x:more
            where more = removeTrailing0s xs

power = (^)

-- in: base
-- in: (a_0, ..., a_n)
-- out: (b_0, ..., b_n)
-- b_j = sum_{i=j}^n {a_i * base^(i-j)}
backAdd :: Integer -> [Integer] -> [Integer]
backAdd _ [a] = [a]
backAdd base (a:as) = (a + base * ba1) : bas
    where bas@(ba1:_) = backAdd base as

-- finds fewest cubes for 1 dimension
-- returned list only contains 1s and 0s
decomposeDim :: Integer -> [Integer]
decomposeDim 0 = []
decomposeDim n = m : decomposeDim d
    where
        -- ~ (d, m) = n `divMod` 2
        m = n .&. 1
        d = shiftR n 1

-- like decomposeDim, but returns False for 0 and True for 1
decomposeDimB :: Integer -> [Bool]
decomposeDimB 0 = []
decomposeDimB n = (m==1) : decomposeDimB d
    where
        (d, m) = n `divMod` 2

pO2s :: Integer -> [Integer]
pO2s n = n : pO2s (n*2)

-- 3 -> [4, 2, 1]
pO2sr :: Integer -> [Integer]
pO2sr 1 = [1]
pO2sr l = r*2 : rs
    where rs@(r:_) = pO2sr (l-1)

{-TTEW-}

prop_backAdd_len base as = not (null as) ==> length as == length (backAdd base as)

prop_backAdd base as = not (null as) ==>
        and [(bas !! j) == sum [(as !! i) * base^(i-j) | i <- [j..n]] | j <- [0..n]]
    where
        bas = backAdd base as
        n = length as - 1

prop_pO2sr l = l > 0 ==> and $ zipWith f xs $ tail xs
    where
        xs = pO2sr l
        f a b = a == 2 * b

prop_decompose_keepsVolume xs = (not $ null xs) && all (>=0) xs ==> sizeb == sizea
    where
        sizeb = product xs
        sizea = sum $ zipWith (\a b -> a * b^dim) decom (pO2s 1)
        dim = length xs
        decom = decompose xs

prop_decomposeDim_only10 a = and [b == 1 || b == 0 | b <- decomposeDim a]