SMO low% routes comparison tool

data Kingdom = Cap | Cascade | Sand | Lake | Wooded | Cloud | Lost | Metro | Snow | Seaside | Luncheon | Ruined | Bowsers | Moon | Mushroom | DarkSide | DarkerSide
    deriving (Show, Eq)

allKingdoms = [Cap , Cascade , Sand , Lake , Wooded , Cloud, Lost , Metro , Snow , Seaside , Luncheon , Ruined , Bowsers , Moon , Mushroom , DarkSide , DarkerSide]

data TravelMethod = Odyssey | Painting | Auto
    deriving (Show, Eq)

travelAPresses :: TravelMethod -> Int
travelAPresses Odyssey = 2 -- not counting before-postgame cappy dialogue, but that doesn't matter for comparative puropses (TODO: does matter if we decide to do lake odyssey skip)
travelAPresses Painting = 1
travelAPresses Auto = 0

type KingdomOrder = (Bool, Bool)

lakeOrWoodedFirst :: KingdomOrder -> Kingdom
lakeOrWoodedFirst (True, _) = Lake
lakeOrWoodedFirst (False, _) = Wooded

lakeOrWoodedSecond :: KingdomOrder -> Kingdom
lakeOrWoodedSecond (lw, ss) = lakeOrWoodedFirst (not lw, ss)

snowOrSeasideFirst :: KingdomOrder -> Kingdom
snowOrSeasideFirst (_, True) = Snow
snowOrSeasideFirst (_, False) = Seaside

snowOrSeasideSecond :: KingdomOrder -> Kingdom
snowOrSeasideSecond (lw, ss) = snowOrSeasideFirst (lw, not ss)

paintingDest :: KingdomOrder -> Kingdom -> Maybe Kingdom
paintingDest ko k = case k of
    Mushroom -> Just $ snowOrSeasideSecond ko
    Cascade -> Just Bowsers
    Bowsers -> Just $ snowOrSeasideFirst ko
    Sand -> Just Metro
    Metro -> Just $ lakeOrWoodedSecond ko
    Luncheon -> Just Mushroom
    _ -> if k == snowOrSeasideFirst ko then
            Just $ lakeOrWoodedFirst ko
        else if k == snowOrSeasideSecond ko then
            Just Cascade
        else if k == lakeOrWoodedFirst ko then
            Just Sand
        else if k == lakeOrWoodedSecond ko then
            Just Luncheon
        else Nothing

-- Which story moons in each kingdom have extra A presses if done before postgame
storyMoonPenalties :: Kingdom -> [Int]
storyMoonPenalties k = case k of
    Sand -> [4,7]
    Lake -> [1]
    Wooded -> [1]
    Luncheon -> [1,4]
    _ -> []

-- after n indicates it's unlocked after n story moons in the above list.
data StoryMoonIndex = After Int | Postgame

-- where the art is, where the moon is, and when it's unlocked
type HintArt = (Kingdom, Kingdom, StoryMoonIndex)

darkSideArt :: Kingdom -> HintArt
darkSideArt k = (DarkSide, k, Postgame)

hintArts :: [HintArt]
hintArts = [
    (Cap, Moon, Postgame),
    (Sand, Bowsers, Postgame),
    (Lake, Cascade, After 1),
    (Wooded, Sand, After 0),
    (Metro, Lake, After 0),
    (Snow, Lost, After 0),
    (Seaside, Metro, After 0),
    (Luncheon, Seaside, After 2),
    (Bowsers, Sand, Postgame),
    (Moon, Wooded, Postgame),
    (Mushroom, Cap, Postgame),
    darkSideArt Cascade,
    darkSideArt Metro,
    darkSideArt Mushroom,
    darkSideArt Cloud,
    darkSideArt Snow,
    darkSideArt Seaside,
    darkSideArt Lost,
    darkSideArt Luncheon,
    darkSideArt Lake,
    darkSideArt Ruined]


