{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE MultiParamTypeClasses #-}
-- |Operation and interface for Genes 
module AI.Heukarya.Gene(
  module Data.Tree
  ,module System.Random
  -- *Gene Operations
  ,genEukarya
  ,crossEukarya
  ,mutateEukarya
  -- *Interfaces or Gene implementation
  ,HeukaryaGene(..)
) where

import Prelude
import Data.Tree 
import Data.List
import System.Random
import Control.Monad
import Control.DeepSeq
import Data.Maybe(isJust,fromJust,catMaybes)
import Data.Text(Text)
import Control.Parallel.Strategies

withStrategyParTraversableRdeepseq :: NFData d => Tree d -> Tree d
withStrategyParTraversableRdeepseq = withStrategy (parTraversable rdeepseq)

-- |Gene's interface
-- use 'Text' to represent types
class (NFData d, Monad m) => HeukaryaGene m d where
  -- |Transform a syntax tree to a expression
  evalTreeGene :: Tree d -> m d
  -- |Get Gene's Type representation
  geneTypeRep  :: d -> m Text
  geneTypeRep = \d -> geneHTypes d >>= return.snd.head
  -- |Get every possible input type and output type
  -- >>> geneHTypes $ toDyn "(+)" ((+)::Int -> Int -> Int)
  -- [([],"Int -> Int -> Int"),(["Int"],"Int -> Int"),(["Int","Int"],"Int")]
  --
  geneHTypes :: d -> m [([Text],Text)]
  -- |get tree's actual output type.
  outputHTreeType :: Tree d -> m Text
  outputHTreeType tree = 
    geneHTypes (rootLabel tree) >>= return.(!!length (subForest tree)) >>= (return.snd)
  -- |the first parameter is just a typeclass indicator that makes which instance to use
  --  get whether the former type is the special case of latter type
  instanceOf :: d -> Text -> Text -> m Bool
  -- |the first parameter is just a typeclass indicator that makes which instance to use
  -- get whether two types are the same
  equalType  :: d -> Text -> Text -> m Bool
  equalType c a b  = (instanceOf c) a b >>= \y->(instanceOf c) b a >>= \x-> return (x && y)
  -- |the first parameter is just a typeclass indicator that makes which instance to use
  -- get whether two types have direct relation.
  directType  :: d -> Text -> Text -> m Bool
  directType c a b = (instanceOf c) a b >>= \y->(instanceOf c) b a >>= \x-> return (x || y)
  -- |the first parameter is just a typeclass indicator that makes which instance to use
  --  return a generalized one if they have direct relation.
  generalizeType :: d -> Text -> Text -> m (Maybe Text)
  generalizeType c a b = do
    x <- (instanceOf c) a b
    y <- (instanceOf c) b a
    return $ if x then Just b else if y then Just a else Nothing

randomList 0 _ = []
randomList n g = g1 : randomList (n-1) g2
  where
  (g1,g2) = split g
randomInfiniteList g = let (g1,g2) = split g in g1 : randomInfiniteList g2
depth = length . levels

-- |generate a tree structured Eukarya. 
genEukarya :: (RandomGen g, HeukaryaGene m d) =>
              g             -- ^RandomGenerator
           -> Int           -- ^depth of tree structure
           -> Text          -- ^String representation of output's Type
           -> [d]           -- ^Ingredient(Gene) of Eukarya for constructing Eukarya
           -> m (Tree d)
genEukarya g n outType geneGlobalList = return.withStrategyParTraversableRdeepseq =<< 
  unfoldTreeM treeUnfolder (g,n,outType)
  where
  treeUnfolder (g,n,outType) = do
    dy  <- dynam
    iTs <- inTypes
    gL3 <- geneList3
    if null gL3 then treeUnfolder (rr!!0,n,outType) else
      return (dy, zip3 (randomInfiniteList (rr!!1)) (cycle [n-1]) iTs)
    where
    rr = randomList 3 g
    i = (fst . next) (rr!!2)
    inTypes = do
      ghtypes <- dynam >>= geneHTypes 
      filterM (\y -> instanceOf dd (snd y) outType) ghtypes >>= return.fst.head
    dynam = geneList3 >>= \z -> return $ z !! (mod i $ length z)
    geneList2 =
      (if n <= 1 then filterM (\x -> geneTypeRep x >>= \y -> instanceOf dd y outType) else return)
      geneGlobalList
    geneList3 = do
      gL2 <- geneList2
      filterM (\y-> geneHTypes y>>=mapM (return.snd)>>=mapM (\z->instanceOf dd z outType)>>=return.or) gL2 
    
    dd = geneGlobalList!!0

extractTree orgTree list = extract orgTree list ([],[],[])
  where
  extract orgTree (k:list) t@(elems,lForests,rForests) =
    if null forest then t else
      extract (forest !! i) list (rootLabel orgTree:elems,take i forest:lForests,drop (i+1) forest:rForests)
    where
    i = mod k (length forest)
    forest = subForest orgTree
  extract orgTree  []  t  = t

archiveTree subTree (elems, lForests, rForests) =
  archive subTree (elems, lForests, rForests)
  where
  archive subTree (elem:elems, lForest:lForests, rForest:rForests) =
    archive (Node elem $ lForest ++ (subTree:rForest)) (elems, lForests, rForests)
  archive subTree   ([],_,_)    = subTree

replaceSubTree tree list subTree =
  archiveTree subTree $ extractTree tree list

getSubTrees tree = map (\x -> (fst x, (reverse . snd) x) ) ( getSubTrees2 (tree,[]) )
getSubTrees2 (tree,trail) = (tree,trail) : concat [ getSubTrees2 (tr, n:trail) | (n, tr) <- zip [0..] (subForest tree) ]

-- |crossover two Eukarya
crossEukarya :: (RandomGen g, HeukaryaGene m d) =>
                g       -- ^RandomGenerator
             -> Int     -- ^depth of tree structure. no functionality here
             -> (Tree d, Tree d) -- ^origin Eukarya pair
             -> m (Tree d, Tree d) -- ^crossovered Eukarya pair
crossEukarya g n o@(tree1, tree2) = do
  psdSubTrs1 <- passedSubTrees1
  psdSubTrs2 <- passedSubTrees2
  ((subtree1,trail1),tp1) <- choiceSubTree1
  ((subtree2,trail2),tp2) <- choiceSubTree2
  zzz <- (instanceOf dd) tp1 tp2
  return $ (\(x,y)-> (withStrategyParTraversableRdeepseq x,withStrategyParTraversableRdeepseq y)) $
    if null psdSubTrs1 then o else if null psdSubTrs2 then o else
      (replaceSubTree tree1 trail1 subtree2, 
       if zzz
         then replaceSubTree tree2 trail2 subtree1
         else tree2)
  where
    tt = mapM (outputHTreeType . fst) . getSubTrees
    comOutTypes = do
      tt1 <- tt tree1
      tt2 <- tt tree2 
      mapM (\(x,y)->generalizeType dd x y) [(a,b) |a<-tt1,b<-tt2] >>= return.catMaybes 
    rr = randomList 2 g
    
    -- tuple subtrees with type
    taggedSubTrees1 = let subtrees = getSubTrees tree1 in 
      mapM (outputHTreeType . fst) subtrees >>= \z -> return $ zip subtrees z
    passedSubTrees1 = do 
      cOTypes   <- comOutTypes
      tgSubTrs1 <- taggedSubTrees1
      filterM (\x -> mapM (equalType dd (snd x)) cOTypes >>= return.or) tgSubTrs1
    choice1 = return.mod ((fst . next) (rr !! 0)) =<< return.length =<< passedSubTrees1
    choiceSubTree1 = do
      psdSubTrs1 <- passedSubTrees1
      c1 <- choice1
      return $ psdSubTrs1 !! c1

    taggedSubTrees2 = let subtrees = getSubTrees tree2 in
      mapM (outputHTreeType . fst) subtrees >>= \z -> return $ zip subtrees z
    passedSubTrees2 = do 
      tgSubTrs2 <- taggedSubTrees2
      ((subtree1,trail1),tp1) <- choiceSubTree1
      filterM (\x -> 
        (return.(&&) (depth (fst$fst x) + depth subtree1 <= n)) =<<
        (instanceOf dd) (snd x) tp1) tgSubTrs2
    choice2 = return.mod ((fst . next) (rr !! 1)) =<< return.length =<< passedSubTrees2
    choiceSubTree2 = do
      psdSubTrs2 <- passedSubTrees2
      c2 <- choice2
      return $ psdSubTrs2 !! c2

    dd = rootLabel tree1

-- |mutate Eukarya
mutateEukarya :: (RandomGen g, HeukaryaGene m d) =>
                 g       -- ^RandomGenerator
              -> Int     -- ^depth of tree structure
              -> [d]     -- ^Ingredient(Gene) of Eukarya for mutating Eukarya
              -> Tree d  -- ^input Eukarya
              -> m (Tree d)  -- ^mutated Eukarya
mutateEukarya g n geneGlobalList tree = do
  tp <- outputHTreeType tree
  eu <- genEukarya g2 (n-depth tree) tp geneGlobalList
  return.withStrategyParTraversableRdeepseq =<< (crossEukarya g1 n (tree, eu) >>= return.fst)
  where
  (g1,g2) = split g