module Utils.AllocationSpec where

import TestImport hiding (Course)

import Utils.Allocation

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

import System.Random (mkStdGen)


data Man = Alpha | Beta | Gamma | Delta
  deriving (Eq, Ord, Bounded, Enum, Read, Show, Generic, Typeable)
instance NFData Man

data Woman = Alef | Bet | Gimel | Dalet
  deriving (Eq, Ord, Bounded, Enum, Read, Show, Generic, Typeable)
  

spec :: Spec
spec = describe "computeMatching" $
  it "produces some expected known matchings" $ do
    example $ do
      let men = Map.fromList $ (, (0, 1)) <$> [Alpha .. Gamma]
          women = Map.fromList $ (, Just 1) <$> [Alef .. Gimel]
          preferences = fmap ((3 -) *** (3 -)) $ Map.fromList
            [ ((Alpha, Alef ), (1, 3))
            , ((Alpha, Bet  ), (2, 2))
            , ((Alpha, Gimel), (3, 1))
            , ((Beta , Alef ), (3, 1))
            , ((Beta , Bet  ), (1, 3))
            , ((Beta , Gimel), (2, 2))
            , ((Gamma, Alef ), (2, 2))
            , ((Gamma, Bet  ), (3, 1))
            , ((Gamma, Gimel), (1, 3))
            ]

          centralNudge _ _ = id

          expectedResult = Set.fromList [(Alpha, Alef), (Beta, Bet), (Gamma, Gimel)]
          ourResult = computeMatching (mkStdGen 0) men women preferences centralNudge
      ourResult `shouldBe` expectedResult

    example $ do
      let men = Map.fromList $ (, (0, 2)) <$> [Alpha,Beta,Delta]
          women = Map.fromList $ (, Just 1) <$> [Alef .. Gimel]
          preferences = fmap ((3 -) *** (3 -)) $ Map.fromList
            [ ((Alpha, Alef ), (1, 3))
            , ((Alpha, Bet  ), (2, 2))
            , ((Alpha, Gimel), (3, 1))
            , ((Beta , Alef ), (3, 1))
            , ((Beta , Bet  ), (1, 3))
            , ((Beta , Gimel), (2, 2))
            , ((Delta, Alef ), (2, 2))
            , ((Delta, Bet  ), (3, 1))
            , ((Delta, Gimel), (1, 3))
            ]

          centralNudge _ _ = id

          expectedResult = Set.fromList [(Alpha, Alef), (Beta, Bet), (Delta, Gimel)]
          ourResult = computeMatching (mkStdGen 0) men women preferences centralNudge
      ourResult `shouldBe` expectedResult
  
    example $ do
      let men = Map.fromList $ (, (0, 2)) <$> [Alpha .. Gamma]
          women = Map.fromList $ (, Just 2) <$> [Alef .. Gimel]
          preferences = fmap ((3 -) *** (3 -)) $ Map.fromList
            [ ((Alpha, Alef ), (1, 3))
            , ((Alpha, Bet  ), (2, 2))
            , ((Alpha, Gimel), (3, 1))
            , ((Beta , Alef ), (3, 1))
            , ((Beta , Bet  ), (1, 3))
            , ((Beta , Gimel), (2, 2))
            , ((Gamma, Alef ), (2, 2))
            , ((Gamma, Bet  ), (3, 1))
            , ((Gamma, Gimel), (1, 3))
            ]

          centralNudge _ _ = id

          expectedResult = Set.fromList [(Alpha, Alef), (Gamma, Alef), (Beta, Bet), (Alpha, Bet), (Beta, Gimel), (Gamma, Gimel)]
          ourResult = computeMatching (mkStdGen 0) men women preferences centralNudge
      ourResult `shouldBe` expectedResult
 
    example $ do
      let men = Map.fromList $ (, (0, 1)) <$> [Alpha .. Delta]
          women = Map.fromList $ (, Just 1) <$> [Alef .. Dalet]
          preferences = fmap ((4 -) *** (4 -)) $ Map.fromList
            [ ((Alpha, Alef ), (1, 3))
            , ((Alpha, Bet  ), (2, 3))
            , ((Alpha, Gimel), (3, 2))
            , ((Alpha, Dalet), (4, 3))
            , ((Beta , Alef ), (1, 4))
            , ((Beta , Bet  ), (4, 1))
            , ((Beta , Gimel), (3, 3))
            , ((Beta , Dalet), (2, 2))
            , ((Gamma, Alef ), (2, 2))
            , ((Gamma, Bet  ), (1, 4))
            , ((Gamma, Gimel), (3, 4))
            , ((Gamma, Dalet), (4, 1))
            , ((Delta, Alef ), (4, 1))
            , ((Delta, Bet  ), (2, 2))
            , ((Delta, Gimel), (3, 1))
            , ((Delta, Dalet), (1, 4))
            ]

          centralNudge _ _ = id

          expectedResult = Set.fromList [(Alpha, Gimel), (Beta, Dalet), (Gamma, Alef), (Delta, Bet)]
          ourResult = computeMatching (mkStdGen 0) men women preferences centralNudge
      ourResult `shouldBe` expectedResult
      
    example $ do
      let men = Map.fromList $ (, (0, 1)) <$> [Alpha .. Delta]
          women = Map.fromList $ (, Just 1) <$> [Alef .. Dalet]
          preferences = fmap ((4 -) *** (4 -)) $ Map.fromList
            [ ((Alpha, Alef ), (1, 3))
            , ((Alpha, Bet  ), (2, 2))
            , ((Alpha, Gimel), (3, 1))
            , ((Alpha, Dalet), (4, 3))
            , ((Beta , Alef ), (1, 4))
            , ((Beta , Bet  ), (2, 3))
            , ((Beta , Gimel), (3, 2))
            , ((Beta , Dalet), (4, 4))
            , ((Gamma, Alef ), (3, 1))
            , ((Gamma, Bet  ), (1, 4))
            , ((Gamma, Gimel), (2, 3))
            , ((Gamma, Dalet), (4, 2))
            , ((Delta, Alef ), (2, 2))
            , ((Delta, Bet  ), (3, 1))
            , ((Delta, Gimel), (1, 4))
            , ((Delta, Dalet), (4, 1))
            ]

          centralNudge _ _ = id

          expectedResult = Set.fromList [(Alpha, Gimel), (Beta, Dalet), (Gamma, Alef), (Delta, Bet)]
          ourResult = computeMatching (mkStdGen 0) men women preferences centralNudge
      ourResult `shouldBe` expectedResult
      
    example $ do
      let students = Map.fromList $ (, (0, 1)) <$> ([1..6] :: [Int])
          colleges = Map.fromList $ (, Just 2) <$> (['A', 'Z', 'C'] :: [Char])
          student_preferences = Map.fromList
            [ ((1, 'A'), 3), ((1, 'Z'), 2), ((1, 'C'), 1)
            , ((2, 'A'), 3), ((2, 'Z'), 1), ((2, 'C'), 2)
            , ((3, 'A'), 3), ((3, 'Z'), 2), ((3, 'C'), 1)
            , ((4, 'A'), 2), ((4, 'Z'), 3), ((4, 'C'), 1)
            , ((5, 'A'), 1), ((5, 'Z'), 3), ((5, 'C'), 2)
            , ((6, 'A'), 2), ((6, 'Z'), 1), ((6, 'C'), 6)
            ]
          preferences = Map.mapWithKey (\(st, _) stPref -> (stPref, 7 - st)) student_preferences

          centralNudge _ _ = id

          expectedResult = Set.fromList [(1, 'A'), (2, 'A'), (3, 'Z'), (4, 'Z'), (5, 'C'), (6, 'C')]
          ourResult = computeMatching (mkStdGen 0) students colleges preferences centralNudge
      ourResult `shouldBe` expectedResult