import Control.Monad (foldM)

import Data.Char (isAlpha, isDigit, isSpace, toLower)

import Data.List (intercalate)

import Data.Map qualified as Map

import Text.Read (readMaybe)



-- Типи даних

data Token

  = TNumber Double

  | TVariable String

  | TOperator Char

  | TAssign

  | TParenOpen

  | TParenClose

  | TEnd

  | TQuestion String

  deriving (Show, Eq)



type ValueTable = Map.Map String Double



type ExprTable = Map.Map String [Token]



data Expression

  = Expr [Token]

  | Assignment String [Token] -- Імʼя змінної, токени виразу

  | Query String -- Змінна для пошуку

  deriving (Show)



precedence :: Char -> Int

precedence '+' = 1

precedence '-' = 1

precedence '*' = 2

precedence '/' = 2

precedence '_' = 3 -- унарний мінус

precedence _ = 0



isOperator :: Char -> Bool

isOperator c = c `elem` "+-*/_"



-- Лексер

tokenize :: String -> Either String [Token]

tokenize input = tokenizeHelper (filter (not . isSpace) input) []

  where

    tokenizeHelper [] acc = Right (reverse (TEnd : acc))

    tokenizeHelper (c : cs) acc

      | isAlpha c =

          let var = [toLower c]

           in tokenizeHelper cs (TVariable var : acc)

      | isDigit c || c == '.' =

          case parseNumber (c : cs) of

            (num, rest) -> case readMaybe num of

              Just n -> tokenizeHelper rest (TNumber n : acc)

              Nothing -> Left $ "Invalid number: " ++ num

      | c `elem` "+-*/" = tokenizeHelper cs (TOperator c : acc)

      | c == '=' = tokenizeHelper cs (TAssign : acc)

      | c == '(' = tokenizeHelper cs (TParenOpen : acc)

      | c == ')' = tokenizeHelper cs (TParenClose : acc)

      | c == '?' =

          case cs of

            (v : rest)

              | isAlpha v ->

                  Right (reverse (TQuestion [toLower v] : acc))

            _ -> Left "Invalid ?variable syntax"

      | otherwise = Left $ "Unknown character: " ++ [c]



parseNumber :: String -> (String, String)

parseNumber s =

  let (numPart, rest) = span (\c -> isDigit c || c == '.') s

   in (numPart, rest)



-- Уточнення токенів (обробка неявного множення та унарного мінусування)

refineTokens :: [Token] -> [Token]

refineTokens tokens = refineHelper tokens []

  where

    refineHelper [] acc = reverse acc

    refineHelper [t] acc = reverse (t : acc)

    refineHelper (t1 : t2 : ts) acc

      | needsMultiplication t1 t2 =

          refineHelper (t2 : ts) (TOperator '*' : t1 : acc)

      | isUnaryMinus t1 t2 =

          refineHelper (TOperator '_' : ts) (TNumber 0 : t1 : acc)

      | otherwise = refineHelper (t2 : ts) (t1 : acc)



needsMultiplication :: Token -> Token -> Bool

needsMultiplication (TNumber _) TParenOpen = True

needsMultiplication TParenClose (TNumber _) = True

needsMultiplication (TNumber _) (TVariable _) = True

needsMultiplication (TVariable _) (TNumber _) = True

needsMultiplication (TVariable _) (TVariable _) = True

needsMultiplication TParenClose (TVariable _) = True

needsMultiplication (TVariable _) TParenOpen = True

needsMultiplication TParenClose TParenOpen = True

needsMultiplication _ _ = False



isUnaryMinus :: Token -> Token -> Bool

isUnaryMinus prev (TOperator '-') =

  case prev of

    TNumber _ -> False

    TVariable _ -> False

    TParenClose -> False

    _ -> True

isUnaryMinus _ _ = False



-- Парсер

parseExpression :: [Token] -> Either String Expression

parseExpression tokens =

  let refined = refineTokens tokens

   in case refined of

        (TVariable var : TAssign : rest) ->

          Right $ Assignment var (init rest)

        (TQuestion var : _) ->

          Right $ Query var

        _

          | TAssign `elem` init refined ->

              Right $ Expr (init refined)

        _ -> Right $ Expr (init refined)



expandVariables :: [Token] -> ExprTable -> ValueTable -> Either String [Token]

