-----------------------------------------------------------------------------
-- |
-- Module      :  Transform.Rules.Lenses.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 lenses involving recursion.
--
-----------------------------------------------------------------------------

module Transform.Rules.Lenses.Rec where

import Data.Type
import Data.Pf
import Data.Spine
import Data.Lens
import Data.Equal
import Transform.Rewriting
import qualified Transform.Rules.PF as PF
import {-# SOURCE #-} Transform.Rules.Lenses
import Transform.Rules.Lenses.Combinators
import Transform.Rules.Lenses.Lists

import Prelude hiding (Functor(..))
import Control.Monad hiding (Functor(..))
import Unsafe.Coerce

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

-- ** In / Out

in_iso_lns = comp_lns in_iso_lns'
in_iso_lns' :: Rule
in_iso_lns' (Lns a b) (COMP_LNS _ INN_LNS OUT_LNS) = do
    Eq <- teq a b
    success "in-Iso-Lns" ID_LNS
in_iso_lns' _ _ = mzero

out_iso_lns = comp_lns out_iso_lns'
out_iso_lns' :: Rule
out_iso_lns' (Lns a b) (COMP_LNS _ OUT_LNS INN_LNS) = do
    Eq <- teq a b
    success "out-Iso-Lns" ID_LNS
out_iso_lns' _ _ = mzero

-- ** Functors

functor_id_lns :: Rule
functor_id_lns (Lns _ _) (FMAP_LNS fctr (Fun a _) ID_LNS) =
    success "functor-Id-Lns" ID_LNS
functor_id_lns _ _ = mzero

functor_comp_lns = comp_lns functor_comp_lns'
functor_comp_lns' :: Rule
functor_comp_lns' (Lns fa fc) (COMP_LNS fb (FMAP_LNS fctr (Fun b c) f) (FMAP_LNS fctr' (Fun a b') g)) = do
    Eq <- feq fctr fctr'
    Eq <- teq b b'
    success "functor-Comp-Lns" $ FMAP_LNS fctr (Fun a c) $ COMP_LNS b f g
functor_comp_lns' _ _ = mzero

functor_def_lns, functor_def_lns' :: Rule
functor_def_lns a x = functor_def_lns' a x >>= success "functor-Def-Lns"
functor_def_lns' (Lns _ _) (FMAP_LNS I _ f) =
    return f
functor_def_lns' (Lns _ _) (FMAP_LNS L _ f) =
    return $ MAP_LNS f
functor_def_lns' (Lns _ _) (FMAP_LNS (K _) _ f) = 
    return $ ID_LNS
functor_def_lns' (Lns _ _) (FMAP_LNS (g:*!:h) t@(Fun c a) f) = do
    l <- functor_def_lns' (Lns (rep g c) (rep g a)) (FMAP_LNS g t f)
    r <- functor_def_lns' (Lns (rep h c) (rep h a)) (FMAP_LNS h t f)
    return $ l `PROD_LNS` r
functor_def_lns' (Lns _ _) (FMAP_LNS (g:+!:h) t@(Fun c a) f) = do
    l <- functor_def_lns' (Lns (rep g c) (rep g a)) (FMAP_LNS g t f)
    r <- functor_def_lns' (Lns (rep h c) (rep h a)) (FMAP_LNS h t f)
    return $ l `SUM_LNS` r
functor_def_lns' (Lns _ _) (FMAP_LNS (g:@!:h) t@(Fun c a) f) = do
    let (hc,ha) = (rep h c,rep h a)
    r <- functor_def_lns' (Lns hc ha) (FMAP_LNS h t f)
    l <- functor_def_lns' (Lns (rep g hc) (rep g ha)) (FMAP_LNS g (Fun hc ha) r)
    return l
functor_def_lns' _ _ = mzero

-- ** Catas

cata_def_lns :: Rule
cata_def_lns (Lns a@(dataFctr -> Just fctr) b) (CATA_LNS g) = do
    guard (not $ isRec fctr)
    Eq <- teq (rep fctr a) (rep fctr b)
    success "cata-Def-Lns" $ COMP_LNS (rep fctr a) g OUT_LNS
cata_def_lns _ _ = mzero

cata_reflex_lns :: Rule
cata_reflex_lns (Lns a b) (CATA_LNS INN_LNS) = do
    Eq <- teq a b
    success "cata-Reflex-Lns" ID_LNS
cata_reflex_lns _ _ = mzero

cata_cancel_lns = comp_lns cata_cancel_lns'
cata_cancel_lns' :: Rule
cata_cancel_lns' (Lns fa b) (COMP_LNS a (ANA_LNS g) INN_LNS) = do
    cata <- ana_shift_lns (Lns a b)  (ANA_LNS g)
    cata_cancel_lns' (Lns fa b) (COMP_LNS a cata INN_LNS)
cata_cancel_lns' (Lns _ b) (COMP_LNS a@(dataFctr -> Just fctr) (CATA_LNS f) INN_LNS) = do
    let fb = rep fctr b
    success "cata-Cancel-Lns" $ COMP_LNS fb f $ FMAP_LNS fctr (Fun a b) (CATA_LNS f)
cata_cancel_lns' _ _ = mzero


cata_fusion_lns = precomp_lns (rightmost_prod_lns ||| rightmost_sum_lns) (cata_fusion_lns' optimise_lns)
--cata_fusion_lns = comp_lns (cata_fusion_lns' optimise_lns)
cata_fusion_lns' :: Rule -> Rule
cata_fusion_lns' r (Lns _ _) (COMP_LNS _ OUT_LNS (CATA_LNS g)) = mzero
cata_fusion_lns' r t@(Lns c@(dataFctr -> Just fctr) a) v@(COMP_LNS b f (CATA_LNS g)) = do
    debug "cata-Fusion-Lns" (Pf t) v
    let (fa,fb) = (rep fctr a,rep fctr b)
        h'      = COMP_LNS b f $ COMP_LNS fb g $ FMAP_LNS fctr (Fun a b) (rconv_lns f)
    h <- r (Lns fa a) h'
    debug "cataRes" (Pf $ Lns fa a) h
    guard $ not $ find (Pf (Lns Any Any)) (rconv_lns TOP) (Pf (Lns fa a)) h
    success "cata-Fusion-Lns" $ CATA_LNS h
cata_fusion_lns' _ _ _ = mzero

cata_shift_lns :: Rule
cata_shift_lns t@(Lns a@(dataFctr -> Just f) b@(dataFctr -> Just g)) v@(CATA_LNS (COMP_LNS gb INN_LNS eta)) = do
    debug "cata-Shift-Lns" (Pf t) v
    Eq <- teq (rep g b) gb
    eta' <- natCoerce_lns f g b eta a
    success "cata-Shift-Lns" $ ANA_LNS $ COMP_LNS (rep f a) eta' OUT_LNS
cata_shift_lns _ _ = mzero


-- ** Anas

ana_def_lns :: Rule
ana_def_lns (Lns a b@(dataFctr -> Just fctr)) (ANA_LNS g) = do
   guard (not $ isRec fctr)
   Eq <- teq (rep fctr b) (rep fctr a)
   success "ana-Def-Lns" $ COMP_LNS (rep fctr b) INN_LNS g
ana_def_lns _ _ = mzero

ana_reflex_lns :: Rule
ana_reflex_lns (Lns a b) (ANA_LNS OUT_LNS) = do
    Eq <- teq a b
    success "ana-Reflex-Lns" ID_LNS
ana_reflex_lns _ _ = mzero

ana_cancel_lns = comp_lns ana_cancel_lns'
ana_cancel_lns' :: Rule
ana_cancel_lns' (Lns b fa) (COMP_LNS a OUT_LNS (CATA_LNS h)) = do
    ana <- cata_shift_lns (Lns b a) (CATA_LNS h)
    ana_cancel_lns' (Lns b fa) (COMP_LNS a OUT_LNS ana)
ana_cancel_lns' (Lns b fa) (COMP_LNS a@(dataFctr -> Just fctr) OUT_LNS (ANA_LNS f)) = do
    let fb = rep fctr b
    success "ana-Cancel-Lns" $ COMP_LNS fb (FMAP_LNS fctr (Fun b a) (ANA_LNS f)) f
ana_cancel_lns' _ _ = mzero

ana_fusion_lns = postcomp_lns (leftmost_prod_lns ||| leftmost_sum_lns) (ana_fusion_lns' optimise_lns)
--ana_fusion_lns = comp_lns (ana_fusion_lns' optimise_lns)
ana_fusion_lns' :: Rule -> Rule
ana_fusion_lns' r (Lns _ _) (COMP_LNS _ (ANA_LNS f) INN_LNS) = mzero
ana_fusion_lns' r t@(Lns a c@(dataFctr -> Just fctr)) v@(COMP_LNS b (ANA_LNS g) f) = do
    debug "ana-Fusion-Lns" (Pf t) v
    let (fa,fb) = (rep fctr a,rep fctr b)
        h'      = COMP_LNS fb (FMAP_LNS fctr (Fun b a) (lconv_lns f)) $ COMP_LNS b g f
    h <- r (Lns a fa) h'
    debug "anaRes" (Pf $ Lns a fa) h
    guard $ not $ find (Pf (Lns Any Any)) (lconv_lns TOP) (Pf (Lns a fa)) h
    success "ana-Fusion-Lns" $ ANA_LNS h
ana_fusion_lns' _ _ _ = mzero

ana_shift_lns :: Rule
ana_shift_lns t@(Lns a@(dataFctr -> Just f) b@(dataFctr -> Just g)) v@(ANA_LNS (COMP_LNS fa eta OUT_LNS)) = do
    debug "ana-Shift-Lns" (Pf t) v
    Eq <- teq (rep f a) fa
    eta' <- natCoerce_lns f g a eta b
    success "ana-Shift-Lns" $ CATA_LNS $ COMP_LNS (rep g b) INN_LNS eta'
ana_shift_lns _ _ = mzero

-- ** Hylos

hylo_shift_lns = comp_lns hylo_shift_lns'
hylo_shift_lns' :: Rule
hylo_shift_lns' q@(Lns a c) v@(COMP_LNS (Data _ fctrf) (CATA_LNS g) (ANA_LNS h)) = do
    debug "hylo-Shift-Lns" (Pf q) v
    COMPF_LNS fctrg c' gold geta <- natSplit_lns c c fctrf g
    Eq <- teq c c'
    let t = Lns (rep fctrf c) (rep fctrg c)
    debug "hyloSplit" (Pf t) geta
    heta <- natCoerce_lns fctrf fctrg c geta a
    success "hylo-Shift-Lns" $ COMP_LNS (fixof fctrg) (CATA_LNS gold) (ANA_LNS $ COMP_LNS (rep fctrf a) heta h)
hylo_shift_lns' _ _ = mzero

hylo_id_lns = comp_lns hylo_id_lns'
hylo_id_lns' :: Rule
hylo_id_lns' t@(Lns c a) v@(COMP_LNS b@(dataFctr -> Just fctr) (CATA_LNS g) (ANA_LNS h)) = do
    Eq <- teq c a
    debug "hylo-Id-Lns" (Pf t) v
    ID_LNS <- optimise_lns (Lns c a) (COMP_LNS (rep fctr c) g h)
    success "hylo-Id-Lns" ID_LNS
hylo_id_lns' _ _ = mzero

-- ** Natural transformations

-- | n . F f = F f . n
natProof_lns :: (Functor f,Functor g) => Fctr f -> Fctr g -> Type a -> Pf (Lens (Rep f a) (Rep g a)) -> Bool
natProof_lns f g a eta = proof optimise_lns t eq1 eq2
    where eq1 = COMP_LNS (rep f a) eta fmapf
          eq2 = COMP_LNS (rep g a) fmapg eta
          fmapf = FMAP_LNS f (Fun a a) BOT
          fmapg = FMAP_LNS g (Fun a a) BOT
          t = Lns (rep f a) (rep g a)
-- ^ We need to prove this property in order to identify natural transformations, since we cannot know such from the types.

-- | Convert a natural transformation applied to some type into a natural transformation over another type
natCoerce_lns :: (MonadPlus m,Functor f,Functor g) => Fctr f -> Fctr g -> Type a
          -> Pf (Lens (Rep f a) (Rep g a)) -> Type b -> m (Pf (Lens (Rep f b) (Rep g b)))
natCoerce_lns f g a eta b = do
    guard (natProof_lns f g a eta)
    return (unsafeCoerce eta)

natSplit_lns :: (Functor f) => Type a -> Type b -> Fctr f -> Pf (Lens (Rep f a) b) -> Rewrite (Pf (Lens (Rep f a) b))
-- Constant
natSplit_lns a b _ ID_LNS = mzero
natSplit_lns a b (K t) f = do
    return $ COMPF_LNS (K b) a ID_LNS f  
-- Sums
natSplit_lns a b fctr@(fctrf :+!: fctrg) v@(EITHER_LNS p f g) = (do
    COMPF_LNS fctrx a' fold feta <- natSplit_lns a b fctrf f
    COMPF_LNS fctry a'' gold geta <- natSplit_lns a b fctrg g
    Eq <- teq a a'
    Eq <- teq a a''
    return $ COMPF_LNS (fctrx :+!: fctry) a (EITHER_LNS p fold gold) (feta -|-<< geta))
    `mplus` (do
    COMPF_LNS fctrx a' fold feta <- natSplit_lns a b fctrf f
    Eq <- teq a a'
    return $ COMPF_LNS (fctrx :+!: fctrg) a (EITHER_LNS p fold g) (feta -|-<< ID_LNS))
    `mplus` (do
    COMPF_LNS fctry a'' gold geta <- natSplit_lns a b fctrg g
    Eq <- teq a a''
    return $ COMPF_LNS (fctrf :+!: fctry) a (EITHER_LNS p f gold) (ID_LNS -|-<< geta))
natSplit_lns a (Either b c) fctr@(fctrf :+!: fctrg) v@(f `SUM_LNS` g) = (do
    COMPF_LNS fctrx a' fold feta <- natSplit_lns a b fctrf f
    COMPF_LNS fctry a'' gold geta <- natSplit_lns a c fctrg g
    Eq <- teq a a'
    Eq <- teq a a''
    return $ COMPF_LNS (fctrx :+!: fctry) a (fold -|-<< gold) (feta -|-<< geta))
    `mplus` (do
    COMPF_LNS fctrx a' fold feta <- natSplit_lns a b fctrf f
    Eq <- teq a a'
    return $ COMPF_LNS (fctrx :+!: fctrg) a (fold -|-<< g) (feta -|-<< ID_LNS))
    `mplus` (do
    COMPF_LNS fctry a'' gold geta <- natSplit_lns a c fctrg g
    Eq <- teq a a''
    return $ COMPF_LNS (fctrf :+!: fctry) a (f -|-<< gold) (ID_LNS -|-<< geta))
-- Products
natSplit_lns a b (fctrf :*!: fctrg) (FST_LNS v) = do
    let old = ID_LNS
        eta = FST_LNS v
    return $ COMPF_LNS fctrf a old eta
natSplit_lns a b (fctrf :*!: fctrg) (SND_LNS v) = do
    let old = ID_LNS
        eta = SND_LNS v
    return $ COMPF_LNS fctrg a old eta
natSplit_lns a (Prod b c) fctr@(fctrf :*!: fctrg) v@(f `PROD_LNS` g) = (do
    COMPF_LNS fctrx a' fold feta <- natSplit_lns a b fctrf f
    COMPF_LNS fctry a'' gold geta <- natSplit_lns a c fctrg g
    Eq <- teq a a'
    Eq <- teq a a''
    return $ COMPF_LNS (fctrx :*!: fctry) a (fold ><<< gold) (feta ><<< geta))
    `mplus` (do
    COMPF_LNS fctrx a' fold feta <- natSplit_lns a b fctrf f
    Eq <- teq a a'
    return $ COMPF_LNS (fctrx :*!: fctrg) a (fold ><<< g) (feta ><<< ID_LNS))
    `mplus` (do
    COMPF_LNS fctry a'' gold geta <- natSplit_lns a c fctrg g
    Eq <- teq a a''
    return $ COMPF_LNS (fctrf :*!: fctry) a (f ><<< gold) (ID_LNS ><<< geta))
-- Composition
natSplit_lns a b fctr e@(COMP_LNS _ _ _) = (do
    COMP_LNS c f g <- rightmost_lns (Lns (rep fctr a) b) e
    COMPF_LNS fctrx a' gold geta <- natSplit_lns a c fctr g
    Eq <- teq a a'
    COMPF_LNS fctry a'' fold feta <- natSplit_lns a b fctrx (COMP_LNS c f gold)
    Eq <- teq a a''
    let old = fold
        eta = COMP_LNS (rep fctrx a) feta geta
    return $ COMPF_LNS fctry a old eta)
    `mplus` (do
    COMP_LNS c f g <- rightmost_lns (Lns (rep fctr a) b) e
    COMPF_LNS fctrx a' gold geta <- natSplit_lns a c fctr g
    Eq <- teq a a'
    let old = COMP_LNS c f gold
        eta = geta
    return $ COMPF_LNS fctrx a old eta)
-- Id and unrecognized cases match here
natSplit_lns a b fctr f = mzero

-- ** Internal converses for fusion rules

rconv_lns = CONV_LNS (Right _L)
lconv_lns = CONV_LNS (Left _L)

-- | f . fº = id
rconv_cancel_lns = comp_lns rconv_cancel_lns'
rconv_cancel_lns' :: Rule
rconv_cancel_lns' t@(Lns a a') (COMP_LNS c (CATA_LNS f) (CONV_LNS (Right _) (ANA_LNS g))) = do
    f' <- ana_shift_lns (Lns c a) (ANA_LNS g)
    rconv_cancel_lns' t (COMP_LNS c (CATA_LNS f) (CONV_LNS (Right _L) f'))
rconv_cancel_lns' t@(Lns a a') (COMP_LNS c (ANA_LNS f) (CONV_LNS (Right _) (CATA_LNS g))) = do
    f' <- cata_shift_lns (Lns c a) (CATA_LNS g)
    rconv_cancel_lns' t (COMP_LNS c (ANA_LNS f) (CONV_LNS (Right _L) f'))
