{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE TypeFamilies #-}

module LamMachine where

import Data.List
import Data.Word
import Data.Int
import Data.Monoid
import Data.Maybe
--import Data.Foldable (toList)
import qualified SplayList as F
import Control.Arrow hiding ((<+>))
import Control.Category hiding ((.), id)
import Control.Monad
import Debug.Trace

import LambdaCube.Compiler.Pretty hiding (expand)

import FreeVars

--------------------------------------------------------------------- data structures

data Lit
    = LInt   !Int
    | LChar  !Char
    | LFloat !Double
    deriving Eq

data Exp
    = Var_  !VarFV
    | Lam_  VarInfo  !(LamFV Exp)
    | Let_  VarInfo  !(Pair Exp (LamFV Exp))
    | Con_  ConInfo !Int (List Exp)
    | Case_ CaseInfo !(Pair Exp (List Exp))
    | Lit  !Lit
    | Op1  !Op1 !Exp
    | Op2_ !Op2 !(Pair Exp Exp)
    deriving (Eq, Show)

data Op1 = HNFMark | YOp | Round
    deriving (Eq, Show)

data Op2 = AppOp | SeqOp | Add | Sub | Mod | LessEq | EqInt
    deriving (Eq, Show)

pattern Y a         = Op1 YOp a
pattern App a b     = Op2 AppOp a b
pattern Seq a b     = Op2 SeqOp a b
pattern Int i       = Lit (LInt i)

pattern Var i       = Var_ (VarFV i)
pattern Lam n i     = Lam_ n (LamFV i)
pattern Let n i x   <- Let_ n (Pair i (LamFV x))
  where Let _ _ (down 0 -> Just x) = x
        Let n i x   =  Let_ n (Pair i (LamFV x))
pattern Con a b i   = Con_ a b (List i)
pattern Case a b c  = Case_ a (Pair b (List c))
pattern Op2 op a b  = Op2_ op (Pair a b)

infixl 4 `App`

data EnvPiece
    = ELet_   VarInfo !Exp
    | EDLet1_ VarInfo !(LamFV DExp)
    | ECase_ CaseInfo !(List Exp)
    | EOp2_1_ !Op2 !Exp
    | EOp2_2_ !Op2 !Exp
    deriving (Eq, Show)

pattern ELet i x   = ELet_ i x
pattern EDLet1 i x = EDLet1_ i (LamFV x)
pattern ECase op b = ECase_ op (List b)
pattern EOp2_1 i x = EOp2_1_ i x
pattern EOp2_2 i x = EOp2_2_ i x

type DEnv = F.SplayList EnvPiece

pattern DEnvNil = F.Nil

pattern DEnvCons :: EnvPiece -> DEnv -> DEnv
pattern DEnvCons p e = F.Cons p e

instance F.Measured EnvPiece where
    type Measure EnvPiece = Nat
    measure = \case
        ELet_{} -> 1
        _ -> 0
--instance HasSFV DExp where


type DExp = DEnv

pattern DExpNil = DEnvNil

pattern DExpCons :: EnvPiece -> DExp -> DExp
pattern DExpCons p e = F.Snoc e p

infixr 4 `DEnvCons`, `DExpCons`

-- state of the machine
data MSt = MSt !DEnv !Exp
    deriving (Eq, Show)

--------- pretty print info

type VarInfo = String
type ConInfo = String
type CaseInfo = [(String, [String])]

--------------------------------------------------------------------- instances

instance HasFV Exp where
    fvLens = \case
        Op1    op e -> appLens (Op1 op) e
        Op2_   op e -> appLens (Op2_ op) e
        Con_ x op e -> appLens (Con_ x op) e
        Case_  op e -> appLens (Case_ op) e
        Let_   op e -> appLens (Let_ op) e
        Lam_   op e -> appLens (Lam_ op) e
        Var_      e -> appLens Var_ e
        Lit l       -> lConst $ Lit l

instance HasFV EnvPiece where
    fvLens = \case
        EOp2_1_ op e -> appLens (EOp2_1_ op) e
        EOp2_2_ op e -> appLens (EOp2_2_ op) e
        ECase_  op e -> appLens (ECase_ op) e
        ELet_   op e -> appLens (ELet_ op) e
        EDLet1_ op e -> appLens (EDLet1_ op) e

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

pattern InHNF a <- (getHNF -> Just a)
  where InHNF a@Lit{}           = a
        InHNF a@Lam{}           = a
        InHNF a@(Op1 HNFMark _) = a
        InHNF a                 = Op1 HNFMark a

getHNF x = case x of
    Lit{}         -> Just x
    Lam{}         -> Just x
    Op1 HNFMark a -> Just a
    _             -> Nothing

initSt :: Exp -> MSt
initSt e = MSt DEnvNil e

----------------------------------------------------------- machine code begins here

justResult :: MSt -> MSt
justResult st = steps st (\_ s -> justResult s) st

hnf = justResult . initSt

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

data StepTag = Begin String | End String | Step String
  deriving Show

focusNotUsed (EDLet1 _ x `DEnvCons` _) = False --not $ usedVar 0 x
focusNotUsed _ = True

inHNF :: MSt -> Bool
inHNF (MSt _ (InHNF _)) = True
inHNF _ = False

steps :: forall e . e -> (StepTag -> MSt -> e) -> MSt -> e
steps nostep bind (MSt vs e) = case e of

    e | EDLet1 _ (down 0 -> Just d) `DEnvCons` vs <- vs -> zipUp e vs d

    Con cn i xs
        | focusNotUsed vs || lz == 0 -> step (cn ++ " hnf") $ MSt vs $ InHNF e
        | otherwise -> step (cn ++ " copy") $ MSt (foldr DEnvCons vs $ ELet "" <$> zs) $ InHNF $ Con cn i xs'  -- share complex constructor arguments
      where
        lz = Nat $ length zs
        (xs', concat -> zs) = unzip $ f 0 xs
        f i [] = []
        f i (x: xs) | simple x = (up 0 lz x, []): f i xs
                    | otherwise = (Var i, [up 0 (lz - i - 1) x]): f (i+1) xs

    -- TODO: handle recursive constants
    Y x             -> step "Y"     $ MSt vs $ App x (Y x)

    Var i           -> begin "var"   $ zipDown i DExpNil vs
    Case cns a cs   -> begin "case"  $ MSt (ECase cns cs `DEnvCons` vs) a
    Op2 op a b      -> begin (show op) $ MSt (EOp2_1 op b `DEnvCons` vs) a

    InHNF a -> case vs of

        DEnvNil -> nostep

        ELet n x `DEnvCons` vs -> step "goUp" $ MSt vs $ InHNF $ Let n x $ InHNF a  -- TODO: speed up?
        x `DEnvCons` vs | Let n y e <- a -> step "pakol" $ MSt (up 0 1 x `DEnvCons` ELet n y `DEnvCons` vs) e   -- TODO: speed up
        x `DEnvCons` vs -> case x of
            ECase cns cs -> end "case" $ case a of
                Con cn i es -> MSt vs $ foldl App (cs !! i) es  -- TODO: remove unused cases
                _           -> MSt vs $ InHNF $ Case cns (InHNF a) cs
            EOp2_1 AppOp b -> end "AppOp" $ case a of
                Lam _ (down 0 -> Just x) -> MSt vs x   -- TODO: remove b
                -- TODO: optimize out Var 0 application
                Lam n x -> MSt (ELet n b `DEnvCons` vs) x
                _       -> MSt vs $ InHNF $ App (InHNF a) b
            EOp2_1 SeqOp b -> end "SeqOp" $ case a of
                Lit{}   -> MSt vs b
                Lam{}   -> MSt vs b   -- TODO: remove a
                Con{}   -> MSt vs b   -- TODO: remove a
                _       -> MSt vs $ InHNF $ Seq (InHNF a) b
            EOp2_1 op b -> next (show op) $ MSt (EOp2_2 op (InHNF a) `DEnvCons` vs) b
            EOp2_2 op b -> end (show op) $ case (b, a) of
                (Int e, Int f) -> MSt vs (int op e f)
                  where
                    int Add a b = Int $ a + b
                    int Sub a b = Int $ a - b
                    int Mod a b = Int $ a `mod` b
                    int LessEq a b = if a <= b then T else F
                    int EqInt a b = if a == b then T else F
                _ -> MSt vs $ InHNF $ Op2 op b (InHNF a)
            EDLet1 n d -> zipUp (up 0 1 a) (ELet n (InHNF a) `DEnvCons` vs) d

  where
    zipDown j e (F.split (> j) -> (t, ELet n z `DEnvCons` zs))
        = MSt (EDLet1 n (f e t) `DEnvCons` zs) z
      where
        f e t = t <> e

    zipUp x xs e = end "var" $ MSt (e <> xs) $ up 0 n x
      where
        n = F.measure e

    begin t = bind (Begin t)
    next t  = bind (Step t)
    end t   = bind (End t)
    step t  = bind (Step t)

simple = \case
    Var{} -> True
    Lit{} -> True
    _ -> False

delElem i xs = take i xs ++ drop (i+1) xs

--------------------------------------------------------------------- toolbox: pretty print

instance PShow Lit where
    pShow (LInt i) = pShow i

instance Show Lit where show = ppShow

class ViewShow a where
    viewShow :: Bool -> a -> Doc

instance ViewShow Exp where
  viewShow vi = pShow where
    pShow e = showFE green vi (fst $ fvLens e :: FV) $
      case {-if vi then Exp_ (selfContract fv) x else-} e of
        Var (Nat i)  -> DVar i
        Let n a b   -> shLet n (pShow a) $ pShow b
        Lam n e     -> shLam n $ pShow e
        Con s i xs  -> foldl DApp (text s) $ pShow <$> xs
        Lit i       -> pShow' i
        Y e         -> "Y" `DApp` pShow e
        InHNF x     -> DBrace (pShow x)
        Op1 o x     -> text (show o) `DApp` pShow x
        Op2 o x y   -> shOp2 o (pShow x) (pShow y)
        Case cn e xs -> shCase cn (pShow e) pShow xs
--        Exp_ u Var_ -> onblue $ pShow' u
        e -> error $ "pShow @Exp: " ++ show e

glueTo = DGlue (InfixR 40)

app_ (Lam _ x) _ = x
app_ x y = App x y

shCase cn e f xs = DPreOp (-20) (ComplexAtom "case" (-10) e (SimpleAtom "of"))
                        $ foldr1 DSemi [foldr DFreshName (DArr_ "->" (foldl DApp (text a) $ DVar <$> reverse [0..length n - 1]) (f $ foldl app_ b $ Var . Nat <$> [0..length n - 1])) $ Just <$> n | ((a, n), b) <- zip cn xs]

shOp2 AppOp x y = DApp x y
shOp2 SeqOp a b = DOp "`seq`" (Infix 1) a b
shOp2 EqInt x y = DOp "==" (Infix 4) x y
shOp2 Add x y   = DOp "+" (InfixL 6) x y
shOp2 o x y     = text (show o) `DApp` x `DApp` y

showFE c True fv | fv /= mempty = DGlue (InfixL 40) (c $ pShow' fv)
showFE _ _ _ = id

pShow' = pShow

instance ViewShow MSt where
    viewShow vi (MSt env e) = g (onred $ white $ pShow e) env
      where
        pShow = viewShow vi

        g y DEnvNil = y
        g y z@(DEnvCons p env) = flip g env $ case p of
            EDLet1 n x -> shLet n y (h 0 x)
            ELet n x -> shLet n (pShow x) y
            ECase cns xs -> shCase cns y pShow xs
            EOp2_1 op x -> shOp2 op y (pShow x)
            EOp2_2 op x -> shOp2 op (pShow x) y

        h n DExpNil = onred $ white $ DVar n
        h n z@(DExpCons p (h (adj p n) -> y)) = showFE blue vi (fst $ fvLens p :: FV) $ case p of
            EDLet1 n x -> shLet n y (h 0 x)
            ELet n x -> shLet n (pShow x) y
            ECase cns xs -> shCase cns y pShow xs
            EOp2_1 op x -> shOp2 op y (pShow x)
            EOp2_2 op x -> shOp2 op (pShow x) y

        adj p n = case p of
            ELet{} -> n + 1
            _ -> n

instance PShow MSt where pShow = viewShow False


shUps a b = DPreOp (-20) (SimpleAtom $ show a) b
shUps' a x b = DPreOp (-20) (SimpleAtom $ show a ++ show x) b

shLam n b = DFreshName (Just n) $ showLam (DVar 0) b

showLam x (DFreshName u d) = DFreshName u $ showLam (DUp 0 x) d
showLam x (DLam xs y) = DLam (DSep (InfixR 11) x xs) y
showLam x y = DLam x y

shLet n a b = DFreshName (Just n) $ showLet (DLet "=" (shVar 0) $ DUp 0 a) b

showLet x (DFreshName u d) = DFreshName u $ showLet (DUp 0 x) d
showLet x (DLet' xs y) = DLet' (DSemi x xs) y
showLet x y = DLet' x y


---------------------------------------------------------------------------------------- examples

pPrint = putStrLn . ppShow

runMachinePure = pPrint . hnf

pattern F = Con "False" 0 []
pattern T = Con "True" 1 []
pattern Nil = Con "[]" 0 []
pattern Cons a b = Con "Cons" 1 [a, b]

caseBool x f t = Case [("False", []), ("True", [])] x [f, t]
caseList x n c = Case [("[]", []), ("Cons", ["c", "cs"])] x [n, c]

id_ = Lam "x" (Var 0)

if_ b t f = caseBool b f t

not_ x = caseBool x T F

test = id_ `App` id_ `App` id_ `App` id_ `App` Int 13

test' = id_ `App` (id_ `App` Int 14)

foldr_ f e = Y $ Lam "g" $ Lam "as" $ caseList (Var 0) (up 0 2 e) (Lam "x" $ Lam "xs" $ up 0 4 f `App` Var 1 `App` (Var 3 `App` Var 0))

filter_ p = foldr_ (Lam "y" $ Lam "ys" $ if_ (p `App` Var 1) (Cons (Var 1) (Var 0)) (Var 0)) Nil

and2 a b = if_ a b F

and_ = foldr_ (Lam "a" $ Lam "b" $ and2 (Var 1) (Var 0)) T

map_ f = foldr_ (Lam "z" $ Lam "zs" $ Cons (f (Var 1)) (Var 0)) Nil

neq a b = not_ $ Op2 EqInt a b

from_ = Y $ Lam "from" $ Lam "n" $ Cons (Var 0) $ Var 1 `App` Op2 Add (Var 0) (Int 1)

idx xs n = caseList xs undefined $ Lam "q" $ Lam "qs" $ if_ (Op2 EqInt n $ Int 0) (Var 1) $ idx (Var 0) (Op2 Sub n $ Int 1)

t = runMachinePure $ idx (from_ `App` Int 3) (Int 5)

takeWhile_ p xs = caseList xs Nil $ Lam "x" $ Lam "xs" $ if_ (p (Var 1)) (Cons (Var 1) $ takeWhile_ p (Var 0)) Nil

sum_ = foldr_ (Lam "a" $ Lam "b" $ Op2 Add (Var 1) (Var 0)) (Int 0)

sum' = Y $ Lam "sum" $ Lam "xs" $ caseList (Var 0) (Int 0) $ Lam "y" $ Lam "ys" $ Op2 Add (Var 1) $ Var 3 `App` Var 0

infixl 4 `sApp`

sApp a b = Lam "s" (Seq (Var 0) (a `App` Var 0)) `App` b

{-
accsum acc [] = acc
accsum acc (x: xs) = let z = acc + x `seq` accsum z xs
-}
accsum = Y $ Lam "accsum" $ Lam "acc" $ Lam "xs" $ caseList (Var 0) (Var 1) $ Lam "y" $ Lam "ys" $ Var 4 `sApp` Op2 Add (Var 3) (Var 1) `App` Var 0

fromTo = Y $ Lam "fromTo" $ Lam "begin" $ Lam "end" $ Cons (Var 1) $ if_ (Op2 EqInt (Var 0) (Var 1)) Nil $ Var 2 `App` Op2 Add (Var 1) (Int 1) `App` Var 0

t' n = sum' `App` (fromTo `App` Int 0 `App` Int n) --  takeWhile_ (\x -> Op2 LessEq x $ Int 3) (from_ `App` Int 0)

t'' n = accsum `App` Int 0 `App` (fromTo `App` Int 0 `App` Int n) --  takeWhile_ (\x -> Op2 LessEq x $ Int 3) (from_ `App` Int 0)

{- TODO

primes :: [Int]
primes = 2:3: filter (\n -> and $ map (\p -> n `mod` p /= 0) (takeWhile (\x -> x <= iSqrt n) primes)) (from 5)


main = primes !! 3000
-}

twice = Lam "f" $ Lam "x" $ Var 1 `App` (Var 1 `App` Var 0)
twice2 = Lam "f" $ Lam "x" $ Var 1 `sApp` (Var 1 `App` Var 0)

inc = Lam "n" $ Op2 Add (Var 0) (Int 1)

test'' = Lam "f" (Int 4) `App` Int 3

twiceTest n = (Lam "twice" $ (iterate (`App` Var 0) (Var 0) !! n) `App` inc `App` Int 0) `App` twice
twiceTest2 n = (Lam "twice" $ (iterate (`App` Var 0) (Var 0) !! n) `App` inc `App` Int 0) `App` twice2

tests
    =   hnf test == hnf (Int 13)
    &&  hnf test' == hnf (Int 14)
    &&  hnf test'' == hnf (Int 4)
    &&  hnf (t' 10) == hnf (Int 55)
--    &&  hnf (t'' 10) == hnf (Int 55)