{-# LANGUAGE OverloadedStrings #-}
-- |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.DeepSeq
import Data.Maybe(isJust,fromJust,catMaybes)
import Data.Text(Text)
import Control.Parallel.Strategies

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

-- |Gene's interface
-- use 'Text' to represent types
class (NFData d, Show d) => HeukaryaGene d where
  -- |Transform a syntax tree to a expression
  evalTreeGene :: Tree d -> d
  -- |Get Gene's Type representation
  geneTypeRep  :: d -> Text
  -- |Get every possible input type and output type
  --
  -- >>> geneHTypes $ toDyn "(+)" ((+)::Int -> Int -> Int)
  -- [([],"Int -> Int -> Int"),(["Int"],"Int -> Int"),(["Int","Int"],"Int")]
  --
  geneHTypes :: d -> [([Text],Text)]
  -- |get tree's actual output type.
  outputHTreeType :: Tree d -> Text
  -- |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 -> 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 -> Bool
  equalType c a b = (instanceOf c) a b && (instanceOf c) b a
  -- |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 -> Bool
  directType c a b = (instanceOf c) a b || (instanceOf c) b a
  -- |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 -> Maybe Text
  generalizeType c a b = 
    if (instanceOf c) a b then Just a else if (instanceOf c) b a then Just b 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 d) =>
              g             -- ^RandomGenerator
           -> Int           -- ^depth of tree structure
           -> Text          -- ^String representation of output's Type
           -> [d]           -- ^Ingredient(Gene) of Eukarya for constructing Eukarya
           -> Tree d
genEukarya g n outType geneGlobalList = withStrategyParTraversableRdeepseq $
  unfoldTree treeUnfolder (g,n,outType)
  where
  treeUnfolder (g,n,outType) =
    if null geneList3 then treeUnfolder (rr!!0,n,outType) else
      (dynam, zip3 (randomInfiniteList (rr!!1)) (cycle [n-1]) inTypes)
    where
    rr = randomList 3 g
    inTypes = fst $ geneHTypes dynam !! inputNum
    inputNum = fromJust $ findIndex (\y->instanceOf dd y outType) $ (map snd.geneHTypes) dynam
    i = (fst . next) (rr!!2)
    dynam = geneList3 !! (mod i $ length geneList3)
    geneList2 =
      (if n <= 1 then filter (\x -> instanceOf dd (geneTypeRep x) outType) else id)
      geneGlobalList
    geneList3 = filter (\x->isJust $ find (\y->instanceOf dd y outType) $ (map snd.geneHTypes) x) geneList2
    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 d) =>
                g       -- ^RandomGenerator
             -> Int     -- ^depth of tree structure. no functionality here
             -> (Tree d, Tree d) -- ^origin Eukarya pair
             -> (Tree d, Tree d) -- ^crossovered Eukarya pair
crossEukarya g n o@(tree1, tree2) = (\(x,y)->
  (withStrategyParTraversableRdeepseq x, 
   withStrategyParTraversableRdeepseq y)) $
  if null passedSubTrees1 then o else if null passedSubTrees2 then o else
  (replaceSubTree tree1 trail1 subtree2, 
  if (instanceOf dd) (snd choiceSubTree1) (snd choiceSubTree2)
    then replaceSubTree tree2 trail2 subtree1
    else tree2)
  where
    tt = map (outputHTreeType . fst) . getSubTrees
    comOutTypes = 
      catMaybes $ map (\(x,y)->generalizeType dd x y) [(a,b) |a<-(tt tree1),b<-(tt tree2)]
    rr = randomList 2 g
    
    -- tuple subtrees with type
    taggedSubTrees1 = let subtrees = getSubTrees tree1 in zip subtrees (map (outputHTreeType . fst) subtrees)
    passedSubTrees1 = filter (\x -> isJust $
      find ((equalType dd) (snd x)) comOutTypes) taggedSubTrees1
    choice1 = mod ((fst . next) (rr !! 0)) (length passedSubTrees1)
    choiceSubTree1 = passedSubTrees1 !! choice1
    ((subtree1,trail1),_) = choiceSubTree1

    taggedSubTrees2 = let subtrees = getSubTrees tree2 in zip subtrees (map (outputHTreeType . fst) subtrees)
    passedSubTrees2 = filter (\x -> 
      depth (fst$fst x) + depth subtree1 <= n &&
      (instanceOf dd) (snd x) (snd choiceSubTree1)) taggedSubTrees2
    choice2 = mod ((fst . next) (rr !! 1)) (length passedSubTrees2)
    choiceSubTree2 = passedSubTrees2 !! choice2
    ((subtree2,trail2),_) = choiceSubTree2
    dd = rootLabel tree1

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