module Game.Chess.Tree (
  positionTree, positionForest
, plyTree, plyForest
, pathTree
) where

import Data.List.NonEmpty (NonEmpty, cons)
import Data.Tree ( Tree(Node), Forest, foldTree )
import Game.Chess.Internal

positionTree :: Position -> Tree Position
positionTree :: Position -> Tree Position
positionTree Position
pos = Position -> Forest Position -> Tree Position
forall a. a -> Forest a -> Tree a
Node Position
pos (Forest Position -> Tree Position)
-> Forest Position -> Tree Position
forall a b. (a -> b) -> a -> b
$ Position -> Forest Position
positionForest Position
pos

positionForest :: Position -> Forest Position
positionForest :: Position -> Forest Position
positionForest Position
pos = Position -> Tree Position
positionTree (Position -> Tree Position)
-> (Ply -> Position) -> Ply -> Tree Position
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Position -> Ply -> Position
unsafeDoPly Position
pos (Ply -> Tree Position) -> [Ply] -> Forest Position
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Position -> [Ply]
legalPlies Position
pos

plyForest :: Position -> Forest Ply
plyForest :: Position -> Forest Ply
plyForest Position
pos = Position -> Ply -> Tree Ply
plyTree Position
pos (Ply -> Tree Ply) -> [Ply] -> Forest Ply
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Position -> [Ply]
legalPlies Position
pos

plyTree :: Position -> Ply -> Tree Ply
plyTree :: Position -> Ply -> Tree Ply
plyTree Position
pos Ply
ply = Ply -> Forest Ply -> Tree Ply
forall a. a -> Forest a -> Tree a
Node Ply
ply (Forest Ply -> Tree Ply)
-> (Position -> Forest Ply) -> Position -> Tree Ply
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Position -> Forest Ply
plyForest (Position -> Tree Ply) -> Position -> Tree Ply
forall a b. (a -> b) -> a -> b
$ Position -> Ply -> Position
unsafeDoPly Position
pos Ply
ply

pathTree :: Tree a -> Tree (NonEmpty a)
pathTree :: Tree a -> Tree (NonEmpty a)
pathTree = (a -> [Tree (NonEmpty a)] -> Tree (NonEmpty a))
-> Tree a -> Tree (NonEmpty a)
forall a b. (a -> [b] -> b) -> Tree a -> b
foldTree ((a -> [Tree (NonEmpty a)] -> Tree (NonEmpty a))
 -> Tree a -> Tree (NonEmpty a))
-> (a -> [Tree (NonEmpty a)] -> Tree (NonEmpty a))
-> Tree a
-> Tree (NonEmpty a)
forall a b. (a -> b) -> a -> b
$ \a
a -> NonEmpty a -> [Tree (NonEmpty a)] -> Tree (NonEmpty a)
forall a. a -> Forest a -> Tree a
Node (a -> NonEmpty a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a) ([Tree (NonEmpty a)] -> Tree (NonEmpty a))
-> ([Tree (NonEmpty a)] -> [Tree (NonEmpty a)])
-> [Tree (NonEmpty a)]
-> Tree (NonEmpty a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Tree (NonEmpty a) -> Tree (NonEmpty a))
-> [Tree (NonEmpty a)] -> [Tree (NonEmpty a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Tree (NonEmpty a) -> Tree (NonEmpty a))
 -> [Tree (NonEmpty a)] -> [Tree (NonEmpty a)])
-> ((NonEmpty a -> NonEmpty a)
    -> Tree (NonEmpty a) -> Tree (NonEmpty a))
-> (NonEmpty a -> NonEmpty a)
-> [Tree (NonEmpty a)]
-> [Tree (NonEmpty a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NonEmpty a -> NonEmpty a)
-> Tree (NonEmpty a) -> Tree (NonEmpty a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (a -> NonEmpty a -> NonEmpty a
forall a. a -> NonEmpty a -> NonEmpty a
cons a
a)