expandVariables tokens varExpressions variables = expandHelper tokens

  where

    expandHelper [] = Right []

    expandHelper (TVariable var : ts) = do

      expanded <- expandVariable var

      rest <- expandHelper ts

      Right (expanded ++ rest)

    expandHelper (t : ts) = do

      rest <- expandHelper ts

      Right (t : rest)



    expandVariable var =

      case Map.lookup var variables of

        Just val -> Right [TNumber val]

        Nothing -> case Map.lookup var varExpressions of

          Just expr -> expandHelper expr

          Nothing -> Right [TVariable var]



popHigherPrecedence :: Char -> [Token] -> [Token] -> ([Token], [Token])

popHigherPrecedence op output [] = (output, [])

popHigherPrecedence op output (TOperator topOp : ops)

  | precedence topOp >= precedence op =

      popHigherPrecedence op (TOperator topOp : output) ops

  | otherwise = (output, TOperator topOp : ops)

popHigherPrecedence op output (t : ops) = (output, t : ops)



popUntilParen :: [Token] -> [Token] -> Maybe ([Token], [Token])

popUntilParen output [] = Nothing

popUntilParen output (TParenOpen : ops) = Just (output, ops)

popUntilParen output (op : ops) =

  case popUntilParen (op : output) ops of

    Just (newOutput, newOps) -> Just (newOutput, newOps)

    Nothing -> Nothing



evaluate :: [Token] -> Either String Double

evaluate tokens = evaluateHelper tokens []

  where

    evaluateHelper [] [result] = Right result

    evaluateHelper [] _ = Left "Invalid expression"

    evaluateHelper (TNumber n : ts) stack = evaluateHelper ts (n : stack)

    evaluateHelper (TOperator op : ts) (b : a : stack) =

      case applyOp op a b of

        Just result -> evaluateHelper ts (result : stack)

        Nothing -> Left $ "Invalid operation: " ++ [op]

    evaluateHelper (TOperator '_' : ts) (a : stack) =

      evaluateHelper ts ((-a) : stack)

    evaluateHelper _ _ = Left "Invalid expression structure"



applyOp :: Char -> Double -> Double -> Maybe Double

applyOp '+' a b = Just (a + b)

applyOp '-' a b = Just (a - b)

applyOp '*' a b = Just (a * b)

applyOp '/' a b = if b /= 0 then Just (a / b) else Nothing

applyOp '_' a b = Just (a - b)

applyOp _ _ _ = Nothing



data CalculatorState = CalculatorState

  { variables :: ValueTable,

    variableExpressions :: ExprTable,

    expressions :: [Expression],

    queries :: [String]

  }

  deriving (Show)



initialState :: CalculatorState

initialState = CalculatorState Map.empty Map.empty [] []



processInput :: [String] -> Either String CalculatorState

processInput inputs = do

  state <- foldM collectExpressions initialState inputs

  finalState <- resolveVariables state

  Right finalState

  where

    collectExpressions state input = do

      tokens <- tokenize input

      expr <- parseExpression tokens

      case expr of

        Assignment var exprTokens ->

          Right $ state {variableExpressions = Map.insert var exprTokens (variableExpressions state)}

        Query var ->

          Right $ state {queries = var : queries state}

        Expr exprTokens ->

          Right $ state {expressions = expr : expressions state}



-- Алгоритм для перетворення інфіксу в постфікс із заміною змінних

shuntingYard :: [Token] -> Either String [Token]

shuntingYard tokens = shuntingYardHelper tokens [] []

  where

    shuntingYardHelper [] output [] = Right (reverse output)

    shuntingYardHelper [] output ops = Right (reverse output ++ reverse ops)

    shuntingYardHelper (t : ts) output ops =

      case t of

        TNumber _ -> shuntingYardHelper ts (t : output) ops

        TVariable var -> Left $ "Unresolved variable: " ++ var

        TOperator op ->

          let (newOutput, newOps) = popHigherPrecedence op output ops

           in shuntingYardHelper ts newOutput (TOperator op : newOps)

        TParenOpen -> shuntingYardHelper ts output (t : ops)

        TParenClose ->

          case popUntilParen output ops of

            Just (newOutput, newOps) -> shuntingYardHelper ts newOutput newOps

            Nothing -> Left "Mismatched parentheses"

        _ -> shuntingYardHelper ts output ops



