{-  Copyright 2010 Dominique Devriese

    This file is part of the grammar-combinators library.

    The grammar-combinators library is free software: you can
    redistribute it and/or modify it under the terms of the GNU
    Lesser General Public License as published by the Free
    Software Foundation, either version 3 of the License, or (at
    your option) any later version.

    Foobar is distributed in the hope that it will be useful, but
    WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General
    Public License along with Foobar. If not, see
    <http://www.gnu.org/licenses/>.
-}
{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}

module Text.GrammarCombinators.Transform.UniformPaull (
    UPDomain ( UPBase, UPHead, UPTail )
  , UPBaseIx, UPHeadIx, UPTailIx
  , UPValue ( UPBV, UPHV, UPTV )
  , unUPBV, unUPHV, unUPTV
  , transformUniformPaull
  , transformUniformPaullP
  , transformUniformPaullE
  , transformUniformPaullLE
  ) where

import Text.GrammarCombinators.Base

import Control.Monad (ap, liftM2, liftM)

import Data.Maybe (isJust, fromMaybe)

import Language.Haskell.TH.Syntax

data UPBaseIx ix
data UPHeadIx ix
data UPTailIx ix 

data UPDomain phi ix where
  UPBase :: phi ix -> UPDomain phi (UPBaseIx ix)
  UPHead :: phi ix -> UPDomain phi (UPHeadIx ix)
  UPTail :: phi ix -> UPDomain phi (UPTailIx ix)

instance (FoldFam phi) => FoldFam (UPDomain phi) where
  foldFam (f :: forall ix. UPDomain phi ix -> b -> b) =
    foldFam (\idx -> f (UPBase idx) . f (UPHead idx) . f (UPTail idx))

instance (ShowFam phi) => ShowFam (UPDomain phi) where
  showIdx (UPBase idx) = showIdx idx
  showIdx (UPHead idx) = showIdx idx ++ "_head"
  showIdx (UPTail idx) = showIdx idx ++ "_tail"