type StoryRequirement = (Kingdom, StoryMoonIndex)
touristOrder :: [StoryRequirement]
touristOrder = [
    (Sand, After 2),
    (Metro, After 0),
    (Cascade, After 0),
    (Luncheon, After 2),
    (Moon, After 0),
    (Mushroom, Postgame),
    (Sand, Postgame)]


hasPainting :: Kingdom -> Bool
hasPainting k = paintingDest (True, True) k /= Nothing

-- for verifying peach is got in all kingdoms before moon
hasPeach :: Kingdom -> Bool
hasPeach k = not (k `elem` [DarkSide, DarkerSide, Moon, Mushroom])

type Route = [(TravelMethod, Kingdom)]
type KingdomIndex = Int

expect :: String -> Maybe a -> a
expect _ (Just x) = x
expect s Nothing = error ("Expected " ++ s)

firstIndexOf :: Kingdom -> Route -> KingdomIndex
firstIndexOf x r = expect (show x) $ foldr (\(_, y) res ->
    if x == y then Just 0 else fmap (+1) res) Nothing r

lastIndexOf :: Kingdom -> Route -> KingdomIndex
lastIndexOf x xs = (length xs - 1) - firstIndexOf x (reverse xs) -- ugly

lastIndexOfBefore x i xs = lastIndexOf x (take i xs)

isPostgame :: KingdomIndex -> Route -> Bool
isPostgame i r = i >= firstIndexOf Mushroom r

postgame :: Route -> Route
postgame r = drop (firstIndexOf Mushroom r) r

kingdoms :: Route -> [Kingdom]
kingdoms = map snd

getKingdomOrder :: Route -> KingdomOrder
getKingdomOrder r = (firstIndexOf Lake r < firstIndexOf Wooded r, firstIndexOf Snow r < firstIndexOf Seaside r)

expectedBeforePostgame :: KingdomOrder -> Route
expectedBeforePostgame ko = [
    (Auto, Cap),
    (Auto, Cascade),
    (Odyssey, Sand),
    (Odyssey, lakeOrWoodedFirst ko),
    (Odyssey, lakeOrWoodedSecond ko),
    (Odyssey, Cloud),
    (Auto, Lost),
    (Odyssey, Metro),
    (Odyssey, snowOrSeasideFirst ko),
    (Odyssey, snowOrSeasideSecond ko),
    (Odyssey, Luncheon),
    (Odyssey, Ruined),
    (Odyssey, Bowsers),
    (Odyssey, Moon),
    (Auto, Mushroom)]


-- TODO: support sand->lake Odyssey Skip
validateBeforePostgame :: Route -> Bool
validateBeforePostgame r = expectedBeforePostgame (getKingdomOrder r) == take ((firstIndexOf Mushroom r) +1) r

validatePostgame :: Route -> Bool
validatePostgame r = all (`elem` (kingdoms $ postgame r)) allKingdoms

validatePaintingDests :: Route -> Bool
validatePaintingDests r =
        all (\i ->
            let ((_, s), (m, d)) = (r!!i, r!!(i+1)) in
                m /= Painting || paintingDest ko s == Just d) [0..length r -2]
    where ko = getKingdomOrder r

validateAllPaintingsUsed :: Route -> Bool
validateAllPaintingsUsed r =
    all (\k -> (Painting, k) `elem` r) (filter hasPainting allKingdoms)

validateTravel :: Route -> Bool
validateTravel r = all (/= Auto) $ map fst $ drop ((firstIndexOf Mushroom r)+1) r

validatePeach :: Route -> Bool
validatePeach r = all (`elem` (kingdoms peachActive)) (filter hasPeach allKingdoms)
    where peachActive = dropWhile (\(m, _) -> m /= Odyssey) mushToMoon
          mushToMoon = postgame $ take (lastIndexOf Moon r) r

singleHintArtReqs :: Route -> HintArt -> StoryRequirement
singleHintArtReqs r (s, d, req) =
    let i = lastIndexOfBefore s (lastIndexOf d r) r in
    if isPostgame i r then (s, After 0) else (s, req)

hintArtReqs :: Route -> [StoryRequirement]
hintArtReqs r = map (singleHintArtReqs r) hintArts

