------------------------------------------------------------------------------
-- | Author : Michael Hanus, Steven Libby
--   Version: September 2025
--
-- Implementation of transforming FlatCurry expressions by applying
-- deterministically defined expressions transformations as long as possible.
------------------------------------------------------------------------------

module FlatCurry.Transform.ExecDet
  ( transformFuncsInProgDet
  , transformExprDet, transformExprMaxDet, showTransformExprDet )
 where

import Data.Tuple.Extra   ( second )
import FlatCurry.Goodies  ( updFuncBody, updProgFuncs )
import FlatCurry.Types
import FlatCurry.Pretty   ( ppExp, Options(..), QualMode(..) )
import Text.Pretty        ( pPrint )

import FlatCurry.Transform.Types
import FlatCurry.Transform.Utils ( ReWriter(..)
                                 , curVar, newVar, replace, update )

------------------------------------------------------------------------------
-- | Transforms the bodies of all functions in a FlatCurry program according
--   to some deterministic(!) expression transformation provided as
--   the first argument.
transformFuncsInProgDet :: ExprTransformationDet -> Prog -> Prog
transformFuncsInProgDet trans =
  updProgFuncs (map (updFuncBody (transformExprDet trans)))

-- | Transforms an expression by applying some deterministic(!)
--   expression transformation as long as possible with a bottom-up strategy.
transformExprDet :: ExprTransformationDet -> Expr -> Expr
transformExprDet = transformExprMaxDet (-1)

-- | The same as 'transformExprDet' but takes the maximum number of
--   transformation steps to be applied as a further argument.
--   If the number is negative, then keep going until no transformation
--   can be applied.
transformExprMaxDet :: Int -> ExprTransformationDet -> Expr -> Expr
transformExprMaxDet n trans e = fst (runTrExpr trans n (newVar e) e)

-- | The same as 'transformExprMaxDet' but returns also a formatted trace of
--   all applied transformation steps as well as its total number.
showTransformExprDet :: Int -> ExprTransformationDet
                     -> Expr -> (Expr,String,Int)
showTransformExprDet n trans e 
  = let (e',steps) = runTrExpr trans n (newVar e) e
    in (e', showTransSteps e steps, length steps)

runTrExpr :: ExprTransformationDet -> Int -> VarIndex -> Expr -> (Expr,[Step])
runTrExpr trans n v e
 | n == 0    = (e,[])
 | otherwise = let (e', s, v', seen) = runRewriter (run trans [] e) v
               in case seen of
                    False -> (e', s)
                    True  -> second (s++) $ runTrExpr trans (n-1) v' e'

run :: ExprTransformationDet -> Path -> Expr -> ReWriter Expr
run _ _ e@(Var _) = return e
run _ _ e@(Lit _) = return e

run trans p (Comb ct n es) = do es' <- mapM runExp (zip [0..] es)
                                runExprTransform trans p (Comb ct n es')
 where runExp (i,e) = run trans (i:p) e

run trans p (Let bs e) = do e'  <- run trans (-1:p) e
                            bs' <- mapM runVar (zip [0..] bs)
                            runExprTransform trans p (Let bs' e')
 where runVar (n,(v,tv,be)) = do be' <- run trans (n:p) be
                                 return (v,tv,be')

run trans p (Free vs e) = do e' <- run trans (0:p) e
                             runExprTransform trans p (Free vs e')

run trans p (Or e1 e2) = do e1' <- run trans (0:p) e1
                            e2' <- run trans (1:p) e2
                            runExprTransform trans p (Or e1' e2')

run trans p (Case ct e bs) = do e'  <- run trans (-1:p) e
                                bs' <- mapM runBranch (zip [0..] bs)
                                runExprTransform trans p (Case ct e' bs')
 where runBranch (n,Branch q b) = do b' <- run trans (n:p) b
                                     return (Branch q b')

run trans p (Typed e te) = do e' <- run trans (0:p) e
                              runExprTransform trans p (Typed e' te)

-- Apply a totally defined expression transformation to an expression.
runExprTransform :: ExprTransformationDet -> Path -> Expr -> ReWriter Expr
runExprTransform trans p e = do
  v <- curVar
  case trans (v,p) e of
    Nothing        -> return e
    Just (e',r,dv) -> do update e' (r,p,e') dv
                         run trans p e'

showTransSteps :: Expr -> [Step] -> String
showTransSteps _ [] = ""
showTransSteps e ((rule, p, rhs):steps) =
  "=> " ++ rule ++ " " ++ show (reverse p) ++ "\n" ++
  pPrint (ppExp (Options 2 QualNone "") e') ++ "\n" ++
  showTransSteps e' steps
 where
  e' = replace e (reverse p) rhs

------------------------------------------------------------------------------