module MLLabs.Validation
      ( TrainAndTestSets
       , ValidationStrategy
       , kfold
       ) where

import           MLLabs.Types

type TrainAndTestSets obj lbl = (Dataset obj lbl, Dataset obj lbl)
type ValidationStrategy obj lbl = Dataset obj lbl -> [TrainAndTestSets obj lbl]

-- kfold 3 [1..6] == [([3,4,5,6],[1,2]),([1,2,5,6],[3,4]),([1,2,3,4],[5,6])]
kfold :: Int -> ValidationStrategy obj lbl
kfold k dat = parts
  where
    n = length dat

    -- size of the first i parts
    size :: Int -> Int
    size i | i < 0 = error "invalid size in kfold"
           | i > k = error "invalid size in kfold"
           | i < q = i * (p+1)
           | True  = q * (p+1) + (i - q) * p
           where (p,q) = n `divMod` k

    splitsAt :: [Int] -> [a] -> [[a]]
    splitsAt [] xs = [xs]
    splitsAt (i:is) xs =
      let (as, bs) = splitAt i xs
      in as : splitsAt is bs

    -- trainᵢ = d[0,sᵢ) ++ d[sᵢ₊₁,n)
    -- testᵢ  = d[sᵢ,sᵢ₊₁)
    parts = flip map [0..k-1] $ \i ->
      let [tr1, te, tr2] = splitsAt [size i, size (i+1) - size i] dat
      in (tr1 ++ tr2, te)
