module Handler.Utils.Allocation
  ( allocationStarted, allocationNotifyNewCourses
  , ordinalPriorities
  , sinkAllocationPriorities
  , MatchingLogRun(..)
  , computeAllocation
  -- , doAllocation -- Use `storeAllocationResult`
  , ppMatchingLog
  , storeAllocationResult
  ) where

import Import

import qualified Data.Map.Strict as Map

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

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

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

import qualified Data.Binary as Binary
import Crypto.Hash.Algorithms (SHAKE256)
import Crypto.Random (drgNewSeed, seedFromBinary)
import Crypto.Error (onCryptoFailure)

import Utils.Allocation

import qualified Data.Conduit.List as C

import Data.Generics.Product.Param

import qualified Crypto.Hash as Crypto

import Language.Haskell.TH (nameBase)


data MatchingExcludedReason
  = MatchingExcludedParticipationExisted
  | MatchingExcludedParticipationExists
  | MatchingExcludedVeto
  | MatchingExcludedLecturer
  deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Typeable)
  deriving anyclass (Universe, Finite, NFData)

nullaryPathPiece ''MatchingExcludedReason $ camelToPathPiece' 2
pathPieceJSON ''MatchingExcludedReason

data MatchingLogRun = MatchingLogRun
  { matchingLogRunCourseRestriction :: Maybe (Set CourseId)
  , matchingLogRunCoursesExcluded :: Set CourseId
  , matchingLogMatchingsExcluded :: Map (UserId, CourseId) (NonNull (Set MatchingExcludedReason))
  , matchingLogRunLog :: Seq (MatchingLog UserId CourseId Natural)
  } deriving (Eq, Ord, Read, Show, Generic, Typeable)
    deriving anyclass (NFData)

deriveJSON defaultOptions
  { fieldLabelModifier = camelToPathPiece' 3
  } ''MatchingLogRun


allocationStarted :: AllocationId -> DB (Maybe UTCTime)
-- ^ Time the first allocation was made
allocationStarted allocId = fmap (E.unValue <=< listToMaybe) . E.select . E.from $ \allocationMatching -> do
  E.where_ $ allocationMatching E.^. AllocationMatchingAllocation E.==. E.val allocId
  return . E.min_ $ allocationMatching E.^. AllocationMatchingTime

allocationNotifyNewCourses :: E.SqlExpr (E.Value AllocationId)
                           -> E.SqlExpr (E.Value UserId)
                           -> E.SqlExpr (E.Value Bool)
allocationNotifyNewCourses allocId uid = ( hasOverride True E.||. hasApplication E.||. isParticipant )
                                   E.&&.   E.not_ (hasOverride False)
  where
    hasOverride overrideVal = E.exists . E.from $ \allocationNotificationSetting ->
      E.where_ $ allocationNotificationSetting E.^. AllocationNotificationSettingUser E.==. uid
           E.&&. allocationNotificationSetting E.^. AllocationNotificationSettingAllocation E.==. allocId
           E.&&. allocationNotificationSetting E.^. AllocationNotificationSettingIsOptOut E.==. E.val (not overrideVal)

    hasApplication = E.exists . E.from $ \application ->
      E.where_ $ application E.^. CourseApplicationAllocation E.==. E.just allocId
           E.&&. application E.^. CourseApplicationUser E.==. uid

    isParticipant = E.exists . E.from $ \allocationUser ->
      E.where_ $ allocationUser E.^. AllocationUserAllocation E.==. allocId
           E.&&. allocationUser E.^. AllocationUserUser E.==. uid


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 Int64
sinkAllocationPriorities allocId = fmap getSum . C.foldMapM . ifoldMapM $ \matr prio ->
  fmap Sum . E.updateCount $ \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 :: Entity Allocation
                  -> Maybe (Set CourseId) -- ^ Optionally restrict allocation to only consider the given courses
                  -> DB ( AllocationFingerprint
                        , Set CourseId
                        , Set (UserId, CourseId)
                        , Seq MatchingLogRun
                        )
