module Test.Agata.Base 
    ( agata, agataWith, agataSC, agataEnum
    , Buildable(..), Builder
    , rebuild, rb, (>=>), (*>), ($>), (.>), graft, inline, automutrec
    , use, construct, autorec, nonrec, mutrec, rec
    ) where


import Test.QuickCheck
import Control.Monad.State.Lazy
import Control.Monad (liftM2)
import Control.Applicative((<$>))

import Data.Maybe(mapMaybe)
import Data.Tagged

import Test.Agata.Common
import Test.Agata.Strategies


agata :: Buildable a => Gen a
agata = agataWith linearSize

agataWith :: Buildable a => Strategy a -> Gen a
agataWith s = do
  dist <- sized $ flip s dimension
  evalImproving (dimension+1,0,[]) $ ii dist undefined
  where
    ii :: Buildable a => Improving () -> a -> Improving a
    ii dist a = currentDimension >>= \lvl -> case unTagged lvl of
      0 -> put (0,0,[]) >> realImp a
      _ -> do
        x <- realImp a 
        dec
        dist
        ii dist x

    dec = get >>= \(lvl,r,[]) -> put (lvl-1,r,[])


evalImproving :: (Dimension a,Int,[Int]) -> Improving a -> Gen a
evalImproving (d,k,ss) = flip evalStateT (unTagged d,k,ss)

agataSC :: Buildable a => Int -> [a]
agataSC = snd . agataEnum

agataEnum :: Buildable a => Int -> (Integer,[a])
agataEnum 0 = (toInteger $ length xs, xs) where 
  xs = concat $ snd $ unzip [benum c 0|c<-build]
agataEnum n
 | n < 0 = (0,[])
 | otherwise = (sum ms, concat xs) where 
  (ms,xs) = unzip [benum c n|c<-build]





class Buildable a where
  build :: [Builder a]
  improve :: a -> Improving a
  improve = return
  dimension :: Dimension a
  dimension = autoDim

data DB a = BuildDebug (Dimension a) [Builder a] deriving Show

db :: Buildable a => DB a 
db = BuildDebug dimension build

rbuild :: Buildable a => Tagged a [Builder a]
rbuild = return build

data Builder a = MkBuilder {
  bskel :: Int -> Improving a,
  benter :: a, 
  benum :: Int -> (Integer,[a]),
  bfields :: [Recursivity a],
  bweight :: Int
  }

instance Show (Builder a) where
  show = show . bfields


brec :: Builder a -> Dimension a -> Bool
brec b d = d > 0 && (not . null $ filter (rc d) (bfields b))



realBuild :: (Buildable a) => Int -> Improving a
realBuild n = do
  c <- currentDimension
  let recs   = [bskel b n|b<- build, brec b c]
  let nrecs  = [bskel b n|b<- build, not $ brec b c]
  let exits   = [bskel b n|b<- build, brec b c, Rec `notElem` bfields b]
  join (lift $ elements $ if n > 0 then 
    if null recs then [get >>= error . show] else recs
    else if null nrecs then recs else nrecs)
      -- FIXME : Get exits
      -- _ -> nrecs
  


-- Determines if a value is defined, should be defined, or left undefined
realImp :: Buildable a => a -> Improving a
realImp a = do
  cur <- currentDimension
  case compare (dimension `taggedWith` a) cur of
     GT -> improve a
     EQ -> if cur == 0 then realBuild 0 else unTagged (bacq a)
     LT -> if (dimension `taggedWith` a) == cur - 1 then unTagged breq else return a

breq :: Buildable a => Tagged a (Improving a)
breq = isAlwaysRecursive >>= \b -> return $ if b then request >> return (error "1") else lift (elements (map benter build)) >>= improve

bacq :: Buildable a => a -> Tagged a (Improving a)
bacq a = isAlwaysRecursive >>= \b -> return $ if b then acquire >>= realBuild else improve a

rebuild :: a -> (a -> Improving b) -> Improving b
rebuild a f = f a 

rb :: Buildable a => a -> (a->b) -> Improving b
rb a f = f <$> realImp a


data Recursivity a = 
    NonRec (Dimension a) 
  | Rec 
  | MutRec 
  | AutoMutRec (Dimension a) 
  | AutoRec (Dimension a) 
    deriving (Show,Eq)

erc r = case r of
       MutRec    -> True
       Rec       -> True
       AutoMutRec _ -> True
       AutoRec _ -> False
       NonRec _  -> False

rc d r = case r of
       MutRec    -> True
       Rec       -> True
       AutoMutRec n -> n >= fromIntegral d
       AutoRec n -> n >= fromIntegral d
       NonRec _  -> False

