qInsertAt :: [a] -> Int -> a -> Maybe [a]
qInsertAt lst k e = if k < 1 then Nothing else qInsertAtI lst k e []
  where
    qInsertAtI :: [a] -> Int -> a -> [a] -> Maybe [a]
    qInsertAtI xs 1 x acc = Just (qReverse acc ++ (x : xs))
    qInsertAtI [] _ _ _ = Nothing
    qInsertAtI (h : t) n x acc = qInsertAtI t (n - 1) x (h : acc)

qRange :: Int -> Int -> [Int]
qRange a b = [a .. b]

qRandomSelect :: [a] -> Int -> IO (Maybe [a])
qRandomSelect lst n = qRandomSelectI lst n []
  where
    qRandomSelectI :: [a] -> Int -> [a] -> IO (Maybe [a])
    qRandomSelectI (h : _) 1 acc = return (Just (h : acc))
    qRandomSelectI [] _ _ = return Nothing
    qRandomSelectI l k acc
      | k < 1 = return Nothing
      | otherwise = do
          idx <- randomRIO (1, qLength l)
          case qRemoveAt l idx of
            (Just e, rest) -> qRandomSelectI rest (k - 1) (e : acc)
            (Nothing, _) -> return Nothing

qLotto :: Int -> Int -> IO (Maybe [Int])
qLotto n m = qRandomSelect [1 .. m] n

qRandomPermutation :: [a] -> IO (Maybe [a])
qRandomPermutation lst = qRandomSelect lst (qLength lst)

qCombination :: [a] -> Int -> Maybe [[a]]
qCombination lst k
  | k < 1 || k > qLength lst = Nothing
  | otherwise = Just (qCombinationI lst k [] [])
  where
    qCombinationI :: [a] -> Int -> [a] -> [[a]] -> [[a]]
    qCombinationI l m c acc = case m of
      0 -> c : acc
      _ -> case l of
        [] -> acc
        h : t ->
          qCombinationI t m c (qCombinationI t (m - 1) (h : c) acc)

qGroup :: [a] -> [Int] -> Maybe [[[a]]]
qGroup lst g
  | sum g > qLength lst || any (< 1) g = Nothing
  | otherwise = mapM (`partList` g) (permu lst)
  where
    partList :: [a] -> [Int] -> Maybe [[a]]
    partList xs p
      | null xs || null p || sum p > qLength xs || any (< 1) p = Nothing
      | otherwise = Just (partListI xs g [] [])
      where
        partListI :: [a] -> [Int] -> [a] -> [[a]] -> [[a]]
        partListI [] _ _ acc = qReverse acc
        partListI _ [] _ acc = qReverse acc
        partListI (lh : lt) (ph : pt) c acc = case ph of
          0 -> partListI (lh : lt) pt [] (qReverse c : acc)
          _ -> partListI lt ((ph - 1) : pt) (lh : c) acc
    permu :: [a] -> [[a]]
    permu xs = permuI xs [] [] []
      where
        permuI :: [a] -> [a] -> [a] -> [[a]] -> [[a]]
        permuI [] [] select acc = select : acc
        permuI [] _ _ acc = acc
        permuI (h : t) cache select acc =
          permuI t (h : cache) select (permuI (t ++ cache) [] (h : select) acc)

