{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ViewPatterns      #-}

module Main where

import           Control.Arrow                         ((***))
import           Data.Maybe                            (fromJust)
import           Data.Text.Prettyprint.Doc.Render.Text (putDoc)
import           System.Environment                    (getArgs)

import           Language.Egison.AST
import           Language.Egison.Parser                (readUTF8File, removeShebang)
import           Language.Egison.Parser.SExpr
import           Language.Egison.Pretty


exprInfix :: [(String, Op)]
exprInfix =
  [ ("**",        Op "^"  8 InfixL False)
  , ("**'",       Op "^'" 8 InfixL False)
  , ("*",         Op "*"  7 InfixL False)
  , ("/",         Op "/"  7 InfixL False)
  , ("*'",        Op "*'" 7 InfixL False)
  , ("/'",        Op "/'" 7 InfixL False)
  , (".",         Op "."  7 InfixL False) -- tensor multiplication
  , (".'",        Op ".'" 7 InfixL False) -- tensor multiplication
  , ("remainder", Op "%"  7 InfixL False) -- primitive function
  , ("+",         Op "+"  6 InfixL False)
  , ("-",         Op "-"  6 InfixL False)
  , ("+'",        Op "+'" 6 InfixL False)
  , ("-'",        Op "-'" 6 InfixL False)
  , ("append",    Op "++" 5 InfixR False)
  , ("cons",      Op "::" 5 InfixR False)
  , ("equal",     Op "="  4 InfixL False) -- primitive function
  , ("lte",       Op "<=" 4 InfixL False) -- primitive function
  , ("gte",       Op ">=" 4 InfixL False) -- primitive function
  , ("lt",        Op "<"  4 InfixL False) -- primitive function
  , ("gt",        Op ">"  4 InfixL False) -- primitive function
  , ("&&",        Op "&&" 3 InfixR False)
  , ("and",       Op "&&" 3 InfixR False)
  , ("||",        Op "||" 2 InfixR False)
  , ("or",        Op "||" 2 InfixR False)
  , ("apply",     Op "$"  0 InfixR False)
  ]

patternInfix :: [(String, Op)]
patternInfix =
  [ ("^",    Op "^"  8 InfixL False)  -- PowerPat
  , ("*",    Op "*"  7 InfixL False)  -- MultPat
  , ("div",  Op "/"  7 InfixL False)  -- DivPat
  , ("+",    Op "+"  6 InfixL False)  -- PlusPat
  , ("cons", Op "::" 5 InfixR False)
  , ("join", Op "++" 5 InfixR False)
  , ("&",    Op "&"  3 InfixR False)
  , ("|",    Op "|"  2 InfixR False)
  ]

lookupVarExprInfix :: String -> Maybe Op
lookupVarExprInfix x = lookup x exprInfix

class SyntaxElement a where
  toNonS :: a -> a

instance SyntaxElement TopExpr where
  toNonS (Define x y)   = Define (toNonS x) (toNonS y)
  toNonS (Test x)       = Test (toNonS x)
  toNonS (Execute x)    = Execute (toNonS x)
  toNonS x              = x

instance SyntaxElement Expr where
  toNonS (VarExpr (lookupVarExprInfix -> Just op)) =
    SectionExpr op Nothing Nothing
  toNonS (VarExpr x) = VarExpr x

  toNonS (IndexedExpr b x ys)  = IndexedExpr  b (toNonS x) (map toNonS ys)
  toNonS (SubrefsExpr b x y)   = SubrefsExpr  b (toNonS x) (toNonS y)
  toNonS (SuprefsExpr b x y)   = SuprefsExpr  b (toNonS x) (toNonS y)
  toNonS (UserrefsExpr b x y)  = UserrefsExpr b (toNonS x) (toNonS y)
  toNonS (TupleExpr xs)      = TupleExpr (map toNonS xs)
  toNonS (CollectionExpr xs) = CollectionExpr (map toNonS xs)
  toNonS (ConsExpr x xs) = InfixExpr cons (toNonS x) (toNonS xs)
    where cons = fromJust $ lookup "cons" exprInfix
  toNonS (JoinExpr x xs) = InfixExpr append (toNonS x) (toNonS xs)
    where append = fromJust $ lookup "append" exprInfix
  toNonS (HashExpr xs)       = HashExpr (map (toNonS *** toNonS) xs)
  toNonS (VectorExpr xs)     = VectorExpr (map toNonS xs)

  toNonS (LambdaExpr xs e)        = LambdaExpr xs (toNonS e)
  toNonS (MemoizedLambdaExpr xs e)  = MemoizedLambdaExpr xs (toNonS e)
  toNonS (CambdaExpr x e)           = CambdaExpr x (toNonS e)
  toNonS (PatternFunctionExpr xs p) = PatternFunctionExpr xs (toNonS p)

  toNonS (IfExpr x y z)         = IfExpr (toNonS x) (toNonS y) (toNonS z)
  toNonS (LetRecExpr xs y)      = LetRecExpr (map toNonS xs) (toNonS y)
  toNonS (WithSymbolsExpr xs y) = WithSymbolsExpr xs (toNonS y)

  toNonS (MatchExpr pmmode m p xs)    = MatchExpr pmmode (toNonS m) (toNonS p) (map toNonS xs)
  toNonS (MatchAllExpr pmmode m p xs) = MatchAllExpr pmmode (toNonS m) (toNonS p) (map toNonS xs)
  toNonS (MatchLambdaExpr p xs)       = MatchLambdaExpr    (toNonS p) (map toNonS xs)
  toNonS (MatchAllLambdaExpr p xs)    = MatchAllLambdaExpr (toNonS p) (map toNonS xs)

  toNonS (MatcherExpr xs) = MatcherExpr (map toNonS xs)
  toNonS (AlgebraicDataMatcherExpr xs) =
    AlgebraicDataMatcherExpr (map (\(s, es) -> (s, map toNonS es)) xs)

  toNonS (QuoteExpr x)        = QuoteExpr (toNonS x)
  toNonS (QuoteSymbolExpr x)  = QuoteSymbolExpr (toNonS x)
  toNonS (WedgeApplyExpr (VarExpr (lookupVarExprInfix -> Just op)) (y:ys)) =
    optimize $ foldl (\acc x -> InfixExpr op' acc (toNonS x)) (toNonS y) ys
      where
        op' = op { isWedge = True }

        optimize (InfixExpr Op{ repr = "*" } (ConstantExpr (IntegerExpr (-1))) e2) =
          PrefixExpr "-" (optimize e2)
        optimize (InfixExpr op e1 e2) =
          InfixExpr op (optimize e1) (optimize e2)
        optimize e = e
  toNonS (WedgeApplyExpr x ys) = WedgeApplyExpr (toNonS x) (map toNonS ys)

  toNonS (DoExpr xs y) = DoExpr (map toNonS xs) (toNonS y)

  toNonS (SeqExpr e1 e2) = SeqExpr (toNonS e1) (toNonS e2)
  toNonS (ApplyExpr (VarExpr (lookupVarExprInfix -> Just op)) (y:ys)) =
    optimize $ foldl (\acc x -> InfixExpr op acc (toNonS x)) (toNonS y) ys
      where
        optimize (InfixExpr Op{ repr = "*" } (ConstantExpr (IntegerExpr (-1))) e2) =
          PrefixExpr "-" (optimize e2)
        optimize (InfixExpr op e1 e2) =
          InfixExpr op (optimize e1) (optimize e2)
        optimize e = e

  toNonS (ApplyExpr x ys) = ApplyExpr (toNonS x) (map toNonS ys)
  toNonS (CApplyExpr e1 e2) = CApplyExpr (toNonS e1) (toNonS e2)
  toNonS (AnonParamFuncExpr n e) =
    case AnonParamFuncExpr n (toNonS e) of
      AnonParamFuncExpr 2 (InfixExpr op (AnonParamExpr 1) (AnonParamExpr 2)) ->
        SectionExpr op Nothing Nothing
      -- TODO(momohatt): Check if %1 does not appear freely in e
      -- AnonParamFuncExpr 1 (InfixExpr op e (AnonParamExpr 1)) ->
      --   SectionExpr op (Just (toNonS e)) Nothing
      -- AnonParamFuncExpr 1 (InfixExpr op (AnonParamExpr 1) e) ->
      --   SectionExpr op Nothing (Just (toNonS e))
      e' -> e'

  toNonS (GenerateTensorExpr e1 e2) = GenerateTensorExpr (toNonS e1) (toNonS e2)
  toNonS (TensorExpr e1 e2) = TensorExpr (toNonS e1) (toNonS e2)
  toNonS (TensorContractExpr e1) = TensorContractExpr (toNonS e1)
  toNonS (TensorMapExpr e1 e2) = TensorMapExpr (toNonS e1) (toNonS e2)
  toNonS (TensorMap2Expr e1 e2 e3) = TensorMap2Expr (toNonS e1) (toNonS e2) (toNonS e3)
  toNonS (TransposeExpr e1 e2) = TransposeExpr (toNonS e1) (toNonS e2)
  toNonS (FlipIndicesExpr _) = error "Not supported: FlipIndicesExpr"

  toNonS x = x

instance SyntaxElement Pattern where
  toNonS (ValuePat e) = ValuePat (toNonS e)
  toNonS (PredPat e) = PredPat (toNonS e)
  toNonS (IndexedPat p es) = IndexedPat (toNonS p) (map toNonS es)
  toNonS (LetPat binds pat) = LetPat (map toNonS binds) (toNonS pat)
  toNonS (InfixPat op p1 p2) = InfixPat op (toNonS p1) (toNonS p2)
  toNonS (NotPat p) = NotPat (toNonS p)
  toNonS (AndPat p1 p2) = InfixPat op (toNonS p1) (toNonS p2)
    where op = fromJust $ lookup "&" patternInfix
  toNonS (OrPat p1 p2) = InfixPat op (toNonS p1) (toNonS p2)
    where op = fromJust $ lookup "|" patternInfix
  toNonS ForallPat{} = error "Not supported: forall pattern"
  toNonS (TuplePat ps) = TuplePat (map toNonS ps)
  toNonS (InductivePat ((`lookup` patternInfix) -> Just op) [p1, p2]) =
    InfixPat op (toNonS p1) (toNonS p2)
  toNonS (InductivePat name ps) = InductivePat name (map toNonS ps)
  toNonS (LoopPat i range p1 p2) = LoopPat i (toNonS range) (toNonS p1) (toNonS p2)
  toNonS (PApplyPat e p) = PApplyPat (toNonS e) (map toNonS p)
  toNonS (SeqConsPat p1 p2) = SeqConsPat (toNonS p1) (toNonS p2)
  toNonS (DApplyPat p ps) = DApplyPat (toNonS p) (map toNonS ps)
  toNonS p = p

instance SyntaxElement PrimitivePatPattern where
  toNonS (PPInductivePat x pps) = PPInductivePat x (map toNonS pps)
  toNonS (PPTuplePat pps) = PPTuplePat (map toNonS pps)
  toNonS pp = pp

instance SyntaxElement PrimitiveDataPattern where
  toNonS (PDInductivePat x pds) = PDInductivePat x (map toNonS pds)
  toNonS (PDTuplePat pds) = PDTuplePat (map toNonS pds)
  toNonS (PDConsPat pd1 pd2) = PDConsPat (toNonS pd1) (toNonS pd2)
  toNonS (PDSnocPat pd1 pd2) = PDSnocPat (toNonS pd1) (toNonS pd2)
  toNonS pd = pd

instance SyntaxElement LoopRange where
  toNonS (LoopRange e1 e2 p) = LoopRange (toNonS e1) (toNonS e2) (toNonS p)

instance SyntaxElement a => SyntaxElement (IndexExpr a) where
  toNonS script = toNonS <$> script

instance SyntaxElement BindingExpr where
  toNonS (Bind pdp x) = Bind (toNonS pdp) (toNonS x)
  toNonS (BindWithIndices var x) = BindWithIndices var (toNonS x)

instance SyntaxElement MatchClause where
  toNonS (pat, body) = (toNonS pat, toNonS body)

instance SyntaxElement PatternDef where
  toNonS (x, y, zs) = (toNonS x, toNonS y, map (\(z, w) -> (toNonS z, toNonS w)) zs)

instance SyntaxElement VarWithIndices where
  toNonS = id

main :: IO ()
main = do
  args <- getArgs
  input <- readUTF8File $ head args
  -- 'ast' is not desugared
  let ast = parseTopExprs (removeShebang True input)
  case ast of
    Left err ->
      print err
    Right ast -> do
      putStrLn "--"
      putStrLn "-- This file has been auto-generated by egison-translator."
      putStrLn "--"
      putStrLn ""
      putDoc $ prettyTopExprs $ map toNonS ast
      putStrLn ""