module Utils.Allocation
  ( computeMatching
  , MatchingLog(..)
  , computeMatchingLog
  ) where

import Import.NoModel hiding (StateT, st, get)

import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Sequence as Seq

import Data.Array.ST (STArray)
import qualified Data.Array.MArray as MArr

import System.Random (RandomGen)
import Control.Monad.Trans.Random.Strict (evalRandT, RandT)
import Control.Monad.Trans.State.Strict (StateT, modify', get, gets, evalStateT)
import Control.Monad.Writer (tell)

import Control.Monad.ST

import Data.List ((!!))


type CourseIndex = Int
type StudentIndex = Int
type CloneIndex = Int

data MatchingLog student course cloneIndex
  = MatchingConsider
    { mlStudent :: student, mlClone :: cloneIndex }
  | MatchingApply
    { mlStudent :: student, mlClone :: cloneIndex, mlCourse :: course }
  | MatchingNoApplyCloneInstability
    { mlStudent :: student, mlClone :: cloneIndex, mlCourse :: course }
  | MatchingLostSpot
    { mlStudent :: student, mlClone :: cloneIndex, mlCourse :: course }
  deriving (Eq, Ord, Read, Show, Generic, Typeable)
instance (NFData student, NFData course, NFData cloneIndex) => NFData (MatchingLog student course cloneIndex)

deriveJSON defaultOptions
  { constructorTagModifier = camelToPathPiece' 1
  , fieldLabelModifier = camelToPathPiece' 1
  } ''MatchingLog

computeMatching :: forall randomGen student course cloneCount cloneIndex capacity studentRatingCourse courseRatingStudent courseRatingStudent'.
                   ( RandomGen randomGen
                   , Ord student, Ord course
                   , NFData student
                   , Ord studentRatingCourse
                   , Ord courseRatingStudent'
                   , Integral cloneCount, Integral capacity, Integral cloneIndex
                   )
                => randomGen -- ^ Source of randomness
                -> Map student (cloneIndex, cloneCount) -- ^ requested number of placements per student
                -> Map course (Maybe capacity) -- ^ capacity of courses
                -> Map (student, course) (studentRatingCourse, courseRatingStudent) -- ^ Mutual preference ordering @(studentRatingCourse, courseRatingStudent)@
                -> (student -> cloneIndex -> courseRatingStudent -> courseRatingStudent') -- ^ Adjust preference ordering of courses (incorporate central priority)
                -> Set (student, course) -- ^ Stable matching
computeMatching g cloneCounts capacities preferences centralNudge
  = view _1 . runWriter $ computeMatchingLog g cloneCounts capacities preferences centralNudge

computeMatchingLog :: forall randomGen student course cloneCount cloneIndex capacity studentRatingCourse courseRatingStudent courseRatingStudent'.
                      ( RandomGen randomGen
                      , Ord student, Ord course
                      , NFData student
                      , Ord studentRatingCourse
                      , Ord courseRatingStudent'
                      , Integral cloneCount, Integral capacity, Integral cloneIndex
                      )
                   => randomGen -- ^ Source of randomness
                   -> Map student (cloneIndex, cloneCount) -- ^ requested number of placements and first cloneIndex per student
                   -> Map course (Maybe capacity) -- ^ capacity of courses
                   -> Map (student, course) (studentRatingCourse, courseRatingStudent) -- ^ Mutual preference ordering @(studentRatingCourse, courseRatingStudent)@
                   -> (student -> cloneIndex -> courseRatingStudent -> courseRatingStudent') -- ^ Adjust preference ordering of courses (incorporate central priority)
                   -> Writer (Seq (MatchingLog student course cloneIndex)) (Set (student, course)) -- ^ Stable matching
computeMatchingLog g cloneCounts capacities preferences centralNudge = writer $ runST computeMatching'
  where
    computeMatching' :: forall s. ST s (Set (student, course), Seq (MatchingLog student course cloneIndex))
    computeMatching' = runWriterT . flip evalRandT g $ do
      stb <- (Map.!) <$> sequence (Map.fromSet (const getRandom) clonedStudents) :: RandT randomGen (WriterT _ (ST s)) ((student, CloneIndex) -> UUID)
      cstb <- (Map.!) <$> sequence (Map.fromSet (const getRandom) courses) :: RandT randomGen (WriterT _ (ST s)) (course -> UUID)

      courses' <- lift . lift . MArr.newListArray courseBounds . map initCourse $ Set.toAscList courses :: RandT randomGen (WriterT _ (ST s)) (STArray s CourseIndex (Either (Set (student, CloneIndex)) (Seq (student, CloneIndex))))

      stPrefs <- lift . lift $ MArr.newArray studentBounds [] :: RandT randomGen (WriterT _ (ST s)) (STArray s (StudentIndex, CloneIndex) [course])
      forM_ clonedStudents $ \(st, cn) ->
        lift . lift . MArr.writeArray stPrefs (st ^. contStudents, cn) $ studentPrefs cstb (st, cn)

      let
        propose :: StateT (Set (student, CloneIndex)) (WriterT _ (ST s)) ()
        propose = (=<< get) . mapM_ $ \(st, cn) -> do
          lift . tell . pure . MatchingConsider st $ fromIntegral cn
          let markDone = modify' $ Set.delete (st, cn)

          options <- lift . lift $ MArr.readArray stPrefs (st ^. contStudents, cn)
          case options of
            []     -> markDone
            c : cs -> do
              lift . lift $ MArr.writeArray stPrefs (st ^. contStudents, cn) cs
              cState <- lift . lift $ MArr.readArray courses' (c ^. contCourses)
              case cState of
                Left pSet
                  | none (\(st', _) -> st == st') pSet -> do
                      lift . tell . pure $ MatchingApply st (fromIntegral cn) c
                      lift . lift . MArr.writeArray courses' (c ^. contCourses) $!! Left (Set.insert (st, cn) pSet)
                      markDone
                Right spots
                  | none (\(st', _) -> st == st') spots -> do
                      courseMatchings <- lift . lift $ MArr.getAssocs courses'
                      let
                        matchingCourse s cn' = listToMaybe $ do
                          (review contCourses -> course, students) <- courseMatchings
                          student <- case students of
                            Left  pSet   -> toList pSet
                            Right spots' -> toList spots'
                          guard $ (s, cn') == student
                          return course

                      let capacity = maybe (error "course without capacity treated as one") fromIntegral . fromMaybe (error "course not found in capacities") $ capacities Map.!? c
                          (worseSpots, betterSpots) = Seq.spanr isWorseSpot spots
                          isWorseSpot existing = case (comparing $ fromMaybe (error "(st, c) not in preferences") . courseRating c &&& stb) existing (st, cn) of
                            EQ -> error "Two student-clones compared equal in the face of stb"
                            GT -> False
                            LT -> True
                          (newSpots, lostSpots) = force . Seq.splitAt capacity $ betterSpots <> Seq.singleton (st, cn) <> worseSpots

                          isUnstableWith :: CloneIndex -> (student, CloneIndex) -> Bool
                          isUnstableWith cn' (stO, cnO) = Just True == (do
                            c' <- matchingCourse st cn'
                            rMe <- courseRating c' (st, cn')
                            rOther <- courseRating c' (stO, cnO)
                            return $ LT == compare (rMe, stb (st, cn')) (rOther, stb (stO, cnO)))

                      if | any (uncurry isUnstableWith) $ (,) <$> [0,1..pred cn] <*> toList lostSpots
                           -> lift . tell . pure $ MatchingNoApplyCloneInstability st (fromIntegral cn) c
                         | Seq.length betterSpots >= capacity
                           -> return ()
                         | otherwise
                           -> do
                             lift . tell . pure $ MatchingApply st (fromIntegral cn) c
                             lift . lift . MArr.writeArray courses' (c ^. contCourses) $ Right newSpots
                             forM_ lostSpots $ \(st', cn') -> do
                               lift . tell . pure $ MatchingLostSpot st' (fromIntegral cn') c
                               modify' $ Set.insert (st', cn')
                             markDone
                _other -> return ()

        proposeLoop = do
          propose
          done <- gets Set.null
          unless done
            proposeLoop

      lift $ evalStateT proposeLoop clonedStudents


      -- let
      --   pairwiseExchange :: ST s ()
      --   pairwiseExchange = do
      --     let possiblePairs = do
      --           (s:ss) <- tails . sortOn stb $ toList clonedStudents
      --           s' <- ss
      --           return (s, s')
      --         matchingCourse (s, c) = do
      --           courseMatchings <- MArr.getAssocs courses'
      --           return . listToMaybe $ do
      --             (course, students) <- courseMatchings
      --             student <- case students of
      --               Left  pSet  -> toList pSet
      --               Right spots -> toList spots
      --             guard $ (s, c) == student
      --             return course
      --     forM_ possiblePairs $ \((a, cna), (b, cnb)) -> void . runMaybeT $ do
      --       ca <- MaybeT $ matchingCourse (a, cna)
      --       cb <- MaybeT $ matchingCourse (b, cnb)

      --       let rank (s, cn) c = Seq.elemIndexL c $ studentPrefs cstb (s, cn)
      --       caRa <- hoistMaybe $ rank (a, cna) ca 
      --       caRb <- hoistMaybe $ rank (b, cnb) ca 
      --       cbRa <- hoistMaybe $ rank (a, cna) cb
      --       cbRb <- hoistMaybe $ rank (b, cnb) cb 

      --       let currentRanks cop = caRa `cop` cbRb
      --           newRanks cop = cbRa `cop` caRb

      --           swapImproves = or
      --             [ currentRanks (+) > newRanks (+)
      --             ,    currentRanks (+) == newRanks (+)
      --               && currentRanks min > newRanks min
      --             ]
      --       lift . when swapImproves $ do
      --         traceM $ show (a, cna) <> " `swap` " <> show (b, cnb)
      --         let
      --           addCourseUser :: course -> (student, CloneIndex) -> ST s ()
      --           addCourseUser c (st, cn) = do
      --             cState <- MArr.readArray courses' c
      --             case cState of
      --               Left pSet ->
      --                 MArr.writeArray courses' c $!! Left (Set.insert (st, cn) pSet)
      --               Right spots ->
      --                 let (worseSpots, betterSpots) = Seq.spanr isWorseSpot spots
      --                     isWorseSpot existing = case (comparing $ courseRating c &&& stb) existing (st, cn) of
      --                       EQ -> error "Two student-clones compared equal in the face of stb"
      --                       GT -> False
      --                       LT -> True
      --                     newSpots = force $ betterSpots <> Seq.singleton (st, cn) <> worseSpots
      --                  in MArr.writeArray courses' c $ Right newSpots
      --           remCourseUser :: course -> (student, CloneIndex) -> ST s ()
      --           remCourseUser c (st, cn) = do
      --             cState <- MArr.readArray courses' c
      --             case cState of
      --               Left pSet ->
      --                 MArr.writeArray courses' c $!! Left (Set.delete (st, cn) pSet)
      --               Right spots ->
      --                 MArr.writeArray courses' c $!! Right (Seq.filter (/= (st, cn)) spots)

      --         remCourseUser ca (a, cna)
      --         remCourseUser cb (b, cnb)
      --         addCourseUser cb (a, cna)
      --         addCourseUser ca (b, cnb)

      -- lift pairwiseExchange


      courseMatchings <- lift . lift $ MArr.getAssocs courses'
      return . Set.fromList $ do
        (review contCourses -> course, students) <- courseMatchings
        student <- case students of
          Left  pSet  -> view _1 <$> toList pSet
          Right spots -> view _1 <$> toList spots
        return (student, course)

    courseRating :: course -> (student, CloneIndex) -> Maybe courseRatingStudent'
    courseRating c (st, cn) = do
      (_, courseRating') <- preferences Map.!? (st, c)
      return $ centralNudge st (fromIntegral cn) courseRating'

    cloneIndices :: cloneIndex -> cloneCount -> Set CloneIndex
    cloneIndices firstClone clones = Set.fromList $ map fromIntegral [firstClone, pred $ firstClone + fromIntegral clones]

    clonedStudents :: Set (student, CloneIndex)
    clonedStudents = Set.fromDistinctAscList $ do
      (student, (firstClone, clones)) <- Map.toAscList cloneCounts
      clone <- Set.toAscList $ cloneIndices firstClone clones
      return (student, clone)

    contStudents :: Iso' student StudentIndex
    contStudents = iso toInt fromInt
      where
        students' = Map.keys cloneCounts

        toInt   = fromMaybe (error "trying to resolve unknown student") . flip elemIndex students'
        fromInt = (!!) students'

    studentBounds :: ((StudentIndex, CloneIndex), (StudentIndex, CloneIndex))
    studentBounds = ((0, 0), (pred $ Map.size cloneCounts, fromMaybe 0 $ maximumOf (folded . to (uncurry cloneIndices) . folded) cloneCounts))

    courses :: Set course
    courses = Set.fromDistinctAscList . map (view _1) . filter (maybe True (> 0) . view _2) $ Map.toAscList capacities
    courseBounds :: (CourseIndex, CourseIndex)
    courseBounds = (0, pred $ Set.size courses)

    contCourses :: Iso' course CourseIndex
    contCourses = iso toInt fromInt
      where
        courses' = Set.toAscList courses

        toInt   = fromMaybe (error "trying to resolve unknown course") . flip elemIndex courses'
        fromInt = (!!) courses'

    initCourse :: course -> Either (Set (student, CloneIndex)) (Seq (student, CloneIndex))
    initCourse c
      | is _Just . join $ Map.lookup c capacities
      = Right Seq.empty
      | otherwise
      = Left Set.empty

    studentPrefs :: forall a. Ord a => (course -> a) -> (student, CloneIndex) -> [course]
    studentPrefs cstb (st, _) = map (view _1) . sortOn (Down . view _2) . mapMaybe (\c -> (c, ) <$> cPref c) $ Set.toList courses
      where
        cPref :: course -> Maybe (studentRatingCourse, a)
        cPref c = do
          (cPref', _) <- Map.lookup (st, c) preferences
          return (cPref', cstb c)