qLengthSort :: (Ord a) => [[a]] -> [[a]]
qLengthSort lst =
  Data.List.foldl' (++) [] $
    map
      (map (\(_, _, e) -> e) . Data.List.sortBy aux3')
      ( Data.List.groupBy (\(lenA, _, _) (lenB, _, _) -> lenA == lenB) $
          Data.List.sortBy aux3 $
            zipWith (\idx l -> (qLength l, idx, l)) [1 :: Int ..] lst
      )
  where
    aux3 :: (Ord a) => (a, b, c) -> (a, b, c) -> Ordering
    aux3 (a, _, _) (b, _, _)
      | a > b = GT
      | a < b = LT
      | otherwise = EQ
    aux3' :: (Ord b) => (a, b, c) -> (a, b, c) -> Ordering
    aux3' (_, a, _) (_, b, _)
      | a > b = GT
      | a < b = LT
      | otherwise = EQ

qLengthFrequencySort :: [[a]] -> [[a]]
qLengthFrequencySort lst =
  Data.List.foldl' (++) [] $
    map
      (((map (\(_, _, e) -> e) . Data.List.sortBy aux3') . Data.List.foldl' (++) []) . map snd)
      ( Data.List.groupBy (\(lenA, _) (lenB, _) -> lenA == lenB) $
          Data.List.sortBy aux2 $
            map (\g -> (qLength g, g)) $
              Data.List.groupBy (\(lenA, _, _) (lenB, _, _) -> lenA == lenB) $
                Data.List.sortBy aux3 $
                  zipWith (\idx l -> (qLength l, idx, l)) [1 :: Int ..] lst
      )
  where
    aux3 :: (Ord a) => (a, b, c) -> (a, b, c) -> Ordering
    aux3 (a, _, _) (b, _, _)
      | a > b = GT
      | a < b = LT
      | otherwise = EQ
    aux3' :: (Ord b) => (a, b, c) -> (a, b, c) -> Ordering
    aux3' (_, a, _) (_, b, _)
      | a > b = GT
      | a < b = LT
      | otherwise = EQ
    aux2 :: (Ord a) => (a, b) -> (a, b) -> Ordering
    aux2 (a, _) (b, _)
      | a > b = GT
      | a < b = LT
      | otherwise = EQ

module Part03 (part03) where

import Qarks.Nnp
  ( qCombination,
    qGroup,
    qInsertAt,
    qLotto,
    qRandomPermutation,
    qRandomSelect,
    qRange,
  )
import Test.Hspec
  ( context,
    describe,
    it,
    shouldBe,
    shouldSatisfy,
  )
import Test.Hspec.Runner (SpecWith)

maybeAllElem :: (Eq a) => [a] -> Maybe [a] -> Bool
maybeAllElem xs lst = case lst of
  Nothing -> False
  Just l -> all (`elem` xs) l

maybeLength :: Maybe [a] -> Int
maybeLength Nothing = 0
maybeLength (Just lst) = length lst

part03 :: SpecWith ()
part03 = describe "Part03" $ do
  context "Qarks.Nnp.qInsertAt" $ do
    it "insert 'X' to \"abcd\" at 2" $ do
      qInsertAt "abcd" 2 'X' `shouldBe` Just "aXbcd"
  context "Qarks.Nnp.qRange" $ do
    it "range from 4 to 9" $ do
      qRange 4 9 `shouldBe` [4 .. 9]
  context "Qarks.Nnp.qRandomSelect" $ do
    it "randomly select 3 from ['a' .. 'h']" $ do
      res <- qRandomSelect ['a' .. 'h'] 3
      res
        `shouldSatisfy` \lst ->
          maybeAllElem ['a' .. 'h'] lst
            && maybeLength lst == 3
  context "Qarks.Nnp.qLotto" $ do
    it "randomly select 6 from [1 .. 49]" $
      do
        res <- qLotto 6 49
        res
          `shouldSatisfy` \lst ->
            maybeAllElem [1 .. 49] lst
              && maybeLength lst == 6
  context "Qarks.Nnp.qRandomPermutation" $ do
    it "random permutation of ['a' .. 'f']" $ do
      res <- qRandomPermutation ['a' .. 'f']
      res
        `shouldSatisfy` \lst ->
          maybeAllElem ['a' .. 'f'] lst
            && maybeLength lst == length ['a' .. 'f']
  context "Qarks.Nnp.qCombination" $ do
    let res = qCombination [1 .. 4 :: Int] 2
    it "combinan of [1 .. 4] with 2 :: length" $ do
      length <$> res `shouldBe` Just 6
    it "combinan of [1 .. 4] with 2 :: sublength" $ do
      all (\l -> length l == 2) <$> res `shouldBe` Just True
    it "combinan of [1 .. 4] with 2 :: subelem" $ do
      all (all (\e -> e `elem` [1 .. 4])) <$> res
        `shouldBe` Just True
  context "Qarks.Nnp.qGroup" $ do
    let res = qGroup [1 .. 3 :: Int] [1, 1]
    it "group [1 .. 3] by [1, 1] :: length" $ do
      length <$> res `shouldBe` Just 6
    it "group [1 .. 3] by [1, 1] :: sublength" $ do
      all (\l -> length l == 2) <$> res `shouldBe` Just True
    it "group [1 .. 3] by [1, 1] :: subelem" $ do
      all (all (\e -> length e == 1 && head e `elem` [1 .. 4])) <$> res
        `shouldBe` Just True
  context "Qarks.Nnp.qLengthSort" $ do
    it "sort [\"abc\", \"de\", \"fgh\", \"de\", \"ijkl\", \"mn\", \"o\"] by length" $ do
      qLengthSort ["abc", "de", "fgh", "de", "ijkl", "mn", "o"]
        `shouldBe` ["o", "de", "de", "mn", "abc", "fgh", "ijkl"]
  context "Qarks.Nnp.qLengthFrequencySort" $ do
    it "sort [\"abc\", \"de\", \"fgh\", \"de\", \"ijkl\", \"mn\", \"o\"] by frequency" $ do
      qLengthFrequencySort ["abc", "de", "fgh", "de", "ijkl", "mn", "o"]
        `shouldBe` ["ijkl", "o", "abc", "fgh", "de", "de", "mn"]