isAlwaysRecursive :: Buildable a => Tagged a Bool
isAlwaysRecursive =
  any erc . concatMap bfields <$> rbuild


-- A type that represents four possible computations on constructors
--  Build a value with a list of sizes for recursive fields
--  Collect informations about the fields of the constructor
--  Enumerate all values to a specific depth
--  Return a value where all fields are undefined
data Application b a = 
    Build (Improving (a,[Int]))
  | Fields [Recursivity b]
  | Enumerate Int Integer [a]
  | Enter a

a $> b = [construct a b]
infixr 8 $>


inline :: Buildable a => (a -> b) -> [Builder b]
inline f = map trans build where
  trans b = MkBuilder
    (\n -> f <$> bskel b n)
    (f $ benter b)
    (\n -> if n <= 0 then (0,[]) else 
             let (m1,ys) = agataEnum (n-1) in 
               if m1 <= 0 then (0,[]) else (m1,[f a|a <- ys]))
    (map refield $ bfields b)
    (bweight b)
    where 
      refield r = case r of
        MutRec    -> MutRec
        Rec       -> Rec
        AutoMutRec n -> AutoMutRec (retag n)
        AutoRec n -> AutoRec (retag n)
        NonRec n  -> NonRec (retag n)

construct :: a -> (Application b a -> Application b b) -> Builder b
construct c f = MkBuilder skel enter enm fields 1 where
  fields = case f $ Fields [] of
    Fields ls -> ls
  recfields lev = filter (rc lev) fields
  isrec lev = not $ null $ recfields lev
  skel n        = do
    rs <- length . recfields <$> currentDimension
    ns <- if rs == 0 then return $ repeat 0 else lift $ piles rs (n-1)
    let Build m = f (Build $ return (c,ns))
    fst <$> m
  enm n = case f $ Enumerate n 1 [c] of
    Enumerate _ m ls -> (m,ls)
  enter = case f $ Enter c of
    Enter x -> x

graft :: Gen a -> (Int -> (Integer,[a])) -> [Builder a]
graft g e = [MkBuilder (lift . flip resize g) undefined e [MutRec] 1]

use :: a -> [Builder a]
use x = [construct x id]


(.>) a b = b . a
(*>) a b = a >=> b

autoDim :: Buildable a => Dimension a
autoDim = do 
  r <- isAlwaysRecursive
  if r then (+1) <$> maxdim else maxdim where
    maxdim :: Buildable a => Dimension a
    maxdim = (maximum . (0:)) <$> (rbuild >>= sequence . mapMaybe dimOf . concatMap bfields)  where
      dimOf r = case r of
        NonRec d  -> Just d
        AutoRec d -> Just d
        _         -> Nothing


def :: Buildable a => Application c (a -> b) -> Application c b
def (Enter f)          = Enter $ f (error "Entry-value")
def (Enumerate n 0 []) = Enumerate n 0 []
def (Enumerate n m xs) = Enumerate n (m1*m) [f a|f <- xs, a <- ys] where
  (m1,ys) = agataEnum (n-1)


mutrec :: Buildable a => Application c (a -> b) -> Application c b
mutrec x = case x of
  Fields xs -> Fields $ MutRec : xs
  Build mf  -> Build $ do
    (f,x:xs) <- mf
    realBuild x >>= \e -> return (f e,xs)
  _         -> def x

rec :: Buildable c => Application c (c -> b) -> Application c b
rec x = case x of
  Fields xs -> Fields $ Rec : xs
  _ -> mutrec x

nonrec :: Buildable a => Application c (a -> b) -> Application c b
nonrec x = case x of
  Fields xs -> Fields $ NonRec (retag $ appDimension x) : xs
  Build mf  -> Build $ do
    (f,ns) <- mf
    realImp undefined >>= \e -> return (f e,ns)
  _         -> def x

autorec :: Buildable a => Application c (a -> b) -> Application c b
autorec x = case x of
  Fields xs -> Fields $ AutoRec (retag $ appDimension x) : xs
  Build mf  -> Build $ do
    c <- currentDimension
    let isRec = appDimension x >= c
    if isRec then unbuild $ mutrec x else unbuild $ nonrec x
    where
      unbuild (Build x) = x
  _         -> def x

automutrec :: Buildable a => Application c (a -> b) -> Application c b
automutrec x = case x of
  Fields xs -> Fields $ AutoMutRec (retag $ appDimension x) : xs
  _         -> autorec x

appDimension :: Buildable a => Application c (a->b) -> Dimension a
appDimension f = dimension