{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ViewPatterns #-}
import Grammar (BinOp, Identifier, Type, HashMap)
import Typed
import Control.Comonad
import Control.Applicative (liftA2)

data Expression t
    = BinOp {lhs :: Expression t, op :: BinOp, rhs :: Expression t, typ :: t}
    | MethodCall (Expression t) Identifier [Expression t] t
    | LitInt Int t
    | LitString String t
    | LitTrue t
    | LitFalse t
    | Identifier {id :: Identifier, typ :: t}
    | LitThis t
    | LitNull t
    | New {obj ::Identifier, typ :: t}
    | Not {expr :: (Expression t), typ :: t}
    | Block [Expression t] t
    | If { predicate :: Expression t
         , body :: Expression t
         , elseBody :: (Maybe (Expression t))
         , typ :: t
         }
    | While { predicate :: Expression t, body :: Expression t, typ :: t }
    | Println { expr :: Expression t, typ :: t }
    | Assign { var :: Identifier, expr :: Expression t, typ :: t }
    | Lambda { var :: Identifier
             , argType :: Type
             , expr :: Expression t
             , returnType :: (Maybe Type)
             , typ :: t
             }
    | Closure { free :: (HashMap Identifier ())
             , var :: Identifier
             , argType :: Type
             , expr :: Expression t
             , returnType :: (Maybe Type)
             , typ :: t
             }
    deriving (Eq, Show)

instance Functor Expression where
    fmap = liftW

instance Comonad Expression where
    extract (MethodCall _ _ _ t) = t
    extract (LitInt _ t) = t
    extract (LitString _ t) = t
    extract (LitTrue t) = t
    extract (LitFalse t) = t
    extract (LitThis t) = t
    extract (LitNull t) = t
    extract (Block _ t) = t
    extract exp = typ exp

    extend f exp@(MethodCall object iden params t) = MethodCall (extend f object) iden (fmap (extend f) params) (f exp)
    extend f exp@(LitInt x _) = LitInt x (f exp)
    extend f exp@(LitString x _) = LitString x (f exp)
    extend f exp@(LitTrue _) = LitTrue (f exp)
    extend f exp@(LitFalse _) = LitFalse (f exp)
    extend f exp@(LitThis _) = LitThis (f exp)
    extend f exp@(LitNull _) = LitNull (f exp)
    extend f exp@(Block x _) = Block (fmap (extend f) x) (f exp)
    extend f (Not exp t) = Not (extend f exp) (f $ Not exp t)
    extend f exp@(If pred body elseB _) = If (e pred) (e body) (fmap e elseB) (f exp)
        where e = extend f
    extend f exp@While{predicate,body} = While (e predicate) (e body) (f exp)
        where e = extend f
    extend f exp@Println{expr} = exp {expr = extend f expr, typ = f exp}
    extend f exp@Assign{expr} = exp {expr = extend f expr, typ = f exp}
    extend f exp@Lambda{expr} = exp {expr = extend f expr, typ = f exp}
    extend f exp@Closure{expr} = exp {expr = extend f expr, typ = f exp}
    extend f exp@BinOp{lhs,rhs} = exp {lhs = extend f lhs, rhs = extend f rhs, typ = f exp}
    extend f exp@(Identifier i _) = Identifier i (f exp)
    extend f exp@(New o _) = New o (f exp)

isLeaf exp = case exp of
    LitInt _ t -> True
    LitString _ t -> True
    LitTrue t -> True
    LitFalse t -> True
    Identifier _ t -> True
    LitThis t -> True
    LitNull t -> True
    New _ t -> True
    _ -> False

extractSimple exp = case exp of
    Not{expr,typ} -> Just (expr,typ)
    Println{expr,typ} -> Just (expr,typ)
    Assign{expr,typ} -> Just (expr,typ)
    Lambda{expr,typ} -> Just (expr,typ)
    Closure{expr,typ} -> Just (expr,typ)
    _ -> Nothing


instance ComonadApply Expression where
    fxp@(extract -> f) <@> exp@(extract -> x) | isLeaf exp = exp $> f x
    (BinOp lf _ rf f) <@> (BinOp lx o rx x) = BinOp (lf <@> lx) o (rf <@> rx) (f x)
    (MethodCall fo _ fp f) <@> (MethodCall xo iden xp x) = MethodCall (fo <@> xo) iden (zipWith (<@>) fp xp) (f x)
    (Not fe f) <@> (Not xe x) = Not (fe <@> xe) (f x)
    (Block fe f) <@> (Block xe x) = Block (zipWith (<@>) fe xe) (f x)
    (If fp fb fe f) <@> (If xp xb xe x) = If (fp <@> xp) (fb <@> xb) (liftA2 (<@>) fe xe) (f x)
    (While fp fb f) <@> (While xp xb x) = While (fp <@> xp) (fb <@> xb) (f x)
    (extractSimple -> Just (fe,f)) <@> (Not xe x) = Not (fe <@> xe) (f x)
    (extractSimple -> Just (fe,f)) <@> (Println xe x) = Println (fe <@> xe) (f x)
    (extractSimple -> Just (fe,f)) <@> (Assign var xe x) = Assign var (fe <@> xe) (f x)
    (extractSimple -> Just (fe,f)) <@> exp@Lambda{expr,typ} = exp {expr = fe <@> expr, typ = f typ}
    (extractSimple -> Just (fe,f)) <@> exp@Closure{expr,typ} = exp {expr = fe <@> expr, typ = f typ}

types :: Expression String -> Expression TType -> TType
types LitInt{} = const TInt
types Not{} = const TBool
types While{} = \(While _ xb _) -> extract xb
types LitTrue{} = const TBool
types (LitThis klass) = const (TClass klass)

types' :: Expression TType -> TType
types' LitInt{} = TInt
types' Not{} = TBool
types' (While _ xb _) = extract xb
types' LitTrue{} = TBool

-- ComonadInject (TM) där man kan byta översta elementet.
class Comonad w => ComonadInject w where
    inject :: a -> w a -> w a

instance ComonadInject Expression where
    inject x e | isLeaf e = x <$ e
    inject x (Not e _) = Not e x
    inject x e = e {typ = x}

--kfix :: ComonadApply w => w (w a -> a) -> w a
--kfix w = fix $ \u -> w <@> duplicate u

--kfix (extend types $ (LitInt 1 ())) :: Expression (Expression TType -> TType) =
--    fix $ \u :: Expression TType -> (extend types $ (LitInt 1 ())) <@> duplicate u

fix f = f (fix f)

-- kfix med 0 sharing.
pfix :: Comonad w => w (w a -> a) -> w a
pfix = fmap wfix . duplicate

bfix :: ComonadApply w => w (w a -> a) -> w a
bfix w = w <@> extend bfix w

bfix2 :: (ComonadApply w, ComonadInject w) => w (w a -> a) -> w a
bfix2 w = let x = w <@> inject undefined (extend bfix2 w) in x

fixtend :: ComonadApply w => (w a -> w b -> b) -> w a -> w b
fixtend f = bfix . extend f

deepen w = While w w (extract w)