module Handler.Utils.SubmissionSpec where

import qualified Yesod

import TestImport

import Handler.Utils.Submission
import ModelSpec ()

import qualified Data.Set as Set
import qualified Data.Map as Map

import Data.List (genericLength)

import qualified Data.CaseInsensitive as CI

import System.IO.Unsafe

import System.Random.Shuffle
import Control.Monad.Random.Class

import Database.Persist.Sql (fromSqlKey)


userNumber :: TVar Natural
userNumber = unsafePerformIO $ newTVarIO 1
{-# NOINLINE userNumber #-}

makeUsers :: Natural -> SqlPersistM [Entity User]
makeUsers (fromIntegral -> n) = do
  users' <- liftIO . replicateM n $ generate arbitrary
  users <- forM users' $ \u -> do
    i <- atomically $ readTVar userNumber <* modifyTVar userNumber succ
    let u' = u { userIdent = CI.mk $ "user." <> tshow i
               , userEmail = CI.mk $ "user." <> tshow i <> "@example.com"
               }
    return u'
  uids <- insertMany users
  return $ zipWith Entity uids users

distributionExample :: SqlPersistM [(Natural, [Maybe Load])] -- ^ Number of submissions and corrector specification
                    -> ([Entity Submission] -> [Entity SheetCorrector] -> SqlPersistM ()) -- ^ Setup hook
                    -> (Map (Maybe SheetCorrector) (Set SubmissionId) -> Expectation)
                    -> YesodExample UniWorX ()
distributionExample mkParameters setupHook cont = do
  situations <- runDB $ do
    term <- liftIO $ generate arbitrary
    void . insert $ term
    school <- liftIO $ generate arbitrary
    void . insert $ school
    course <- liftIO $ generate arbitrary <&> \c -> c { courseTerm = TermKey $ termName term, courseSchool = SchoolKey $ schoolShorthand school }
    cid <- insert course

    steps <- mkParameters
    let subsN = maybe 0 maximum . fromNullable $ map fst steps
        correctorsN = maybe 0 maximum . fromNullable $ map (genericLength . snd) steps
    participants <- makeUsers subsN
    correctors <- makeUsers correctorsN

    situations <- forM steps $ \(subsN', loads) -> do
      sheet <- liftIO $ generate arbitrary <&> \s -> s { sheetCourse = cid }
      sid <- insert sheet

      participants' <- liftIO $ take (fromIntegral subsN') <$> shuffleM participants
      let loads' = loads ++ replicate (fromIntegral $ correctorsN - genericLength loads) Nothing
      
      submissions <- forM participants' $ \(Entity uid _) -> do
        sub@(Entity subId _) <- insertEntity $ Submission
          sid
          Nothing
          Nothing
          Nothing
          Nothing
          Nothing
        void . insert $ SubmissionUser uid subId
        return sub

      let sheetCorrectors = [ SheetCorrector corr sid load CorrectorNormal | (Entity corr _, Just load) <- zip correctors loads']
      scIds <- insertMany sheetCorrectors
      let sheetCorrectors' = zipWith Entity scIds sheetCorrectors

      return (sid, (submissions, sheetCorrectors'))

    mapM_ (uncurry setupHook) $ map snd situations

    return situations

  let subIds = concatMap (\(_, (subs, _)) -> map entityKey subs) situations

  results <- runHandler . Yesod.runDB . mapM (\sid -> assignSubmissions sid Nothing) $ map fst situations

  submissions <- fmap concat . forM results $ \(assigned, unassigned) -> runDB $ selectList ([ SubmissionId <-. Set.toList assigned ] ||. [ SubmissionId <-. Set.toList unassigned ]) []

  liftIO $ do
    let (assigned, unassigned) = bimap concat concat $ unzip results
    Set.union assigned unassigned `shouldBe` Set.fromList subIds
    cont . Map.fromListWith mappend $ do
      Entity subId Submission{..} <- submissions
      let key = listToMaybe . filter (\(Entity _ (SheetCorrector uid _ _ _)) -> Just uid == submissionRatingBy) $ concatMap (\(_, (_, corrs)) -> corrs) situations
      return (entityVal <$> key, Set.singleton subId)
  

spec :: Spec
spec = withApp . describe "Submission distribution" $ do
  it "is fair" $
    distributionExample
      (return [(500, replicate 10 (Just $ Load Nothing 1))])
      (\_ _ -> return ())
      (\result -> do
          let countResult = Map.map Set.size result
          countResult `shouldNotSatisfy` Map.member Nothing
          countResult `shouldSatisfy` all (== 50)
      )
  it "follows distribution" $
    distributionExample
      (return [(500, replicate 6 (Just $ Load Nothing 1) ++ replicate 2 (Just $ Load Nothing 2))])
      (\_ _ -> return ())
      (\result -> do
          let countResult = Map.map Set.size result
          countResult `shouldNotSatisfy` Map.member Nothing
          countResult `shouldSatisfy` all (\(Just (SheetCorrector _ _ (Load _ prop) _), subsSet) -> (== 50 * prop) $ fromIntegral subsSet) . Map.toList
      )
  it "follows cumulative distribution over multiple sheets" $ do
    ns <- liftIO . replicateM 5 . fmap fromInteger $ getRandomR (0, 100)
    let ns' = ns ++ [500 - sum ns]
        loads = replicate 6 (Just $ Load Nothing 1) ++ replicate 2 (Just $ Load Nothing 2)
    distributionExample
      (return [ (n, loads) | n <- ns' ])
      (\_ _ -> return ())
      (\result -> do
          let countResult = Map.map Set.size result
              countResult' = Map.mapKeysWith (+) (fmap $ \SheetCorrector{..} -> (fromSqlKey sheetCorrectorUser, byProportion sheetCorrectorLoad)) countResult
          countResult `shouldNotSatisfy` Map.member Nothing
          countResult' `shouldSatisfy` all (\(Just (_, prop), subsSet) -> (== 50 * prop) $ fromIntegral subsSet) . Map.toList
      )