-----------------------------------------------------------------------------
-- |
-- Module      :  Transform.Rules.PF.Rec
-- Copyright   :  (c) 2010 University of Minho
-- License     :  BSD3
--
-- Maintainer  :  hpacheco@di.uminho.pt
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Pointless Rewrite:
-- automatic transformation system for point-free programs
-- 
-- Combinators for the rewriting of point-free functions involving recursion.
--
-----------------------------------------------------------------------------

module Transform.Rules.PF.Rec where
    
import Data.Type
import Data.Pf
import Data.Equal
import Transform.Rewriting
import Transform.Rules.PF.Combinators
import {-# SOURCE #-} Transform.Rules.PF
import Transform.Rules.Lenses.Lists
import Transform.Rules.PF.Sums
import Transform.Rules.PF.Products

import Prelude hiding (Functor(..))
import Control.Monad hiding (Functor(..))
import Control.Monad.RWS hiding (Functor(..),Any)
import Unsafe.Coerce

import Generics.Pointless.Combinators hiding (comp)
import Generics.Pointless.Functors hiding (rep)
import Generics.Pointless.Lenses

-- ** In / Out

in_iso = comp in_iso'
in_iso' :: Rule
in_iso' (Fun a b) (COMP _ INN OUT) = do
    Eq <- teq a b
    success "in-Iso" ID
in_iso' _ _ = mzero

out_iso = comp out_iso'
out_iso' :: Rule
out_iso' (Fun a b) (COMP _ OUT INN) = do
    Eq <- teq a b
    success "out-Iso" ID
out_iso' _ _ = mzero

-- ** Functors

functor_id :: Rule
functor_id (Fun _ _) (FMAP fctr (Fun a _) ID) =
    success "functor-Id" ID
functor_id _ _ = mzero

functor_comp = comp functor_comp'
functor_comp' :: Rule
functor_comp' (Fun fa fc) (COMP fb (FMAP fctr (Fun b c) f) (FMAP fctr' (Fun a b') g)) = do
    Eq <- feq fctr fctr'
    Eq <- teq b b'
    success "functor-Comp" $ FMAP fctr (Fun a c) $ COMP b f g
functor_comp' _ _ = mzero

functor_def, functor_def' :: Rule
functor_def a x = functor_def' a x >>= success "functor-Def"
functor_def' (Fun _ _) (FMAP I _ f) =
    return f
functor_def' (Fun _ _) (FMAP (K _) _ f) = 
    return ID
functor_def' (Fun _ _) (FMAP L _ f) = 
    return $ MAP f
functor_def' (Fun _ _) (FMAP (g:*!:h) t@(Fun c a) f) = do
    l <- functor_def' (Fun (rep g c) (rep g a)) (FMAP g t f)
    r <- functor_def' (Fun (rep h c) (rep h a)) (FMAP h t f)
    return $ l `PROD` r
functor_def' (Fun _ _) (FMAP (g:+!:h) t@(Fun c a) f) = do
    l <- functor_def' (Fun (rep g c) (rep g a)) (FMAP g t f)
    r <- functor_def' (Fun (rep h c) (rep h a)) (FMAP h t f)
    return $ l `SUM` r
functor_def' (Fun _ _) (FMAP (g:@!:h) t@(Fun c a) f) = do
    let hc = rep h c
    let ha = rep h a
    r <- functor_def' (Fun hc ha) (FMAP h t f)
    l <- functor_def' (Fun (rep g hc) (rep g ha)) (FMAP g (Fun hc ha) r)
    return l
functor_def' _ _ = mzero

fzip_def :: Rule
fzip_def (Fun _ _) (FZIP I _ f) =
    success "fzip-Def" ID
fzip_def (Fun _ _) (FZIP (K t) _ f) =
    success "fzip-Def" FST
fzip_def (Fun _ _) (FZIP L (Fun a b) f) =
    success "fzip-Def" $ listzip a b
fzip_def (Fun _ _) (FZIP (fctrf :*!: fctrg) (Fun a c) f) = do
    let (fa,fc) = (rep fctrf a,rep fctrf c)
        (ga,gc) = (rep fctrg a,rep fctrg c)
        t = (Prod (Prod fa fc) (Prod ga gc))
    f' <- fzip_def (Fun (Prod fa fc) (rep fctrf (Prod a c))) (FZIP fctrf (Fun a c) f)
    g' <- fzip_def (Fun (Prod ga gc) (rep fctrg (Prod a c))) (FZIP fctrg (Fun a c) f)
    success "fzip-Def" $ COMP t (f' ><= g') distp_pf
fzip_def (Fun _ _) (FZIP (fctrf :+!: fctrg) (Fun a c) f) = do
    let (fa,fc) = (rep fctrf a,rep fctrf c)
        (ga,gc) = (rep fctrg a,rep fctrg c)
        t = (Either (Either (Prod fa fc) (Prod fa gc)) (Either (Prod ga fc) (Prod ga gc)))
    f' <- fzip_def (Fun (Prod fa fc) (rep fctrf (Prod a c))) (FZIP fctrf (Fun a c) f)
    g' <- fzip_def (Fun (Prod ga gc) (rep fctrg (Prod a c))) (FZIP fctrg (Fun a c) f)
    let l = f' \/= (COMP fa (FMAP fctrf (Fun a (Prod a c)) (ID /\= f)) FST)
        r = (COMP ga (FMAP fctrg (Fun a (Prod a c)) (ID /\= f)) FST) \/= g'
    success "fzip-Def" $ COMP t (l -|-= r) $ dists_pf $ Prod (Either fa ga) (Either fc gc)
fzip_def (Fun _ _) (FZIP (fctrf :@!: fctrg) (Fun a c) f) = do
    let (fa,fc,fac)  = (rep fctrf a,rep fctrf c,rep fctrf (Prod a c))
        (ga,gc,gac) = (rep fctrg a,rep fctrg c,rep fctrg (Prod a c))
        t = (rep fctrf (Prod ga gc))
    f' <- fzip_def (Fun (Prod (rep fctrf ga) (rep fctrf gc)) t) (FZIP fctrf (Fun ga gc) (FMAP fctrg (Fun a c) f))
    g' <- fzip_def (Fun (Prod ga gc) gac) (FZIP fctrg (Fun a c) f)
    success "fzip-Def" $ COMP t (FMAP fctrf (Fun (Prod ga gc) gac) g') f'
fzip_def _ _ = mzero

-- ** Catas

cata_def :: Rule
cata_def (Fun a@(dataFctr -> Just fctr) b) (CATA g) = do
    guard (not $ isRec fctr)
    Eq <- teq (rep fctr a) (rep fctr b)
    success "cata-Def" $ COMP (rep fctr a) g OUT
cata_def _ _ = mzero

cata_reflex :: Rule
cata_reflex (Fun a b) (CATA INN) = do
    Eq <- teq a b
    success "cata-Reflex" ID
cata_reflex _ _ = mzero

cata_cancel = comp cata_cancel'
cata_cancel' :: Rule
cata_cancel' (Fun fa b) (COMP a (ANA g) INN) = do
    cata <- ana_shift (Fun a b)  (ANA g)
    cata_cancel' (Fun fa b) (COMP a cata INN)
cata_cancel' t@(Fun _ b) v@(COMP a@(dataFctr -> Just fctr) (CATA g) INN) = do
    debug "cata-Cancel" (Pf t) v
    let fb = rep fctr b
    success "cata-Cancel" $ COMP fb g $ FMAP fctr (Fun a b) (CATA g)
cata_cancel' _ _ = mzero

cata_fusion = precomp (rightmost_prod ||| rightmost_sum) cata_fusion'
cata_fusion' :: Rule
cata_fusion' (Fun _ _) (COMP _ OUT (CATA g)) = mzero
cata_fusion' t@(Fun (dataFctr -> Just fctr) a) v@(COMP b f (CATA g)) = do
    debug "cataFusion" (Pf t) v
    let (fa,fb) = (rep fctr a,rep fctr b)
        h'      = COMP b f $ COMP fb g $ FMAP fctr (Fun a b) (rconv f)
    h <- optimise_pf (Fun fa a) h'
    debug "cataRes" (Pf $ Fun fa a) h
    guard $ not $ find (Pf (Fun Any Any)) (rconv TOP) (Pf (Fun fa a)) h
    success "cata-Fusion" $ CATA h
cata_fusion' _ _ = mzero

cata_shift :: Rule
cata_shift t@(Fun a@(dataFctr -> Just f) b@(dataFctr -> Just g)) v@(CATA (COMP gb INN eta)) = do
    debug "cata-Shift" (Pf t) v
    Eq <- teq (rep g b) gb
    eta' <- natCoerce f g b eta a
    success "cata-Shift" $ ANA $ COMP (rep f a) eta' OUT
cata_shift _ _ = mzero

-- ** Paras

para_def :: Rule
para_def (Fun a@(dataFctr -> Just fctr) c) (PARA g) = do
   guard (not $ isRec fctr)
   Eq <- teq (rep fctr a) (rep fctr (Prod c a))
   success "para-Def" $ COMP (rep fctr a) g OUT
para_def _ _ = mzero

para_reflex :: Rule
para_reflex (Fun a b) (PARA (COMP _ INN (FMAP _ _ FST))) = do
    Eq <- teq a b
    success "para-Reflex" ID
para_reflex _ _ = mzero

para_cancel = comp para_cancel'
para_cancel' :: Rule
para_cancel' (Fun faa c) (COMP a@(dataFctr -> Just fctr) (PARA g) INN) = do
    Eq <- teq (rep fctr a) faa
    let p = (PARA g `SPLIT` ID)
    success "para-Cancel" $ COMP (rep fctr (Prod c a)) g $ FMAP fctr (Fun a (Prod c a)) p
para_cancel' _ _ = mzero

para_cata = comp para_cata'
para_cata' :: Rule
para_cata' (Fun a@(dataFctr -> Just fctr) b) (PARA f) = do
    let (fb,fba) = (rep fctr b,rep fctr (Prod b a))
        g' = COMP fba f $ FMAP fctr (Fun b (Prod b a)) (rconv FST)
    g <- optimise_pf (Fun fb b) g'
    guard $ not $ find (Pf $ Fun Any Any) (rconv TOP) (Pf $ Fun fb b) g
    success "para-Cata" $ CATA g
para_cata' _ _ = mzero

para_fusion = comp para_fusion'
para_fusion' :: Rule
para_fusion' (Fun _ _) (COMP _ OUT (PARA g)) = mzero
para_fusion' t@(Fun c@(dataFctr -> Just fctr) a) v@(COMP b f (PARA g)) = do
    debug "paraRes!!" (Pf $ Fun c a) v
    let (fbc,fac) = (rep fctr (Prod b c),rep fctr (Prod a c))
        h'          = COMP b f $ COMP fbc g $ FMAP fctr (Fun (Prod a c) (Prod b c)) (rconv f `PROD` ID)
    h <- optimise_pf (Fun fac a) h'
    debug "paraRes" (Pf $ Fun fac a) h
    guard $ not $ find (Pf (Fun Any Any)) (rconv TOP) (Pf (Fun fac a)) h
    success "para-Fusion" $ PARA h
para_fusion' _ _ = mzero

-- ** Anas

ana_def :: Rule
ana_def (Fun a b@(dataFctr -> Just fctr)) (ANA g) = do
   guard (not $ isRec fctr)
   Eq <- teq (rep fctr b) (rep fctr a)
   success "ana-Def" $ COMP (rep fctr b) INN g
ana_def _ _ = mzero

ana_reflex :: Rule
ana_reflex (Fun a b) (ANA OUT) = do
    Eq <- teq a b
    success "ana-Reflex" ID
ana_reflex _ _ = mzero

ana_cancel = comp ana_cancel'
ana_cancel' :: Rule
ana_cancel' (Fun b fa) (COMP a OUT (CATA h)) = do
    ana <- cata_shift (Fun b a) (CATA h)
    ana_cancel' (Fun b fa) (COMP a OUT ana)
ana_cancel' (Fun b fa) (COMP a@(dataFctr -> Just fctr) OUT (ANA h)) = do
    Eq <- teq fa (rep fctr a)
    let fb = rep fctr b
    success "ana-Cancel" $ COMP fb (FMAP fctr (Fun b a) (ANA h)) h
ana_cancel' _ _ = mzero

ana_fusion = postcomp (leftmost_prod ||| leftmost_sum) ana_fusion'
ana_fusion' :: Rule
ana_fusion' (Fun _ _) (COMP _ (ANA f) INN) = mzero
ana_fusion' t@(Fun a c@(dataFctr -> Just fctr)) v@(COMP b (ANA g) f) = do
    debug "ana-Fusion" (Pf t) v
    let (fa,fb) = (rep fctr a,rep fctr b)
        h'      = COMP fb (FMAP fctr (Fun b a) (lconv f)) $ COMP b g f
    h <- optimise_pf (Fun a fa) h'
    debug "anaRes" (Pf $ Fun a fa) h
    guard $ not $ find (Pf (Fun Any Any)) (lconv TOP) (Pf (Fun a fa)) h
    success "ana-Fusion" $ ANA h
ana_fusion' _ _ = mzero

ana_shift :: Rule
ana_shift t@(Fun a@(dataFctr -> Just f) b@(dataFctr -> Just g)) v@(ANA (COMP fa eta OUT)) = do
    debug "ana-Shift" (Pf t) v
    Eq <- teq (rep f a) fa
    eta' <- natCoerce f g a eta b
    success "ana-Shift" $ CATA $ COMP (rep g b) INN eta'
ana_shift _ _ = mzero

-- ** Hylos

hylo_shift = comp hylo_shift'
hylo_shift' :: Rule
hylo_shift' q@(Fun a c) v@(COMP b@(Data _ fctrf) (CATA g) (ANA h)) = do
    debug "hylo-Shift" (Pf q) v
    COMPF fctrg c' gold geta <- natSplit c c fctrf g
    Eq <- teq c c'
    let t = Fun (rep fctrf c) (rep fctrg c)
    debug "hyloSplit" (Pf t) geta
    heta <- natCoerce fctrf fctrg c geta a
    success "hylo-Shift" $ COMP (fixof fctrg) (CATA gold) (ANA $ COMP (rep fctrf a) heta h)
hylo_shift' _ _ = mzero

hylo_id = comp hylo_id'
hylo_id' :: Rule
hylo_id' t@(Fun c a) v@(COMP b@(dataFctr -> Just fctr) (CATA g) (ANA h)) = do
    Eq <- teq c a
    debug "hylo-Id" (Pf t) v
    ID <- optimise_pf (Fun c a) (COMP (rep fctr c) g h)
    success "hylo-Id" ID
hylo_id' _ _ = mzero

-- ** Natural transformations

natProof :: (Functor f,Functor g) => Fctr f -> Fctr g -> Type a -> Pf (Rep f a -> Rep g a) -> Bool
natProof f g a eta = proof optimise_pf t eq1 eq2
    where eq1 = COMP (rep f a) eta fmapf
          eq2 = COMP (rep g a) fmapg eta
          fmapf = FMAP f (Fun a a) BOT
          fmapg = FMAP g (Fun a a) BOT
          t = Fun (rep f a) (rep g a)

natCoerce :: (MonadPlus m,Functor f,Functor g) => Fctr f -> Fctr g -> Type a
          -> Pf (Rep f a -> Rep g a) -> Type b -> m (Pf (Rep f b -> Rep g b))
natCoerce f g a eta b = if (natProof f g a eta) then return (unsafeCoerce eta) else mzero

-- Separates the natural part from the type-dependent one in a functor transformation
natSplit :: (Functor f) => Type a -> Type b -> Fctr f -> Pf ((Rep f a) -> b) -> Rewrite (Pf ((Rep f a) -> b))
-- Constant
natSplit a b _ ID = mzero
natSplit a b (K t) f = do
    return $ COMPF (K b) a ID f  
-- Sums
natSplit a b fctr@(fctrf :+!: fctrg) v@(EITHER f g) = (do
    COMPF fctrx a' fold feta <- natSplit a b fctrf f
    COMPF fctry a'' gold geta <- natSplit a b fctrg g
    Eq <- teq a a'
    Eq <- teq a a''
    return $ COMPF (fctrx :+!: fctry) a (EITHER fold gold) (feta -|-= geta))
    `mplus` (do
    COMPF fctrx a' fold feta <- natSplit a b fctrf f
    Eq <- teq a a'
    return $ COMPF (fctrx :+!: fctrg) a (EITHER fold g) (feta -|-= ID))
    `mplus` (do
    COMPF fctry a'' gold geta <- natSplit a b fctrg g
    Eq <- teq a a''
    return $ COMPF (fctrf :+!: fctry) a (EITHER f gold) (ID -|-= geta))
natSplit a (Either b c) fctr@(fctrf :+!: fctrg) v@(f `SUM` g) = (do
    COMPF fctrx a' fold feta <- natSplit a b fctrf f
    COMPF fctry a'' gold geta <- natSplit a c fctrg g
    Eq <- teq a a'
    Eq <- teq a a''
    return $ COMPF (fctrx :+!: fctry) a (fold -|-= gold) (feta -|-= geta))
    `mplus` (do
    COMPF fctrx a' fold feta <- natSplit a b fctrf f
    Eq <- teq a a'
    return $ COMPF (fctrx :+!: fctrg) a (fold -|-= g) (feta -|-= ID))
    `mplus` (do
    COMPF fctry a'' gold geta <- natSplit a c fctrg g
    Eq <- teq a a''
    return $ COMPF (fctrf :+!: fctry) a (f -|-= gold) (ID -|-= geta))
-- Products
natSplit a b (fctrf :*!: fctrg) FST = do
    return $ COMPF fctrf a ID FST
natSplit a b (fctrf :*!: fctrg) SND = do
    return $ COMPF fctrg a ID SND
natSplit a b fctr v@(_ `SPLIT` _) = do
    v' <- (prod_undef ||| (prod_unfusion >>> comp1 (try prod_undef))) (Fun (rep fctr a) b) v
    natSplit a b fctr v'
natSplit a (Prod b c) fctr@(fctrf :*!: fctrg) v@(f `PROD` g) = (do
    COMPF fctrx a' fold feta <- natSplit a b fctrf f
    COMPF fctry a'' gold geta <- natSplit a c fctrg g
    Eq <- teq a a'
    Eq <- teq a a''
    return $ COMPF (fctrx :*!: fctry) a (fold ><= gold) (feta ><= geta))
    `mplus` (do
    COMPF fctrx a' fold feta <- natSplit a b fctrf f
    Eq <- teq a a'
    return $ COMPF (fctrx :*!: fctrg) a (fold ><= g) (feta ><= ID))
    `mplus` (do
    COMPF fctry a'' gold geta <- natSplit a c fctrg g
    Eq <- teq a a''
    return $ COMPF (fctrf :*!: fctry) a (f ><= gold) (ID ><= geta))
-- Composition
natSplit a b fctr e@(COMP _ _ _) = (do
    COMP c f g <- rightmost (Fun (rep fctr a) b) e
    COMPF fctrx a' gold geta <- natSplit a c fctr g
    Eq <- teq a a'
    COMPF fctry a'' fold feta <- natSplit a b fctrx (COMP c f gold)
    Eq <- teq a a''
    let old = fold
        eta = COMP (rep fctrx a) feta geta
    return $ COMPF fctry a old eta)
    `mplus` (do
    COMP c f g <- rightmost (Fun (rep fctr a) b) e
    COMPF fctrx a' gold geta <- natSplit a c fctr g
    Eq <- teq a a'
    let old = COMP c f gold
        eta = geta
    return $ COMPF fctrx a old eta)
-- Id and unrecognized cases match here
natSplit a b fctr f = mzero


-- ** Internal converses for fusion rules

rconv = CONV (Right _L)
lconv = CONV (Left _L)

rconv_cancel = comp rconv_cancel'
rconv_cancel' :: Rule
rconv_cancel' t@(Fun a a') (COMP c (CATA f) (CONV (Right _) (ANA g))) = do
    f' <- ana_shift (Fun c a) (ANA g)
    rconv_cancel' t (COMP c (CATA f) (CONV (Right _L) f'))
