-- SPDX-FileCopyrightText: 2022 Gregor Kleen <gregor.kleen@ifi.lmu.de>,Sarah Vaupel <vaupel.sarah@campus.lmu.de>,Steffen Jost <jost@tcs.ifi.lmu.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-}

module Handler.Utils.SubmissionSpec where

import qualified Yesod

import TestImport
-- import qualified Test.HUnit.Base as HUnit

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)

import qualified Database.Esqueleto.Legacy as E

import Data.Monoid (First(..))

import Utils (guardOn)

import Control.Lens.Extras (is)

import Data.Maybe (fromJust)


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
                    -> (Natural -> [Entity Submission] -> [Entity SheetCorrector] -> SqlPersistM ()) -- ^ Setup hook
                    -> (Map (Maybe SheetCorrector) (Set (SubmissionId, Maybe Natural)) -> 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 (zip [1..] steps) $ \(i, (subsN', loads)) -> do
      sheet <- liftIO $ generate arbitrary <&> \s -> s { sheetName = CI.mk $ "Sheet " <> tshow (i :: Integer), 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_ (\(n, (subs, corrs)) -> setupHook n subs corrs) . zip [1..] $ map snd situations

    situations' <- forM situations $ \(sid, (submissions, sheetCorrectors)) -> (sid, ) <$> do
      submissions' <- mapM (fmap fromJust . getEntity . entityKey) submissions
      sheetCorrectors' <- mapM (fmap fromJust . getEntity . entityKey) sheetCorrectors
      return (submissions', sheetCorrectors')

    return situations'

  let
    subIds :: [SubmissionId]
    subIds = concatMap (\(_, (subs, _)) -> mapMaybe (\(Entity subId Submission{..}) -> guardOn (is _Nothing submissionRatingBy) subId) 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
          sheet = getFirst . foldMap (\(n, (sid, _)) -> First $ guardOn (sid == submissionSheet) n) $ zip [1..] situations
      return (entityVal <$> key, Set.singleton (subId, sheet))
  

spec :: Spec
spec = withApp . describe "Submission distribution" $ do
  it "is fair" $
    distributionExample
      (return [(500, replicate 10 (Just $ Load Nothing 1 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 1) ++ replicate 2 (Just $ Load Nothing 2 1))])
      (\_ _ _ -> 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 1) ++ replicate 2 (Just $ Load Nothing 2 1)
    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
      )
  it "follows non-constant cumulative distribution over multiple sheets" $ do
    let ns = replicate 4 100
        loads = do
          (onesBefore, onesAfter) <- zip [0,2..6] [6,4..0]
          return $ replicate onesBefore (Just $ Load Nothing 1 1)
                ++ replicate 2          (Just $ Load Nothing 2 1)
                ++ replicate onesAfter  (Just $ Load Nothing 1 1)
    distributionExample
      (return $ zip ns loads)
      (\_ _ _ -> return ())
      (\result -> do
          let countResult = Map.map Set.size result
              countResult' = Map.mapKeysWith (+) (fmap $ \SheetCorrector{..} -> fromSqlKey sheetCorrectorUser) countResult
          countResult `shouldNotSatisfy` Map.member Nothing
          countResult' `shouldSatisfy` all (\(Just _, subsSet) -> subsSet == 50) . Map.toList
      )
  it "handles tutorials with proportion" $ do
    ns <- liftIO . replicateM 5 . fmap fromInteger $ getRandomR (0, 100)
    let ns' = ns ++ [500 - sum ns]
        loads = replicate 6 (Just $ Load (Just True) 1 1) ++ replicate 2 (Just $ Load (Just True) 2 1)
    tutSubIds <- liftIO $ newTVarIO Map.empty
    distributionExample
      (return [ (n, loads) | n <- ns' ])
      (\_ subs corrs -> do
          tutSubmissions <- liftIO $ getRandomR (5,10)
          subs' <- liftIO $ shuffleM subs
          forM_ (take tutSubmissions subs') $ \(Entity subId Submission{..}) -> do
            Entity _ SheetCorrector{..} <- liftIO $ uniform corrs
            atomically . modifyTVar tutSubIds . Map.insertWith mappend sheetCorrectorUser $ Set.singleton subId
            Sheet{..} <- getJust submissionSheet
            tut <- liftIO $ generate arbitrary <&> \c -> c { tutorialName = CI.mk $ "Tut for " <> tshow (fromSqlKey subId), tutorialCourse = sheetCourse }
            tutId <- insert tut
            void . insert $ Tutor tutId sheetCorrectorUser
            E.insertSelect . E.from $ \submissionUser -> do
              E.where_ $ submissionUser E.^. SubmissionUserSubmission E.==. E.val subId
              return $ TutorialParticipant E.<# E.val tutId E.<&> (submissionUser E.^. SubmissionUserUser)
      )
      (\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) -> fromIntegral subsSet == 50 * prop) . Map.toList
          
          -- -- Does not currently work, because `User`s are reused within `distributionExample`, so submissions end up having more associated course-tutors, because the same user might be a member of a tutorial created for another submission
          -- 
          -- let subs = fold tutSubIds'
          -- forM_ subs $ \subId -> do
          --   let tutors = Map.keysSet $ Map.filter (Set.member subId) tutSubIds'
          --       assignedTo = Set.map (sheetCorrectorUser . fromJust) . Map.keysSet $ Map.filter (Set.member subId) result
          --   HUnit.assertEqual ("Submission " <> show (fromSqlKey subId) <> " assigned to multiple correctors") 1 $ Set.size assignedTo
          --   HUnit.assertEqual ("Submission " <> show (fromSqlKey subId) <> " assigned to non-tutors (" <> show (Set.map fromSqlKey tutors) <> ")") Set.empty (Set.map fromSqlKey $ assignedTo `Set.difference` tutors)
      )
  it "allows disabling deficit consideration" $
    distributionExample
      (return . replicate 2 $ (500, replicate 2 (Just $ Load Nothing 1 0)))
      (\n subs corrs -> if
          | n < 2
          , Entity _ SheetCorrector{ sheetCorrectorUser = corrId } : _ <- corrs
            -> forM_ subs $ \(Entity subId _) ->
                 update subId [SubmissionRatingBy =. Just corrId]
          | otherwise -> return ()
      )
      (\result -> do
          let secondResult = Map.map (Set.size . Set.filter (views _2 (== Just 2))) result
              allEqual [] = True
              allEqual ((_, c) : xs) = all (\(_, c') -> c == c') xs
          secondResult `shouldSatisfy` allEqual . Map.toList
      )
  it "allows disabling deficit consideration with unequal proportions" $
    distributionExample
      (return . replicate 2 $ (550, [Just (Load Nothing 1 0), Just (Load Nothing 10 0)]))
      (\n subs corrs -> if
          | n < 2
          , Entity _ SheetCorrector{ sheetCorrectorUser = corrId } : _ <- corrs
            -> forM_ subs $ \(Entity subId _) ->
                 update subId [SubmissionRatingBy =. Just corrId]
          | otherwise -> return ()
      )
      (\result -> do
          let secondResult = Map.map (Set.size . Set.filter (views _2 (== Just 2))) result
              secondResultNorm = imap go secondResult
                where go Nothing x = fromIntegral x
                      go (Just SheetCorrector{..}) x = fromIntegral x / prop
                        where prop = byProportion sheetCorrectorLoad
              allEqual [] = True
              allEqual ((_, c) : xs) = all (\(_, c') -> c == c') xs
          secondResultNorm `shouldSatisfy` allEqual . Map.toList
      )