instance (LiftFam phi) => LiftFam (UPDomain phi) where
  liftIdxE (UPBase idx) = AppE (ConE 'UPBase) $ liftIdxE idx
  liftIdxE (UPHead idx) = AppE (ConE 'UPHead) $ liftIdxE idx
  liftIdxE (UPTail idx) = AppE (ConE 'UPTail) $ liftIdxE idx
  liftIdxP (UPBase idx) = ConP 'UPBase [liftIdxP idx]
  liftIdxP (UPHead idx) = ConP 'UPHead [liftIdxP idx]
  liftIdxP (UPTail idx) = ConP 'UPTail [liftIdxP idx]

instance (EqFam phi) => EqFam (UPDomain phi) where
  overrideIdx f (UPBase idx) v (UPBase idx') =
    unSubVal $ overrideIdx (MkSubVal . f . UPBase) idx (MkSubVal v) idx'
  overrideIdx f (UPHead idx) v (UPHead idx') =
    unSubVal $ overrideIdx (MkSubVal . f . UPHead) idx (MkSubVal v) idx'
  overrideIdx f (UPTail idx) v (UPTail idx') =
    unSubVal $ overrideIdx (MkSubVal . f . UPTail) idx (MkSubVal v) idx'
  overrideIdx f _ _ idx' = f idx'

instance DomainMap (UPDomain phi) phi UPBaseIx where
  supIx = UPBase
  subIx (UPBase idx) = idx
  
instance DomainMap (UPDomain phi) phi UPHeadIx where
  supIx = UPHead
  subIx (UPHead idx) = idx

instance DomainMap (UPDomain phi) phi UPTailIx where
  supIx = UPTail
  subIx (UPTail idx) = idx

instance (MemoFam phi) => MemoFam (UPDomain phi) where
  data Memo (UPDomain phi) v = UPMemo (Memo phi (SubVal UPBaseIx v)) (Memo phi (SubVal UPHeadIx v)) (Memo phi (SubVal UPTailIx v)) 
  fromMemo (UPMemo mb _ _) (UPBase idx) = unSubVal $ fromMemo mb idx
  fromMemo (UPMemo _ mh _) (UPHead idx) = unSubVal $ fromMemo mh idx
  fromMemo (UPMemo _ _ mt) (UPTail idx) = unSubVal $ fromMemo mt idx 
  toMemo f = UPMemo (toMemo $ MkSubVal . f . UPBase) (toMemo $ MkSubVal . f . UPHead) (toMemo $ MkSubVal . f . UPTail)

instance Domain phi => Domain (UPDomain phi)

data family UPValue (r :: * -> *) ix
data instance UPValue r (UPBaseIx ix) = UPBV { unUPBV :: r ix } deriving (Show)
data instance UPValue r (UPHeadIx ix) = UPHV { unUPHV :: r ix } deriving (Show)
data instance UPValue r (UPTailIx ix) = UPTV { unUPTV :: r ix -> r ix }

data TransformUPIntRule p surrIx (phi :: * -> *) (r :: * -> *) v = 
  MkTUPIR {
    tlclwRecursionLimitActive :: forall ix. phi ix -> Bool,
    tlclwEmpty :: forall ix. phi ix -> Maybe (p v),
    tlclwHead :: forall ix. phi ix -> p v,
    tlclwTail :: phi surrIx -> [(Bool, p (r surrIx -> v))],
    tlclwFull :: p v
    }

newtype TransformUPWrapper p surrIx unused1 unused2 (phi :: * -> *) ixT (r :: * -> *) t v =
  MkTUPW {
    tUPRuleForGrammar ::
       TransformUPGrammar p surrIx phi ixT r t ->
       TransformUPIntRule p surrIx phi r v
    }

type TransformUPGrammar p surrIx phi ixT r t =
  forall ix. phi ix ->
  TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t (r ix)

mkSimpleTUPW :: (ProductionRule p) => p v ->
                 TransformUPWrapper p surrIx unused1 unused2 phi ixT r t v
mkSimpleTUPW r =
  MkTUPW $ \_ -> MkTUPIR (const False) (const Nothing) (const r) (const []) r 

mkEpsTUPW :: (ProductionRule p, EpsProductionRule p) => v ->
               TransformUPWrapper p surrIx unused1 unused2 phi ixT r t v
mkEpsTUPW v =
  MkTUPW $ \_ -> MkTUPIR (const False) (const $ Just $ epsilon v) (const die) (const [(True, epsilon (const v))]) $ epsilon v

mkEpsLTUPW :: (ProductionRule p, LiftableProductionRule p) => v -> Q Exp ->
              TransformUPWrapper p surrIx unused1 unused2 phi ixT r t v
mkEpsLTUPW v q =
  MkTUPW $ \_ -> MkTUPIR (const False) (const $ Just $ epsilonL v q) (const die) (const [(True, epsilonL (const v) [| const $(q) |])]) $ epsilonL v q
instance (ProductionRule p, LiftableProductionRule p) =>
         ProductionRule (TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t) where
  (ra :: TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t (a -> b)) >>>
    rb = MkTUPW $ \g ->
    let (MkTUPIR rlaa eas ha tas fa) = tUPRuleForGrammar ra g
        (MkTUPIR rlab ebs hb tbs fb) = tUPRuleForGrammar rb g
        rla :: phi ix -> Bool
        rla idx = rlaa idx || ((isJust $ eas idx) && rlab idx)
        es :: phi ix -> Maybe (p b)
        es idx = liftM2 (>>>) (eas idx) (ebs idx)
        hForEmptyA :: phi ix -> p b
        hForEmptyA idx = case eas idx of Nothing -> die
                                         Just rea -> rea >>> hb idx
        h :: phi ix -> p b
        h idx =   hForEmptyA idx
              ||| ha idx >>> fb
        ts :: phi surrIx -> [(Bool, p (r surrIx -> b))]
        ts surrIdx =
          do (ea, ta) <- tas surrIdx
             if ea
               then do (eb,tb) <- tbs surrIdx
                       return (eb, epsilonL ap [| ap |] >>> ta >>> tb)
               else return (False, epsilonL flip [| flip |] >>> ta >>> fb)
        f = fa >>> fb
    in MkTUPIR rla es h ts f
  (ra :: TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t a) ||| rb = MkTUPW $ \g -> 
    let (MkTUPIR rlaa eas ha tas fa) = tUPRuleForGrammar ra g
        (MkTUPIR rlab ebs hb tbs fb) = tUPRuleForGrammar rb g
        rla :: phi ix -> Bool
        rla idx = rlaa idx || rlab idx
        es :: phi ix -> Maybe (p a)
        es idx = liftM2 (|||) (eas idx) (ebs idx)
        h :: phi ix -> p a
        h idx = ha idx ||| hb idx
        ts surrIdx = tas surrIdx ++ tbs surrIdx
    in MkTUPIR rla es h ts $ fa ||| fb
  endOfInput = mkSimpleTUPW endOfInput
  die = MkTUPW $ \_ -> MkTUPIR (const False) (const Nothing) (const die) (const []) die

instance (EpsProductionRule p) =>
         EpsProductionRule (TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t) where
  epsilon = mkEpsTUPW

instance (LiftableProductionRule p) =>
         LiftableProductionRule (TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t) where
  epsilonL = mkEpsLTUPW

instance (PenaltyProductionRule p) =>
         PenaltyProductionRule (TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t) where
  penalty p (r :: TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t a) = MkTUPW $ \g ->
    let (MkTUPIR rla es h ts f) = tUPRuleForGrammar r g
        es' :: phi ix -> Maybe (p a)
        es' idx = liftM (penalty p) (es idx)
        h' :: phi ix -> p a
        h' idx = penalty p (h idx)
    in MkTUPIR rla es' h' ts $ penalty p f

instance (TokenProductionRule p t, ProductionRule p) =>
         TokenProductionRule (TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t) t where
  token tt = mkSimpleTUPW $ token tt
  anyToken = mkSimpleTUPW anyToken 

tlclTailRef :: (LiftableProductionRule p, LoopProductionRule p (UPDomain phi) (UPValue r)) =>
               phi ix -> p ([r ix -> r ix])
tlclTailRef idx = epsilonL (map unUPTV) [|map unUPTV|] >>> manyRef (UPTail idx) 

data WrapTransformUPWrapper p surrIx phi ixT r t ix = WrapTUPW {
  unWrapTUPW :: TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t (r ix)
  }

failHeadRefsTo :: (EqFam phi, ProductionRule p) =>
                  TransformUPGrammar p surrIx phi ixT r t -> phi ix ->
                  TransformUPGrammar p surrIx phi ixT r t 
failHeadRefsTo g idx idx' = 
  let
    nr = MkTUPW $ \g' -> 
      let (MkTUPIR _ _ _ _ rf) = tUPRuleForGrammar (g' idx) g'
      in (MkTUPIR (eqIdx idx) (const Nothing) (\_ -> die) (\_ -> []) rf)
  in unWrapTUPW $ overrideIdx (WrapTUPW . g) idx (WrapTUPW nr) idx'

succeedTailRefs :: forall p surrIx phi ixT r t . (EqFam phi, ProductionRule p, LiftableProductionRule p) =>
                  TransformUPGrammar p surrIx phi ixT r t -> phi surrIx ->
                  TransformUPGrammar p surrIx phi ixT r t 
succeedTailRefs g idx idx' = 
  let
    nr = MkTUPW $ \g' -> 
      let (MkTUPIR rla es rh _ rf) = tUPRuleForGrammar (g' idx) g'
      in (MkTUPIR rla es rh (\_ -> [(False, epsilonL id [|id|])]) rf)
  in unWrapTUPW $ overrideIdx (WrapTUPW . g) idx (WrapTUPW nr) idx'

procTailRefs :: forall a. a -> [a -> a] -> a
procTailRefs = foldl $ flip ($)
-- procTailRefs z [] = z
-- procTailRefs z (x : xs) = procTailRefs (x z) xs

instance (RecProductionRule p (UPDomain phi) (UPValue r),
          LiftableProductionRule p,
          EqFam phi,
          LoopProductionRule p (UPDomain phi) (UPValue r)) => 
         RecProductionRule (TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t) phi r where
  ref (idx :: phi ix) =
    MkTUPW $ \g -> 
    let g' :: TransformUPGrammar p surrIx phi ixT r t
        g' = failHeadRefsTo g idx
        MkTUPIR rla eas ha tas _ = tUPRuleForGrammar (g idx) g'
        h :: forall ix'. phi ix' -> p (r ix)
        h idx' = if rla idx' -- use True to turn off optimization
                 then epsilonL procTailRefs [| procTailRefs |] >>>
                      (hForEmptyHead idx' ||| ha idx') >>>
                      tlclTailRef idx
                 else f
        es :: forall ix'. phi ix' -> Maybe (p (r ix))
        es idx' = if rla idx' -- use True to turn off optimization
                  then eas idx'
                  else Nothing
        hForEmptyHead :: forall ix' . phi ix' -> p (r ix)
        hForEmptyHead idx' = fromMaybe die $ eas idx'
        f = epsilonL unUPBV [|unUPBV|] >>> ref (UPBase idx)
    in MkTUPIR rla es h tas f 

data WrapListOfTailHeadManys p surrIx phi ixT r t ix = WLOTHM {
  unWLOTHM :: [(Bool, p (r ix -> [r surrIx]))]
  }
instance (EqFam phi,
          LiftableProductionRule p,
          LoopProductionRule p (UPDomain phi) (UPValue r)) => 
         LoopProductionRule (TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t) phi r where
  -- TODO: fix!
  manyRef (idx :: phi ix) =
    MkTUPW $ \_ -> 
    let tix :: p (r ix -> [r ix])
        tix = (flip (:) . map unUPBV, [| flip (:) . map unUPBV |]) $|>> manyRef (UPBase idx)
        ts :: phi surrIx -> [(Bool, p (r surrIx -> [r ix]))]
        ts surrIdx = unWLOTHM $ overrideIdx (\_ -> WLOTHM []) idx (WLOTHM [(False, tix)]) surrIdx
        f = (map unUPBV, [|map unUPBV|]) $|>> manyRef (UPBase idx)
    in MkTUPIR (\_ -> False) (const $ Just $ epsilonL [] [| [] |]) (const f) ts f 
  many1Ref (idx :: phi ix) =
    MkTUPW $ \_ -> 
    let tix :: p (r ix -> [r ix])
        tix = (flip (:) . map unUPBV, [| flip (:) . map unUPBV |]) $|>> many1Ref (UPBase idx)
        ts :: phi surrIx -> [(Bool, p (r surrIx -> [r ix]))]
        ts surrIdx = unWLOTHM $ overrideIdx (\_ -> WLOTHM []) idx (WLOTHM [(False, tix)]) surrIdx
        f = (map unUPBV, [|map unUPBV|]) $|>> many1Ref (UPBase idx)
    in MkTUPIR (\_ -> False) (const Nothing) (const f) ts f 

transformUniformPaull' ::
  forall p phi t r ix. 
  (Domain phi, 
   RecProductionRule p (UPDomain phi) (UPValue r),
   LiftableProductionRule p,
   LoopProductionRule p (UPDomain phi) (UPValue r),
   TokenProductionRule p t) =>
  (forall ixT ix' surrIx.
   phi ix' ->
   TransformUPWrapper p surrIx (UPDomain phi) (UPValue r) phi ixT r t (r ix')) ->
  UPDomain phi ix -> p (UPValue r ix)
transformUniformPaull' _ (UPBase idx) = 
  let
    ruleHead = epsilonL unUPHV [|unUPHV|] >>> ref (UPHead idx)
    br = epsilonL procTailRefs [|procTailRefs|] >>> ruleHead >>> tlclTailRef idx
  in epsilonL UPBV [|UPBV|] >>> br
transformUniformPaull' bgram (UPHead (idx :: phi ix'')) = 
  let
    intRule :: TransformUPIntRule p ix'' phi r (r ix'')
    intRule = tUPRuleForGrammar (bgram idx) (bgram `failHeadRefsTo` idx)
    ur :: p (r ix'')
    ur =   tlclwHead intRule idx
       ||| (fromMaybe die $ tlclwEmpty intRule idx)
  in epsilonL UPHV [| UPHV |] >>> ur
transformUniformPaull' bgram (UPTail (idx :: phi ix')) = 
  let
    intRule = tUPRuleForGrammar (bgram idx) (bgram `succeedTailRefs` idx)
    ur = foldr ((|||) . snd) die $ filter (not . fst) $ tlclwTail intRule idx
  in epsilonL UPTV [|UPTV|] >>> ur

-- | Apply a uniform variant of the classic Paull transformation to a given grammar,
-- removing direct and indirect left recursion.
transformUniformPaull ::
  forall phi t r. Domain phi =>
  ProcessingContextFreeGrammar phi t r ->
  ProcessingExtendedContextFreeGrammar (UPDomain phi) t (UPValue r)
transformUniformPaull gram idx = transformUniformPaull' gram idx

transformUniformPaullP ::
  forall phi t r. Domain phi =>
  ProcessingPenaltyContextFreeGrammar phi t r ->
  ProcessingPenaltyExtendedContextFreeGrammar (UPDomain phi) t (UPValue r)
transformUniformPaullP gram idx = transformUniformPaull' gram idx

-- | Apply a uniform variant of the classic Paull transformation to a given extended grammar,
-- removing direct and indirect left recursion.
transformUniformPaullE ::
  forall phi t r. Domain phi =>
  ProcessingExtendedContextFreeGrammar phi t r ->
  ProcessingExtendedContextFreeGrammar (UPDomain phi) t (UPValue r)
transformUniformPaullE gram idx = transformUniformPaull' gram idx

-- | Apply a uniform variant of the classic Paull transformation to a given extended liftable grammar,
-- removing direct and indirect left recursion.
transformUniformPaullLE ::
  forall phi t r. Domain phi =>
  ProcessingLExtendedContextFreeGrammar phi t r ->
  ProcessingLExtendedContextFreeGrammar (UPDomain phi) t (UPValue r)
transformUniformPaullLE gram idx = transformUniformPaull' gram idx