-----------------------------------------------------------------------------
-- |
-- Module      :  Generics.Putlenses.Examples.Color
-- Copyright   :  (C) 2014 Hugo Pacheco
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Hugo Pacheco <hpacheco@nii.ac.jp>
-- Stability   :  provisional
--
-- Graph coloring examples.
-- 
--
--
----------------------------------------------------------------------------

module Generics.Putlenses.Examples.Color where

import Data.Int
import Test.QuickCheck hiding ((><))
import Test.QuickCheck.Gen
import Test.QuickCheck.Random

import Control.Monad
import Data.Maybe
import Data.List as List
import Control.Exception.Base

import Generics.Putlenses.Language
import Generics.Putlenses.Putlens
import Generics.Putlenses.Examples.Examples

import Control.Monad.Reader (Reader(..),ReaderT(..))
import qualified Control.Monad.Reader as Reader
import Control.Monad.State (State(..),StateT(..))
import qualified Control.Monad.State as State

-- * Colored graphs and utility functions

data Color = Red | Yellow | Green | Blue deriving (Eq,Ord,Show)
type Vertex = Int
type ColoredGraph = [((Vertex,Color),Edges)]

type Edges = [Vertex]
type Graph = [(Vertex,Edges)]

lookupColor :: MonadPlus m => Vertex -> ColoredGraph -> m Color
lookupColor v [] = mzero
lookupColor v (((v',c),edges):g) = if v == v' then return c else lookupColor v g

colors = [Red,Yellow,Green,Blue]

red,yellow,green,blue :: MonadPlus m => m Color
red = return Red
yellow = return Yellow
green = return Green
blue = return Blue

genColors :: MonadPlus m => (Color -> Bool) -> [Color] -> m Color
genColors p [] = mzero
genColors p (c:cs) = mfilter p (return c) `mplus` genColors p cs

allColors :: MonadPlus m => (Color -> Bool) -> m Color
allColors p = genColors p colors

uncolor :: ColoredGraph -> Graph
uncolor = map (\(x,y) -> (fst x,y))

differentColor :: Color -> ColoredGraph -> Vertex -> Bool
differentColor c g e = maybe True (/=c) (lookupColor e g)

isCorrect :: ColoredGraph -> Bool
isCorrect [] = True
isCorrect (((v,c),edges):g) = b && isCorrect g
	where b = List.and $ map (differentColor c g) edges

isCorrectNode :: ((Vertex,Color),Edges) -> ColoredGraph -> Bool
isCorrectNode ((v,c),edges) g = b
	where b = List.and $ map (differentColor c g) edges

--recolor2 :: MonadPlus m => ColoredGraph -> Graph -> m ColoredGraph
--recolor2 s [] = return []
--recolor2 s ((v,edges):g) = do
--	cg <- recolor2 s g
--	c <- mkColor (\c -> isCorrectNode ((v,c),edges) cg) v s
--	return $ ((v,c),edges) : cg

mkColor :: MonadPlus m => (Color -> Bool) -> Vertex -> ColoredGraph -> m Color
mkColor p v s = do { color <- lookupColor v s; genColors p (color : (colors \\ [color])) } `mplus` allColors p

---- a more naive version

-- colors a graph, basically repairing the original colors if they are erroneous
--recolor1 :: MonadPlus m => ColoredGraph -> Graph -> m ColoredGraph
--recolor1 s v = mfilter isCorrect (recolor' s v)

--recolor' :: MonadPlus m => ColoredGraph -> Graph -> m ColoredGraph
--recolor' s v = mapM (\(vertex,edges) -> liftM (\c -> ((vertex,c),edges)) (mkColor' vertex s)) v

mkColorNaive :: MonadPlus m => Vertex -> ColoredGraph -> m Color
mkColorNaive v s = do { color <- lookupColor v s; genColorsNaive (color : (colors \\ [color])) } `mplus` allColorsNaive

genColorsNaive :: MonadPlus m => [Color] -> m Color
genColorsNaive [] = mzero
genColorsNaive (c:cs) = (return c) `mplus` genColorsNaive cs

allColorsNaive :: MonadPlus m => m Color
allColorsNaive = genColorsNaive colors

--- tests

--exg :: Graph
--exg = [(1,[2,3]),(2,[1,3,4]),(3,[1,2,4]),(4,[2,3])]

--excg :: ColoredGraph
--excg = [((1,Red),[2,3]),((2,Blue),[1,3,4]),((3,Yellow),[1,2,4]),((4,Yellow),[2,3])]

--test :: Maybe ColoredGraph
--test = recolor2 excg exg
  

-- arbitrary graph generation

genGraph :: Int -> Int -> Gen Graph
genGraph i j = do
	let vertices = [0..i]
	edges <- genEdges vertices j
	return $ genGraph' vertices edges
	
genGraph' :: [Int] -> [(Int,Int)] -> Graph
genGraph' [] [] = []
genGraph' (v:vs) es = (v,ves) : genGraph' vs es'
	where (ves,es') = (\(x,y) -> (map snd x,y)) $ partition (\(i,j) -> i == v) es
	
genEdges :: [Int] -> Int -> Gen [(Int,Int)]
genEdges vs 0 = return []
genEdges vs i = do
	(x,y) <- genEdge vs
	es <- genEdges vs (i-1)
	return ((x,y):(y,x):es)
genEdge :: [Int] -> Gen (Int,Int)
genEdge vs = do
	i <- elements vs
	j <- elements (vs \\ [i])
	return $ (i,j)
	
-- * Lenses

-- | Putlens that colors a graph in the backward direction
colorPut :: MonadPlus m => PutlensM m ColoredGraph Graph
colorPut = runReaderPut (\mg _ -> return $ maybe [] id mg) colorPut' where
	colorPut' :: MonadPlus m => PutlensReaderM m ColoredGraph ColoredGraph Graph
	colorPut' = foldrPut $ innPut .< (idPut -|-< (sublPut ><< idPut) .< sublPut .< addsndPutUnsafe f)
		where f _ ((v,es),cg) = do { s <- Reader.ask; mkColor (\c -> isCorrectNode ((v,c),es) cg) v s }

-- | Putlens that colors a graph in the backward direction (slow version)
colorPutNaive :: MonadPlus m => PutlensM m ColoredGraph Graph
colorPutNaive = runReaderPut (\mg _ -> return $ maybe [] id mg) colorPutNaive' where
	colorPutNaive' :: MonadPlus m => PutlensReaderM m ColoredGraph ColoredGraph Graph
	colorPutNaive' = mfilterPut isCorrect $ mapPut (keepsndOrPut colorVertex ><< idPut)
		where colorVertex v = Reader.ask >>= mkColorNaive v

testColor :: IO ()
testColor = do
	stdgen <- newQCGen
	let g = unGen (genGraph 10 20) stdgen 1
	putStrLn "colorPut"
	res <- evaluate $ take 1 $ put (put2lensM (colorPut :: PutlensM [] ColoredGraph Graph)) [] g
	putStrLn "colorPutNaive"
	resNaive <- evaluate $ take 1 $ put (put2lensM (colorPutNaive :: PutlensM [] ColoredGraph Graph)) [] g
	return ()