{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE UndecidableInstances   #-}

module Generics.MultiRec.Rewriting.AC where

import Generics.MultiRec.Base
import Generics.MultiRec.Constructor
import Generics.MultiRec.Eq

-- * Associative/commutative terms generator class
class (HEq s f) => ACterms s f where
  hACterms :: s ix -> [String] -> Bool -> 
               (forall ix'. (Fam s) => s ix' -> I0 ix' -> [I0 ix'])
                  -> f I0 ix -> [f I0 ix]

instance (Eq x) => ACterms s (K x) where
  hACterms _ _ _ _ f = [f]

instance ACterms s U where
  hACterms _ _ _ _ U = [U]

instance (Constructor c, ACterms s f) => ACterms s (C c f) where
  hACterms s ac _ r c@(C f) = map C $ hACterms s ac (elem (conName c) ac) r f

instance (Fam s, El s xi) => ACterms s (I xi) where
  hACterms _ _ _ r (I ix) = map I $ r proof ix

instance (ACterms s f, ACterms s g) => ACterms s (f :+: g) where
  hACterms s ac b r (L x) = map L $ hACterms s ac b r x
  hACterms s ac b r (R x) = map R $ hACterms s ac b r x

instance (ACterms s f, ACterms s g) => ACterms s (f :*: g) where
  hACterms s ac b r (f :*: g) = [ x :*: y
                                | x <- hACterms s ac b r f
                                , y <- hACterms s ac b r g]

-- TODO: fix the undefined below
instance (HEq s (PF s), Fam s, El s ix, El s xi, EqS s) 
    => ACterms s ((I ix) :*: (I xi)) where
  hACterms s ac b r (i1@(I (I0 x1)) :*: i2@(I (I0 x2))) = 
    [ x :*: y 
    | x <- hACterms s ac b r i1
    , y <- hACterms s ac b r i2 ] ++
      case (eqS (proof :: s ix) (proof :: s xi), b) of
        (Just Refl, True) -> if (eq (proof :: s ix) x1 x2) 
                             then []
                             else [ y :*: x
                                  | x <- hACterms s ac b r i1
                                  , y <- hACterms s ac b r i2 ]
        _                 -> []

instance (ACterms s f) => ACterms s (f :>: ix) where
  hACterms s ac b r (Tag x) = map Tag $ hACterms s ac b r x

-- * Associative/commutative terms generator main function
acTerms :: forall s ix. (Fam s, ACterms s (PF s), EqS s)
        => s ix -> [String] -> ix -> [ix]
acTerms s ac = map (to s) 
             . hACterms s ac False (\s' (I0 x) -> map I0 (acTerms s' ac x))
             . from s