rconv_cancel' t@(Fun a a') (COMP c (ANA f) (CONV (Right _) (CATA g))) = do
    f' <- cata_shift (Fun c a) (CATA g)
    rconv_cancel' t (COMP c (ANA f) (CONV (Right _L) f'))
rconv_cancel' (Fun a a') (COMP c f (CONV (Right _) f')) = do
    Eq <- teq a a'
    guard $ geq (Pf (Fun c a)) f f'
    success "rconv-Cancel" ID
rconv_cancel' _ _ = mzero

lconv_cancel = comp lconv_cancel'
lconv_cancel' :: Rule
lconv_cancel' t@(Fun a a') (COMP c (CONV (Left _) (ANA g)) (CATA f)) = do
    f' <- ana_shift (Fun a' c) (ANA g)
    lconv_cancel' t $ COMP c (CONV (Left _L) f') (CATA f)
lconv_cancel' t@(Fun a a') v@(COMP c (CONV (Left _) (CATA g)) (ANA f)) = do
    f' <- cata_shift (Fun a' c) (CATA g)
    lconv_cancel' t $ COMP c (CONV (Left _L) f') (ANA f)
lconv_cancel' (Fun c c') (COMP a (CONV (Left _) f') f) = do
    Eq <- teq c c'
    guard $ geq (Pf (Fun c a)) f f'
    success "rconv-Cancel" ID
lconv_cancel' _ _ = mzero

conv_comp :: Rule
conv_comp (Fun _ _) (CONV e (COMP b f g)) =
    success "conv-Comp" $ COMP b (CONV e g) (CONV e f)
conv_comp _ _ = mzero

conv_conv :: Rule
conv_conv _ (CONV _ (CONV _ f)) =
    success "conv-Conv" f
conv_conv _ _ = mzero

conv_id :: Rule
conv_id _ (CONV _ ID) =
    success "conv-Id" ID
conv_id _ _ = mzero

conv_prod :: Rule
conv_prod _ (CONV e (PROD f g)) =
    success "conv-Prod" $ PROD (CONV e f) (CONV e g)
conv_prod _ _ = mzero


conv_sum :: Rule
conv_sum _ (CONV l (SUM f g)) =
    success "conv-Sum" $ SUM (CONV l f) (CONV l g)
conv_sum _ _ = mzero

convs :: Rule
convs = top rconv_cancel ||| top lconv_cancel
    ||| top conv_comp ||| top conv_conv ||| top conv_id
    ||| top conv_prod ||| top conv_sum

recs :: Rule
recs = top in_iso ||| top out_iso
   ||| top functor_id ||| top functor_comp ||| top functor_def ||| top fzip_def
   ||| top cata_def ||| top cata_reflex
   ||| top para_def ||| top para_reflex ||| top para_cancel
   ||| top ana_def ||| top ana_reflex