-- Розв'язати всі змінні, обчисливши їхні вирази

resolveVariables :: CalculatorState -> Either String CalculatorState

resolveVariables state = do

  resolvedVars <- resolveAllVariables (variableExpressions state) Map.empty 10 -- Ліміт ітерації

  Right $ state {variables = resolvedVars}

  where

    resolveAllVariables :: ExprTable -> ValueTable -> Int -> Either String ValueTable

    resolveAllVariables varExprs resolved 0 = Right resolved

    resolveAllVariables varExprs resolved maxIter

      | Map.null varExprs = Right resolved

      | otherwise = do

          (newResolved, remaining) <- resolveOnePass varExprs resolved

          if Map.size newResolved == Map.size resolved

            then Right resolved -- деякі змінні можуть залишатися символічними

            else resolveAllVariables remaining newResolved (maxIter - 1)



    resolveOnePass :: ExprTable -> ValueTable -> Either String (ValueTable, ExprTable)

    resolveOnePass varExprs resolved = do

      let attempts = Map.toList varExprs

      foldM tryResolve (resolved, varExprs) attempts



    tryResolve :: (ValueTable, ExprTable) -> (String, [Token]) -> Either String (ValueTable, ExprTable)

    tryResolve (resolved, remaining) (var, tokens) =

      case expandVariables tokens remaining resolved >>= shuntingYard >>= evaluate of

        Right value -> Right (Map.insert var value resolved, Map.delete var remaining)

        Left _ -> Right (resolved, remaining)



runCalculator :: [String] -> IO ()

runCalculator inputs =

  case processInput inputs of

    Left err -> putStrLn $ "Error: " ++ err

    Right finalState -> do

      -- Обробити запити

      mapM_ (processQuery finalState) (reverse $ queries finalState)

      -- Обробити вирази

      mapM_ (processExpression finalState) (zip [1 ..] (reverse $ expressions finalState))



processQuery :: CalculatorState -> String -> IO ()

processQuery state var =

  case Map.lookup var (variables state) of

    Just val -> putStrLn $ var ++ " = " ++ show val

    Nothing ->

      case Map.lookup var (variableExpressions state) of

        Just tokens -> do

          case expandVariables tokens (variableExpressions state) (variables state) of

            Right expandedTokens -> putStrLn $ var ++ " = " ++ showTokens expandedTokens

            Left err -> putStrLn $ "Error evaluating " ++ var ++ ": " ++ err

        Nothing -> putStrLn $ var ++ " = " ++ var



showTokens :: [Token] -> String

showTokens tokens = unwords $ map showToken tokens

  where

    showToken (TNumber n) = show n

    showToken (TVariable v) = v

    showToken (TOperator op) = [op]

    showToken TParenOpen = "("

    showToken TParenClose = ")"

    showToken _ = ""



processExpression :: CalculatorState -> (Int, Expression) -> IO ()

processExpression state (num, Expr tokens) =

  case expandVariables tokens (variableExpressions state) (variables state) >>= shuntingYard >>= evaluate of

    Left err -> putStrLn $ "Error in expression " ++ show num ++ ": " ++ err

    Right result -> putStrLn $ "The result of expression #" ++ show num ++ ": " ++ show result

processExpression _ _ = return ()



testInput0 = ["((51+28)*56) -7="]



testInput1 =

  [ "A = 5 * 6 + B * (C + D)",

    "C = 3.2 - D",

    "D = S +8",

    "X = S",

    "S = 18",

    "?X",

    "?A"

  ]



testInput2 =

  [ "((a+b)*c)-7=",

    "a=3*b",

    "b=-(5/5)",

    "c=a*b",

    "?c",

    "?a",

    "?b"

  ]



main :: IO ()

main = do

  putStrLn "Running calculator with test input:"

  mapM_ putStrLn testInput1

  putStrLn "\nResults:"

  runCalculator testInput1



  putStrLn $ "\n" ++ replicate 40 '-' ++ "\n"



  putStrLn "For interactive use, call runCalculator with your input list"

  putStrLn "Example: runCalculator [\"x=5+3\", \"y=x*2\", \"?y\"]"



--------



Ось я розробив калькулятор, що працює з такими запитами як в testInput.

Тепер розпиши відгук для лабораторної роботи про те, за яким принципом працює ця програма, з яких частин складається, які є цікаві деталі реалізації