{-# OPTIONS_GHC -XTemplateHaskell #-}
module Igor2.RuleDevelopment.Cata where

import Igor2.Data.IgorMonad
import Igor2.Data.CallDependencies
import Igor2.Data.Rules
import Igor2.RuleDevelopment.ListCata

import Syntax
import Igor2.Ppr
import Igor2.Logging

import Control.Monad.Error
import Data.List ((\\), groupBy)
import Data.Function (on)
import Data.Maybe
import Data.Either (rights)
import Data.Util (insertAt)
import Generics.Pointless.Functors (Mu)
import Generics.Pointless.Combinators ((\/), _L)
import Generics.Pointless.RecursionPatterns (cata, para)

tyMorphIntro  :: CovrRule -> IM [(CovrRules,[Call])]
tyMorphIntro cr = do 
  p <- usePara 
  (if p then paraIntro else cataIntro) cr

cataIntro :: CovrRule -> IM [(CovrRules,[Call])]
cataIntro cr = do  
  waypointS $ text "Catamorphism Introduction"
  logNO $ text "Check if 'cata' applies to: " <^> pretty cr

  tyMorphIdcs cr >>= gmsg >>= oneOrNone (genericMorph cr)
  `catchError`
  (\e1 -> logNO (text e1) >> (smsg (lstCataArgIs cr) >>=
                                   oneOrNone (listCata cr)))
  `catchError`
  \e2 -> logNO (text e2) >> return []

    where
    gmsg i = logNO ( linebreak <>
                      text "Try generic catamorphism for argument indices" <+>
                      pretty i) >> return i
    smsg i = logNO ( linebreak <> text "Try special catamorphism on lists for argument indices" <+>
                      pretty i) >> return i

paraIntro :: CovrRule -> IM [(CovrRules,[Call])]
paraIntro cr = do

    waypointS $ text "Paramorphism Introduction"
    logNO $ text "Check if 'para' applies to: " <^> pretty cr

    (liftM (lstCataArgIs cr ++) (tyMorphIdcs cr) >>=
                 gmsg >>= oneOrNone (genericMorph cr))
    `catchError`
    \e2 -> logNO (text e2) >> return []

    where
    gmsg i = logNO ( linebreak <>
                      text "Try generic paraemorphism for argument indices" <+>
                      pretty i) >> return i              

oneOrNone :: (a -> IM b) -> [a] -> IM b
oneOrNone _ []     = throwError . strMsg $ "No more applicable!"
oneOrNone f (x:xs) = f x `catchError` \e -> logNO (text e) >> oneOrNone f xs

-- get the indices of all arguments In reverse order which are applicable for
-- catamorhpism or paramorphisms
tyMorphIdcs :: CovrRule -> IM [Int]
tyMorphIdcs cr = 
    mapM doit (zip (lhs . crul $ cr) [0..]) >>= return . reverse . rights
    where
    doit :: (TExp,Int) -> IM (Either String Int)
    doit x = applyC $ (match (fst x) muVar) >> return (snd x)
-- not used yet
tyFuncIdcs cr = 
    mapM doit (zip (lhs . crul $ cr) [0..]) >>= return . reverse . rights
    where
    doit :: (TExp,Int) -> IM (Either String Int)
    doit x = applyC $ (match (fst x) funcVar) >> return (snd x)

lstCataArgIs :: CovrRule -> [Int]
lstCataArgIs = 
    (map fst) . (filter (hasListT . snd)) . (zip [0..]) . lhs . crul

-- fails if nothing can be found
genericMorph :: CovrRule -> Int -> IM [(CovrRules,[Call])]
genericMorph cr i = do 
  withP <- usePara
  logNO $ msg withP
  evi   <- breakupM cr
  (dataTypeName i (crul cr) >>= ctorsOf >>= partByCtors i evi >>= mkMorphism withP i cr)  `catchError` (\e -> noMorph withP e)
    where
    -- data type Name of the  argument of the covering rule
    dataTypeName = ((dataName . typeOf) .) . ithArg
    msg p = linebreak <> text "Check argument" <+> 
            int i <+> text "for general morphism:" <+>
            text (if p then "'para'" else "cata'")
    noMorph p e = throwError . strMsg $ (if p then "Para" else "Cata") ++
                           "morphism not applicable for argument " ++
                           show i ++ "! Universal Property not satisfied: " ++
                           show e

partByCtors :: (Monad m) => Int -> [CovrRule] -> [Name] -> m [[CovrRule]]
partByCtors i crs ns = mapM (byCtor crs) ns
    where
    byCtor crs n = let rs = filter ((hasCtor n) . (ithArg i) . crul) crs in
                   if null rs then noPartition else return rs
    --hasCtor n (TVarE _ _)        = True
    hasCtor n (TConE n' _)       = n == n'
    hasCtor n (TAppE e _)        = hasCtor n e
    hasCtor n  c                 = False
    noPartition = fail $ "No Partition possible for argument " ++ show i ++ "! "

-- make the catamorphisms with all necessary mediating functions
mkCata, mkPara :: Int -> CovrRule -> [[CovrRule]] -> IM [(CovrRules,[Call])]
mkCata = mkMorphism False 
mkPara = mkMorphism True     

mkMorphism :: Bool -> Int -> CovrRule -> [[CovrRule]] -> IM [(CovrRules,[Call])]
mkMorphism para i cr parts = do
     ios  <- mapM (abduceIOs i) parts
     mfns <- mapM checkIOs ios >> mapM (addIO . rules) ios
     mfcr <- mapM coverAll mfns
     let mfuns = map (mkMediators (butIthArg i $ crul cr)) mfcr     
     let fExp = foldTAppE  morphEx [botEx, (mkInfxs mfuns), ithArg i $ crul cr]
     let cr' = modifycrul cr (\r -> rule (lhs r) fExp)
     let clls = map  (\mfn -> (name cr',mfn, EQ)) mfns
     logNO $ text "Generic" <+> text morphStr <+>
              text "applicable on argument" <+> int i
     logDE $ text "With partitions" <+> pretty parts
     return [(covrRules (cr':mfcr),clls)]
     where
     morphEx = if para then paraEx else cataEx
     morphStr = if para then "'para'" else "'cata'"
     mkMediators as r = 
        let fty = (map typeOf) . lhs . crul $ r
            fex = tConE (name r) fty
        in foldTAppE fex as
     mkInfxs [l,r] = 
        let lty = typeOf l; rty = typeOf r
            ety = foldAppT eitherCon $ map (head . unArrowT) [lty,rty]
        in  tInfixE l (tConE '(\/) [lty, rty, ety, last . unArrowT $ lty]) r
     mkInfxs (l:rs) =
        let r = mkInfxs rs; lty = typeOf l; rty = typeOf r
            ety = foldAppT eitherCon $ map (head . unArrowT) [lty, rty]
        in  tInfixE l (tConE '(\/) [lty, rty, ety, last . unArrowT $ rty]) r


checkIOs :: [Rule] -> IM ()
checkIOs = mbfail . (all allEqual). 
           (groupBy ((((all $ uncurry equal) . ) . zip) `on` lhs))
    where
    mbfail b = if b then return () 
               else fail "Contradiction in abduced IOs: Not a function!"
    allEqual l = all (uncurry (equal `on` rhs)) $ zip l (tail l)
    
-- abduce IO examples for one mediating function, add them to the evidence and 
-- return the name of the mediating function. The first argument 'i' is the 
-- index of the catamorphism-enabled argument, the second are all examples of 
-- the initial rule with the same constructor at position 'i'
abduceIOs :: Int -> [CovrRule] -> IM [Rule]
abduceIOs i crs = mapM (abduceIO i) crs 

-- abduce a single example given the index of the morphism-argument and the 
-- example of the initial rule.
abduceIO :: Int -> CovrRule -> IM Rule 
abduceIO i cr = 
    case subterms carg of
      [] -> return . iorule $ wild
           -- the morphism-arg is a constant constructor with no arguments
           -- so it is not required for the function and replaced by a 
           -- wildcard
      as -> usePara >>= \p -> mapM (mkMArg p) as >>= 
           return . iorule . nestedTupE 
          -- the morphism-arg is a constructor with at least one argument
    where
    args = butIthArg i $ crul cr
    carg = ithArg i $ crul cr 
    wild = tVarE "_x" (typeOf carg)
    iorule ca = rule (args ++ [ca])(rhs . crul $ cr)
    mkMArg para a 
        | (typeOf a)==(typeOf carg) = 
            evalIO (name cr) (insertAt i a args) >>= 
            -- 'simulate recursive call, by evaluating against IOs
                   maybe (fail "Insufficient support in Examples") (mu para a)
        | otherwise                 = return a
    mu para a ar 
       | para      = return $ tTupE  [ar,a]
       -- here happens the paramorphism! Pair of input 'a' and 
       -- result of recursive call 'ar'
       | otherwise = return ar
       -- when cata then we just return the result of the 
       --recurisve call

-- This is not DRY!!! copied from Cata, only small changes, 
    

                        
            
             

-- end not DRY                                      
                                      

nestedTupE [a]    = a -- tTupE [a]
nestedTupE [a,b]  = tTupE [a,b]
nestedTupE (x:xs) = tTupE [x,nestedTupE xs]

muVar = tVarE "a" $ forallT [(''Mu, "a")] (varT "a")
funcVar = tVarE "a" $ forallT [(''Functor, "a")] (varT "a")

botEx = TConE ('_L) (varT "a")
               
cataEx = TConE ('cata) $ 
         ForallT [ mkPred ''Mu "a"
                 , Pred fctrname (AppT pctr (varT "a"))]
                 (arrowT [varT "a", arrowT [AppT (AppT fctr (varT "a")) (varT "b"), varT "b"], varT "a", varT "b"])

paraEx =  TConE ('para) $
         ForallT [ mkPred ''Mu "a"
                 , Pred fctrname (AppT pctr (varT "a"))]
                 (arrowT [varT "a", arrowT [AppT (AppT fctr (tupT [varT "b", varT "a"])) (varT "b"), varT "b"], varT "a", varT "b"])

fctrname = mkName "Generics.Pointless.Functors.Functor"
fctr = conT fctrname
pctr = conT (mkName "Generics.Pointless.Functors.PF")
