module Y2022.Day21.Solution where

import CommonParsers (Parser, nameP, numberP, runParser)
import Data.List (nub)
import Data.Map.Lazy (Map)
import qualified Data.Map.Lazy as Map
import qualified Text.Megaparsec as P
import qualified Text.Megaparsec.Char as PC

yearNr :: Int
yearNr = 2022

dayNr :: Int
dayNr = 21

run :: IO ()
run = do
  putStrLn $ "YEAR " <> show yearNr <> "/ DAY " <> show dayNr

  input <- loadInput

  let result1 = part1 input
  putStrLn $ "\t Part 1: " ++ show result1

  let result2 = part2 input
  putStrLn $ "\t Part 2: " ++ show result2

  putStrLn "---\n"

----------------------------------------------------------------------
-- solutions

part1 :: Input -> Number
part1 inp = ymap Map.! "root"
 where
  ymap = yellMap $ initMonkeyMap inp

-- should be 3375719472770
part2 :: Input -> Number
part2 inp =
  case solve (getRootEq inp) of
    (AVariable, AConst r) -> r
    _ -> error "did not solve"

adjustPart2 :: MonkeyMap -> MonkeyMap
adjustPart2 = Map.adjust toEqual "root" . Map.insert "humn" Variable
 where
  toEqual (Const _) = error "no monkeys involved"
  toEqual Variable = error "no monkey involved"
  toEqual (Add a b) = Equals a b
  toEqual (Subtract a b) = Equals a b
  toEqual (Multiply a b) = Equals a b
  toEqual (Divide a b) = Equals a b
  toEqual op@(Equals _ _) = op

----------------------------------------------------------------------
-- data model

type Input = [Monkey]

type MonkeyName = String

type Number = Rational

type Monkey = (MonkeyName, NameOperation)

data Operation a
  = Const Number
  | Variable
  | Add a a
  | Subtract a a
  | Multiply a a
  | Divide a a
  | Equals a a
  deriving (Show)

type NameOperation = Operation MonkeyName

involved :: NameOperation -> [MonkeyName]
involved (Const _) = []
involved Variable = []
involved (Add n1 n2) = [n1, n2]
involved (Subtract n1 n2) = [n1, n2]
involved (Multiply n1 n2) = [n1, n2]
involved (Divide n1 n2) = [n1, n2]
involved (Equals n1 n2) = [n1, n2]

type MonkeyMap = Map MonkeyName NameOperation

initMonkeyMap :: Input -> MonkeyMap
initMonkeyMap = Map.fromList

monkeys :: MonkeyMap -> [MonkeyName]
monkeys = nub . concatMap allNames . Map.toList
 where
  allNames (n, op) = n : involved op

yellMap :: MonkeyMap -> Map MonkeyName Number
yellMap mmap = ymap
 where
  ymap = Map.fromList [(monkeyName, calcYell monkeyName) | monkeyName <- monkeys mmap]
  calcYell monkeyName =
    case operation of
      Const n -> n
      Variable -> error "this one does need to think"
      Add n1 n2 -> ymap Map.! n1 + ymap Map.! n2
      Subtract n1 n2 -> ymap Map.! n1 - ymap Map.! n2
      Multiply n1 n2 -> ymap Map.! n1 * ymap Map.! n2
      Divide n1 n2 -> ymap Map.! n1 / ymap Map.! n2
      Equals _ _ -> error "should not equal"
   where
    operation = mmap Map.! monkeyName

data Ast
  = AConst Number
  | AVariable
  | AAdd Ast Ast
  | ASubtract Ast Ast
  | AMultiply Ast Ast
  | ADivide Ast Ast
  | AEquals Ast Ast
  deriving (Show, Eq)

astMap :: MonkeyMap -> Map MonkeyName Ast
astMap mmap = aMap
 where
  aMap = Map.fromList [(monkeyName, calcAst monkeyName) | monkeyName <- monkeys mmap]
  calcAst :: MonkeyName -> Ast
  calcAst monkeyName =
    case operation of
      Const n -> AConst n
      Variable -> AVariable
      Add n1 n2 -> simplify $ AAdd (aMap Map.! n1) (aMap Map.! n2)
      Subtract n1 n2 -> simplify $ ASubtract (aMap Map.! n1) (aMap Map.! n2)
      Multiply n1 n2 -> simplify $ AMultiply (aMap Map.! n1) (aMap Map.! n2)
      Divide n1 n2 -> simplify $ ADivide (aMap Map.! n1) (aMap Map.! n2)
      Equals n1 n2 -> simplify $ AEquals (aMap Map.! n1) (aMap Map.! n2)
   where
    operation = mmap Map.! monkeyName

simplify :: Ast -> Ast
simplify ast =
  let ast' = go ast
   in if ast' == ast then ast else simplify ast'
 where
  go c@((AConst _)) = c
  go v@AVariable = v
  go ((AAdd ((AConst a)) ((AConst b)))) = AConst (a + b)
  go ((ASubtract ((AConst a)) ((AConst b)))) = AConst (a - b)
  go ((AMultiply ((AConst a)) ((AConst b)))) = AConst (a * b)
  go ((ADivide ((AConst a)) ((AConst b)))) = AConst (a / b)
  go other = other

getRootEq :: Input -> Equation
getRootEq inp =
  case aMap Map.! "root" of
    AEquals l r -> (l, r)
    _ -> error "no equation"
 where
  aMap = astMap . adjustPart2 $ initMonkeyMap inp

type Equation = (Ast, Ast)

solve :: Equation -> Equation
solve eq =
  let eq' = go eq
   in if eq' == eq then eq else solve eq'
 where
  go (c@(AConst _), other) = solve (other, c)
  go (ADivide l (AConst d), AConst r) = (l, AConst (r * d))
  go (AAdd l (AConst a), AConst r) = (l, AConst (r - a))
  go (AAdd (AConst a) l, AConst r) = (l, AConst (r - a))
  go (ASubtract l (AConst a), AConst r) = (l, AConst (r + a))
  go (ASubtract (AConst a) l, AConst r) = (l, AConst (a - r))
  go (AMultiply (AConst a) l, AConst r) = (l, AConst (r / a))
  go (AMultiply l (AConst a), AConst r) = (l, AConst (r / a))
  go other = other

----------------------------------------------------------------------
-- load and parse input

loadInput :: IO Input
loadInput = loadFile "input.txt"

loadExample :: IO Input
loadExample = loadFile "example.txt"

loadFile :: FilePath -> IO Input
loadFile file = do
  txt <- readFile ("./src/Y" <> show yearNr <> "/Day" <> show dayNr <> "/" <> file)
  pure $ parseText txt

parseText :: String -> Input
parseText = map (runParser monkeyP) . lines

monkeyP :: Parser Monkey
monkeyP = do
  n <- monkeyNameP <* PC.string ": "
  act <- operationP
  pure (n, act)

monkeyNameP :: Parser MonkeyName
monkeyNameP = nameP

operationP :: Parser NameOperation
operationP =
  P.choice [Const . toRational <$> (numberP :: Parser Int), binOperationP]

binOperationP :: Parser NameOperation
binOperationP = do
  n1 <- monkeyNameP
  op <- opP
  op n1 <$> monkeyNameP
 where
  opP =
    P.choice
      [ Add <$ PC.string " + "
      , Subtract <$ PC.string " - "
      , Multiply <$ PC.string " * "
      , Divide <$ PC.string " / "
      ]