computeAllocation (Entity allocId Allocation{allocationMatchingSeed}) cRestr = do
  allocations <- selectList [ CourseParticipantAllocated ==. Just allocId, CourseParticipantState ==. CourseParticipantActive ] []
  let allocations' = allocations
        & map ((, Sum 1) . courseParticipantUser . entityVal)
        & Map.fromListWith (<>)

  deregistrations <- E.select . E.from $ \(allocationDeregister `E.InnerJoin` courseParticipant) -> do
    E.on $ courseParticipant E.^. CourseParticipantUser E.==. allocationDeregister E.^. AllocationDeregisterUser
     E.&&. E.just (courseParticipant E.^. CourseParticipantCourse) E.==. allocationDeregister E.^. AllocationDeregisterCourse
    E.where_ $ courseParticipant E.^. CourseParticipantState E.!=. E.val CourseParticipantActive
         E.&&. courseParticipant E.^. CourseParticipantAllocated E.==. E.just (E.val allocId)
    return $ allocationDeregister E.^. AllocationDeregisterUser
  let deregistrations' = deregistrations
        & map ((, Sum 1) . E.unValue)
        & Map.fromListWith (<>)

  users' <- selectList [ AllocationUserAllocation ==. allocId ] []
  let users'' = users'
        & mapMaybe ( runMaybeT $ do
                       user <- lift $ allocationUserUser . entityVal
                       totalCourses <- lift $ allocationUserTotalCourses . entityVal
                       priority <- MaybeT $ allocationUserPriority . entityVal

                       let Sum allocated = Map.findWithDefault 0 user allocations' <> Map.findWithDefault 0 user deregistrations'

                       guard $ totalCourses > allocated

                       return (user, ((allocated, totalCourses - allocated), priority))
                   )
        & Map.fromList
      cloneCounts = Map.map (views _1 $ bimap fromIntegral fromIntegral) 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.&&. participant E.^. CourseParticipantState E.==. E.val CourseParticipantActive
          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'
      eligibleCourses = setOf (folded . _1 . _entityVal . _allocationCourseCourse) courses'

  applications' <- selectList [ CourseApplicationAllocation ==. Just allocId ] []
  excludedMatchings <- flip execStateT mempty . forM_ applications' $ \(Entity _ CourseApplication{..}) -> do
    let
      tellExcluded :: MatchingExcludedReason -> StateT _ _ ()
      tellExcluded reason = State.modify' $ Map.insertWith (<>) (courseApplicationUser, courseApplicationCourse) (opoint reason :: NonNull (Set MatchingExcludedReason))

    when (courseApplicationRatingVeto || maybe False not (courseApplicationRatingPoints ^? _Just . passingGrade . _Wrapped)) $
      tellExcluded MatchingExcludedVeto

    allocStarted <- lift $ allocationStarted allocId
    whenIsJust allocStarted $ \allocStarted' -> do
      let partDeleted = lift $ or2M
            (exists [ TransactionLogInfo ==. toJSON (TransactionCourseParticipantDeleted courseApplicationCourse courseApplicationUser), TransactionLogTime >=. allocStarted' ])
            (exists [ CourseParticipantCourse ==. courseApplicationCourse, CourseParticipantUser ==. courseApplicationUser, CourseParticipantState !=. CourseParticipantActive ])
      whenM partDeleted $
        tellExcluded MatchingExcludedParticipationExisted

    let partExists :: StateT _ DB Bool
        partExists = lift $ exists [ CourseParticipantCourse ==. courseApplicationCourse, CourseParticipantUser ==. courseApplicationUser, CourseParticipantState ==. CourseParticipantActive ]
    whenM partExists $
      tellExcluded MatchingExcludedParticipationExists

    let lecturerExists = lift $ exists [ LecturerCourse ==. courseApplicationCourse, LecturerUser ==. courseApplicationUser ]
    whenM lecturerExists $
      tellExcluded MatchingExcludedLecturer
  let applications'' = applications'
        & map entityVal
        & filter (\CourseApplication{..} -> Map.notMember (courseApplicationUser, courseApplicationCourse) excludedMatchings)
  let preferences = Map.fromList $ do
        CourseApplication{..} <- applications''
        guard $ Map.member courseApplicationCourse capacities
        return ((courseApplicationUser, courseApplicationCourse), (courseApplicationAllocationPriority, courseApplicationRatingPoints))

  gradeScale <- getsYesod $ view _appAllocationGradeScale
  gradeOrdinalProportion <- getsYesod $ view _appAllocationGradeOrdinalProportion
  let ordinalUsers = getSum . flip foldMap users'' $ \(_, prio) -> case prio of
        AllocationPriorityOrdinal{} -> Sum 1
        _other -> mempty
      gradeOrdinalPlaces :: Natural
      gradeOrdinalPlaces = round . abs $ ordinalUsers * gradeOrdinalProportion

  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' = maybe (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
    inputs = Binary.encode (users'', capacities, preferences, gradeScale, gradeOrdinalPlaces)

    fingerprint :: AllocationFingerprint
    fingerprint = Crypto.hashlazy inputs

    g = onCryptoFailure (\_ -> error "Could not create DRG") id . fmap drgNewSeed . seedFromBinary $ kmaclazy @(SHAKE256 320) (encodeUtf8 . (pack :: String -> Text) $ nameBase 'computeAllocation) allocationMatchingSeed inputs

  let
    doAllocationWithout :: Set CourseId -> Writer (Seq (MatchingLog UserId CourseId Natural)) (Set (UserId, CourseId))
    doAllocationWithout cs = 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 :: Set CourseId -> Writer (Seq MatchingLogRun) (Set (UserId, CourseId))
    allocationLoop cs = do
      allocs <- mapWriter (over _2 $ pure . MatchingLogRun cRestr cs excludedMatchings) $ doAllocationWithout cs
      let
        belowMin = catMaybes . flip map courses' $ \(Entity _ AllocationCourse{..}, _, E.Value minCap) -> do
          guard . not $ Set.member allocationCourseCourse cs
          guard $ Set.size (Set.filter (\(_, c) -> c == allocationCourseCourse) allocs) < minCap
          return allocationCourseCourse
      if
        | not $ null belowMin -> allocationLoop $ cs <> Set.fromList belowMin
        | otherwise           -> return allocs

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


doAllocation :: AllocationId
             -> UTCTime
             -> Set (UserId, CourseId)
             -> DB ()
doAllocation allocId now regs =
  forM_ regs $ \(uid, cid) -> do
    void $ upsert
      (CourseParticipant cid uid now (Just allocId) CourseParticipantActive)
      [ CourseParticipantRegistration =. now
      , CourseParticipantAllocated =. Just allocId
      , CourseParticipantState =. CourseParticipantActive
      ]
    audit $ TransactionCourseParticipantEdit cid uid

ppMatchingLog :: Seq MatchingLogRun -> Text
ppMatchingLog = unlines . map prettyRun . otoList
  where
    prettyRun MatchingLogRun{..} = unlines
      [ "----- STARTING RUN -----"
      , "Course restriction: " <> tshow (Set.toAscList <$> matchingLogRunCourseRestriction)
      , "Courses excluded: " <> tshow (Set.toAscList matchingLogRunCoursesExcluded)
      , "Matchings excluded (user, course): "
      , unlines . map ("  " <>) . flip ifoldMap matchingLogMatchingsExcluded $ \(uid, cid) (otoList -> reasons) -> pure $
          "(" <> tshow (fromSqlKey uid) <> ", " <> tshow (fromSqlKey cid) <> ") " <> intercalate ", " (map tshow reasons) :: [Text]
      , "------------------------"
      , unlines . map (tshow . pretty) $ otoList matchingLogRunLog
      , "------ RUN ENDED -------"
      ]

    pretty :: MatchingLog UserId CourseId Natural -> MatchingLog Int64 Int64 Natural
    pretty = over (param @1) fromSqlKey
           . over (param @2) fromSqlKey

storeAllocationResult :: AllocationId
                      -> UTCTime
                      -> (AllocationFingerprint, Set (UserId, CourseId), Seq MatchingLogRun)
                      -> DB ()
storeAllocationResult allocId now (allocFp, allocMatchings, ppMatchingLog -> allocLog) = do
  FileReference{..} <- sinkFile $ File "matchings.log" (Just . yield $ encodeUtf8 allocLog) now
  insert_ . AllocationMatching allocId allocFp now $ fromMaybe (error "allocation result stored without fileReferenceContent") fileReferenceContent

  doAllocation allocId now allocMatchings