touristReqs :: Route -> [StoryRequirement]
touristReqs r = snd $ foldr (\req@(k, _) (idxNext, reqs) ->
    let idxThis = lastIndexOfBefore k idxNext r in
        (idxThis, if isPostgame idxThis r then reqs else req:reqs)) (length r, []) touristOrder

aPressesStoryMoons :: [StoryRequirement] -> Kingdom -> Int
aPressesStoryMoons reqs k = (0:storyMoonPenalties k) !! i
    where i = foldr max 0 reqs'
          reqs' = map (\(_,x) -> case x of
            Postgame -> error ((show k) ++ " needs to be postgame")
            After n -> n) $ filter (\(k',_) -> k==k') reqs

validations :: [Route -> Bool]
validations = [validateBeforePostgame, validatePostgame, validatePaintingDests, validateAllPaintingsUsed, validateTravel, validatePeach]

countAPresses :: Route -> Int
countAPresses r = if not $ all (\v -> v r) validations
    then error "Route invalid"
    else let reqs = hintArtReqs r ++ touristReqs r in
    sum (map (aPressesStoryMoons reqs) allKingdoms)
    + sum (map (travelAPresses . fst) r)

speedrunRoute :: Route
speedrunRoute = [
    (Auto, Cap),
    (Auto, Cascade),
    (Odyssey, Sand),
    (Odyssey, Wooded),
    (Odyssey, Lake),
    (Odyssey, Cloud),
    (Auto, Lost),
    (Odyssey, Metro),
    (Odyssey, Snow),
    (Odyssey, Seaside),
    (Odyssey, Luncheon),
    (Odyssey, Ruined),
    (Odyssey, Bowsers),
    (Odyssey, Moon),
    (Auto, Mushroom),
    (Painting, Seaside),
    (Odyssey, Cap),
    (Odyssey, DarkSide),
    (Odyssey, Moon),
    (Odyssey, Seaside),
    (Painting, Cascade),
    (Painting, Bowsers),
    (Painting, Snow),
    (Painting, Wooded),
    (Painting, Sand),
    (Painting, Metro),
    (Painting, Lake),
    (Painting, Luncheon),
    (Painting, Mushroom),
    (Odyssey, DarkerSide),
    (Odyssey, Bowsers),
    (Odyssey, Ruined),
    (Odyssey, Lost),
    (Odyssey, Cloud),
    (Odyssey, Moon),
    (Odyssey, Mushroom),
    (Odyssey, Sand)]


aaronSuggestedRoute :: Route
aaronSuggestedRoute = [
    (Auto, Cap),
    (Auto, Cascade),
    (Odyssey, Sand),
    (Odyssey, Wooded),
    (Odyssey, Lake),
    (Odyssey, Cloud),
    (Auto, Lost),
    (Odyssey, Metro),
    (Odyssey, Snow),
    (Odyssey, Seaside),
    (Odyssey, Luncheon),
    (Odyssey, Ruined),
    (Odyssey, Bowsers),
    (Odyssey, Moon),
    (Auto, Mushroom),
    (Odyssey, Cap),
    (Odyssey, DarkSide),
    (Odyssey, Moon),
    (Odyssey, Wooded),
    (Painting, Sand),
    (Painting, Metro),
    (Painting, Lake),
    (Painting, Luncheon),
    (Painting, Mushroom),
    (Painting, Seaside),
    (Painting, Cascade),
    (Painting, Bowsers),
    (Painting, Snow),
    (Painting, Wooded),
    (Odyssey, DarkerSide),
    (Odyssey, Luncheon),
    (Odyssey, Ruined),
    (Odyssey, Lost),
    (Odyssey, Cloud),
    (Odyssey, Moon),
    (Odyssey, Mushroom),
    (Odyssey, Sand)]


main :: IO ()
main = do
    print $ countAPresses speedrunRoute -- 68
    print $ countAPresses aaronSuggestedRoute -- 56
-- these 12 saved a presses come from being able to do certain story moons in postgame