module Handler.Utils.Allocation
  ( allocationDone
  , ordinalPriorities
  , sinkAllocationPriorities
  , computeAllocation
  , doAllocation
  , ppMatchingLog
  , storeAllocationResult
  ) where

import Import

import qualified Data.Map.Strict as Map

import qualified Database.Esqueleto as E
import qualified Database.Esqueleto.Utils as E

import qualified Control.Monad.State.Class as State (get, modify')

import Data.List (genericLength, elemIndex)
import qualified Data.Vector as Vector
import Data.Vector.Lens (vector)
import qualified Data.Set as Set

import System.Random (mkStdGen)

import Utils.Allocation

import qualified Data.Conduit.List as C

import Data.Generics.Product.Param

import qualified Crypto.Hash as Crypto
import qualified Data.Binary as Binary

import qualified Data.ByteArray as BA (convert)


allocationDone :: AllocationId -> DB (Maybe UTCTime)
allocationDone allocId = fmap (E.unValue <=< listToMaybe) . E.select . E.from $ \participant -> do
  E.where_ $ participant E.^. CourseParticipantAllocated E.==. E.just (E.val allocId)
  return . E.max_ $ participant E.^. CourseParticipantRegistration


ordinalPriorities :: Monad m => ConduitT UserMatriculation (Map UserMatriculation AllocationPriority) m ()
ordinalPriorities = evalStateC 0 . C.mapM $ \matr -> singletonMap matr <$> (AllocationPriorityOrdinal <$> State.get <* State.modify' succ)
  
sinkAllocationPriorities :: AllocationId
                         -> ConduitT (Map UserMatriculation AllocationPriority) Void DB ()
sinkAllocationPriorities allocId = C.mapM_ . imapM_ $ \matr prio ->
  E.update $ \allocationUser -> do
    E.set allocationUser [ AllocationUserPriority E.=. E.val (Just prio) ]
    E.where_ $ allocationUser E.^. AllocationUserAllocation E.==. E.val allocId
    E.where_ . E.exists . E.from $ \user ->
      E.where_ $ user E.^. UserId E.==. allocationUser E.^. AllocationUserUser
           E.&&. user E.^. UserMatrikelnummer E.==. E.val (Just matr)

  
computeAllocation :: AllocationId
                  -> Maybe (Set CourseId) -- ^ Optionally restrict allocation to only consider the given courses
                  -> DB ( AllocationFingerprint
                        , Set (UserId, CourseId)
                        , Seq (MatchingLog UserId CourseId Natural)
                        )
computeAllocation allocId cRestr = do
  allocations <- selectList [ CourseParticipantAllocated ==. Just allocId ] []
  let allocations' = allocations
        & map ((, Sum 1) . courseParticipantUser . entityVal)
        & Map.fromListWith (<>)
        & fmap getSum
  
  users' <- selectList [ AllocationUserAllocation ==. allocId ] []
  let users'' = users'
        & mapMaybe ( runMaybeT $ do
                       user <- lift $ allocationUserUser . entityVal
                       totalCourses <- lift $ allocationUserTotalCourses . entityVal
                       priority <- MaybeT $ allocationUserPriority . entityVal

                       let allocated = Map.findWithDefault 0 user allocations'

                       guard $ totalCourses > allocated

                       return (user, (totalCourses - allocated, priority))
                   )
        & Map.fromList
      cloneCounts = Map.map (view _1) users''
      allocationPrio = view _2 . (Map.!) users''

  courses' <- E.select . E.from $ \(allocationCourse `E.InnerJoin` course) -> do
    E.on $ allocationCourse E.^. AllocationCourseCourse E.==. course E.^. CourseId
     E.&&. allocationCourse E.^. AllocationCourseAllocation E.==. E.val allocId

    let participants = E.subSelectCount . E.from $ \participant -> do
          E.where_ $ participant E.^. CourseParticipantCourse E.==. course E.^. CourseId
          E.where_ . E.not_ . E.exists . E.from $ \lecturer ->
            E.where_ $ lecturer E.^. LecturerCourse E.==. course E.^. CourseId
                 E.&&. lecturer E.^. LecturerUser E.==. participant E.^. CourseParticipantUser

    whenIsJust cRestr $ \restrSet ->
      E.where_ $ course E.^. CourseId `E.in_` E.valList (Set.toList restrSet)

    return ( allocationCourse
           , E.maybe E.nothing (\c -> E.just $ c E.-. participants) (course E.^. CourseCapacity)
           , allocationCourse E.^. AllocationCourseMinCapacity E.-. participants
           )
  let capacities = Map.filter (maybe True (> 0)) . Map.fromList $ (view (_1 . _entityVal . _allocationCourseCourse) &&& view (_2 . _Value)) <$> courses'

  applications' <- selectList [ CourseApplicationAllocation ==. Just allocId ] []
  let alreadyAssigned (Entity _ CourseApplication{..}) = orM
        [ exists [ TransactionLogInfo ==. toJSON (TransactionCourseParticipantDeleted courseApplicationCourse courseApplicationUser) ]
        , exists [ CourseParticipantCourse ==. courseApplicationCourse, CourseParticipantUser ==. courseApplicationUser ]
        ]
  applications'' <- applications'
    & filter ((\CourseApplication{..} -> not courseApplicationRatingVeto && fromMaybe True (courseApplicationRatingPoints ^? _Just . passingGrade . _Wrapped)) . entityVal)
    & filterM (fmap not . alreadyAssigned)
  let preferences = Map.fromList $ do
        Entity _ CourseApplication{..} <- applications''
        guard $ Map.member courseApplicationCourse capacities
        return ((courseApplicationUser, courseApplicationCourse), (courseApplicationAllocationPriority, courseApplicationRatingPoints))

  gradeScale <- getsYesod $ view _appAllocationGradeScale
  gradeOrdinalPlaces <- getsYesod $ view _appAllocationGradeOrdinalPlaces
  let centralNudge user cloneIndex grade = case allocationPrio user of
        AllocationPriorityNumeric{..}
          -> let allocationPriorities' = under vector (sortOn Down) allocationPriorities
                 minPrio | Vector.null allocationPriorities' = 0
                         | otherwise = Vector.last allocationPriorities'
              in AllocationPriorityComparisonNumeric . withNumericGrade . fromInteger . fromMaybe minPrio $ allocationPriorities Vector.!? fromIntegral cloneIndex
        AllocationPriorityOrdinal{..}
          | gradeOrdinalPlaces > 0
            -> let allocationOrdinal' = gradeScale / fromIntegral gradeOrdinalPlaces * fromIntegral allocationOrdinal 
                in AllocationPriorityComparisonOrdinal (Down cloneIndex) $ withNumericGrade allocationOrdinal'
        AllocationPriorityOrdinal{..}
          -> AllocationPriorityComparisonOrdinal (Down cloneIndex) $ fromIntegral allocationOrdinal
        where
          withNumericGrade :: Rational -> Rational
          withNumericGrade
            | Just grade' <- grade
            = let numberGrade' = fromMaybe (error "non-passing grade") (fromIntegral <$> elemIndex grade' passingGrades) / pred (genericLength passingGrades)
                  passingGrades = sort $ filter (view $ passingGrade . _Wrapped) universeF
                  numericGrade = -gradeScale + numberGrade' * 2 * gradeScale
               in (+) numericGrade
            | otherwise
            = id

  let
    fingerprint :: AllocationFingerprint
    fingerprint = Crypto.hash . toStrict $ Binary.encode (users'', capacities, preferences, gradeScale, gradeOrdinalPlaces)
    
    g = mkStdGen $ hash (BA.convert fingerprint :: ByteString)

  let
    doAllocationWithout cs = runWriter $ computeMatchingLog g cloneCounts capacities' preferences' centralNudge
      where
        capacities'  = Map.filterWithKey (\    c  _ -> Set.notMember c cs) capacities
        preferences' = Map.filterWithKey (\(_, c) _ -> Set.notMember c cs) preferences

    allocationLoop cs
      | not $ null belowMin = doAllocationWithout $ cs <> Set.fromList belowMin
      | otherwise           = (allocs, mLog)
      where
        (allocs, mLog) = doAllocationWithout cs
        belowMin = catMaybes . flip map courses' $ \(Entity _ AllocationCourse{..}, _, E.Value minCap) ->
          guardOn (Set.size (Set.filter (\(_, c) -> c == allocationCourseCourse) allocs) < minCap) allocationCourseCourse

  return . (\(ms, mLog) -> (fingerprint, ms, mLog)) $!! allocationLoop Set.empty


doAllocation :: AllocationId
             -> Set (UserId, CourseId)
             -> DB ()
doAllocation allocId regs = do
  now <- liftIO getCurrentTime
  forM_ regs $ \(uid, cid) -> do
    mField <- (courseApplicationField . entityVal =<<) . listToMaybe <$> selectList [CourseApplicationCourse ==. cid, CourseApplicationUser ==. uid, CourseApplicationAllocation ==. Just allocId] []
    void . insertUnique $ CourseParticipant cid uid now mField (Just allocId)

ppMatchingLog :: forall mono.
                 ( MonoFoldable mono
                 , Element mono ~ MatchingLog UserId CourseId Natural
                 )
              => mono -> Text
ppMatchingLog = unlines . map (tshow . pretty) . otoList
  where
    pretty :: MatchingLog UserId CourseId Natural -> MatchingLog Int64 Int64 Natural
    pretty = over (param @1) fromSqlKey
           . over (param @2) fromSqlKey

storeAllocationResult :: AllocationId
                      -> (AllocationFingerprint, Set (UserId, CourseId), Seq (MatchingLog UserId CourseId Natural))
                      -> DB ()
storeAllocationResult allocId (allocFp, allocMatchings, ppMatchingLog -> allocLog) = do
  now <- liftIO getCurrentTime
  insert_ . AllocationMatching allocId allocFp <=< insert $ File "matchings.log" (Just $ encodeUtf8 allocLog) now

  doAllocation allocId allocMatchings