rconv_cancel_lns' t@(Lns a a') v@(COMP_LNS c f (CONV_LNS (Right _) f')) = do
    Eq <- teq a a'
    guard $ geq (Pf (Lns c a)) f f'
    success "rconv-Cancel-Lns" $ ID_LNS
rconv_cancel_lns' _ _ = mzero

-- | fº . f = id
lconv_cancel_lns = comp_lns lconv_cancel_lns'
lconv_cancel_lns' :: Rule
lconv_cancel_lns' t@(Lns a a') (COMP_LNS c (CONV_LNS (Left _) (ANA_LNS g)) (CATA_LNS f)) = do
    f' <- ana_shift_lns (Lns a' c) (ANA_LNS g)
    lconv_cancel_lns' t $ COMP_LNS c (CONV_LNS (Left _L) f') (CATA_LNS f)
lconv_cancel_lns' t@(Lns a a') v@(COMP_LNS c (CONV_LNS (Left _) (CATA_LNS g)) (ANA_LNS f)) = do
    f' <- cata_shift_lns (Lns a' c) (CATA_LNS g)
    lconv_cancel_lns' t $ COMP_LNS c (CONV_LNS (Left _L) f') (ANA_LNS f)
lconv_cancel_lns' (Lns c c') (COMP_LNS a (CONV_LNS (Left _) f') f) = do
    Eq <- teq c c'
    guard $ geq (Pf (Lns c a)) f f'
    success "lconv-Cancel-Lns" $ ID_LNS
lconv_cancel_lns' _ _ = mzero

conv_comp_lns :: Rule
conv_comp_lns (Lns _ _) (CONV_LNS e (COMP_LNS b f g)) =
    success "conv-Comp-Lns" $ COMP_LNS b (CONV_LNS e g) (CONV_LNS e f)
conv_comp_lns _ _ = mzero

conv_conv_lns :: Rule
conv_conv_lns _ (CONV_LNS _ (CONV_LNS _ f)) =
    success "conv-Conv-Lns" f
conv_conv_lns _ _ = mzero

conv_id_lns :: Rule
conv_id_lns (Lns a c) (CONV_LNS _ ID_LNS) = do
    success "conv-Iso-Lns" ID_LNS
conv_id_lns _ _ = mzero

conv_prod_lns :: Rule
conv_prod_lns _ (CONV_LNS e (PROD_LNS f g)) =
    success "conv-Prod-Lns" $ PROD_LNS (CONV_LNS e f) (CONV_LNS e g)
conv_prod_lns _ _ = mzero

conv_sum_lns :: Rule
conv_sum_lns _ (CONV_LNS e (SUM_LNS f g)) =
    success "conv-Sum-Lns" $ SUM_LNS (CONV_LNS e f) (CONV_LNS e g)
conv_sum_lns _ _ = mzero

convs :: Rule
convs = top rconv_cancel_lns ||| top lconv_cancel_lns
    ||| top conv_comp_lns ||| top conv_conv_lns ||| top conv_id_lns
    ||| top conv_prod_lns ||| top conv_sum_lns

recs :: Rule
recs  = top in_iso_lns ||| top out_iso_lns
    ||| top functor_id_lns ||| top functor_comp_lns ||| top functor_def_lns
    ||| top cata_def_lns ||| top cata_reflex_lns ||| top cata_cancel_lns
    ||| top ana_def_lns ||| top ana_reflex_lns ||| top ana_cancel_lns

-- ** Lists

list_ana_cancel_lns, list_ana_cancel_lns' :: Rule
list_ana_cancel_lns = comp_lns list_ana_cancel_lns'
list_ana_cancel_lns' (Lns b fa) (COMP_LNS a@(dataFctr -> Just fctr) OUT_LNS ana) = do
    ANA_LNS g <- list_anas_defs_lns (Lns b a) ana
    success "list-ana-Cancel-Lns" $ COMP_LNS (rep fctr b) (FMAP_LNS fctr (Fun b a) ana) g
list_ana_cancel_lns' _ _ = mzero

list_ana_fusion_lns :: Rule
list_ana_fusion_lns = postcomp_lns (leftmost_prod_lns ||| leftmost_sum_lns) $
    comp1_lns list_anas_defs_lns >>> (ana_fusion_lns' optimise_all_lns)

list_cata_cancel_lns, list_cata_cancel_lns' :: Rule
list_cata_cancel_lns = comp_lns list_cata_cancel_lns'
list_cata_cancel_lns' (Lns fa b) (COMP_LNS a@(dataFctr -> Just fctr) cata INN_LNS) = do
    CATA_LNS f <- list_catas_defs_lns (Lns a b) cata
    success "list-cata-Cancel-Lns" $ COMP_LNS (rep fctr b) f $ FMAP_LNS fctr (Fun a b) cata
list_cata_cancel_lns' _ _ = mzero

list_cata_fusion_lns :: Rule
list_cata_fusion_lns = precomp_lns (rightmost_prod_lns ||| rightmost_sum_lns) $
    comp2_lns list_catas_defs_lns >>> (cata_fusion_lns' optimise_all_lns)

list_hylo_fusion_lns :: Rule
list_hylo_fusion_lns =
     (postcomp_lns (leftmost_prod_lns ||| leftmost_sum_lns) $ precomp_lns list_hylos_defs_lns (ana_fusion_lns' optimise_all_lns))
 ||| (precomp_lns (rightmost_prod_lns ||| rightmost_sum_lns) $ postcomp_lns list_hylos_defs_lns (cata_fusion_lns' optimise_all_lns))