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

module Generics.Putlenses.Examples.Graph where

import System.Random
import Generics.Putlenses.Putlens
import Generics.Putlenses.Language
import Generics.Putlenses.Examples.Examples
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.Monad.Reader
import Control.Monad.State (State(..),StateT(..))
import qualified Control.Monad.State as State
import Control.Exception.Base
import Control.Monad.Identity
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.Identity

type Vertex = Int

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

-- 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
	e <- genEdge vs
	es <- genEdges vs (i-1)
	return (e:swap e:es)
  where swap (x,y) = (y,x)
genEdge :: [Int] -> Gen (Int,Int)
genEdge vs = do
	i <- elements vs
	j <- elements (vs \\ [i])
	return $ (i,j)


deleteVertex :: (Monad m) => Vertex -> PutlensM m Graph Graph
deleteVertex v = filterPut ((/=v) . fst)  .< mapPut (idPut ><< filterPut (/=v))

testDelete = do
	let lns = put2lensM (deleteVertex 3 :: PutlensM Maybe Graph Graph)
	get lns exg

genGraphPut :: PutlensM (MaybeT Gen) Graph (Int,Int)
genGraphPut = (verticesPut .< lengthVerticesPut) `unforkPut` edgesPut

edgesPut :: PutlensM (MaybeT Gen) Graph Int
edgesPut = runReaderPut (\s v -> return $ maybe [] (map fst) s) $ mapPut (keepfstPut .< lengthEdgesFromPut) .< sumEdgesPut

lengthEdgesFromPut :: PutlensM (ReaderT [Vertex] (MaybeT Gen)) [Vertex] Int
lengthEdgesFromPut = unfoldrPut (keepfstOrPut f .< predPut) 0
	where f i = ask >>= lift . lift . elements

verticesPut :: PutlensM (MaybeT Gen) Graph [Vertex]
verticesPut = runReaderPutV' $ mapPut (addsndPutUnsafe $ \s v -> ask >>= \vs -> return $ filter (\x -> elem x vs) $ maybe [] snd s)

lengthVerticesPut :: PutlensM (MaybeT Gen) [Vertex] Int
lengthVerticesPut = runStatePut (\s v -> return []) $ unfoldrPut (updateStatePut (\s (v,_) vs -> return $ v:vs) $ keepfstOrPut f .< predPut) 0
	where f i = State.get >>= \vs -> lift $ lift $ suchThat positiveint (\x -> not $ elem x vs)

-- | Updates the sum of a list (distributes the difference by dividing it by the length of the original list, always preserving the size of the original list even when the view is zero)
-- the source always contains positive numbers
sumEdgesPut :: PutlensM (ReaderT [Vertex] (MaybeT Gen)) [Int] Int
sumEdgesPut = (nilPut .< ignorePut 0) `unionPut` unfoldr1Put (splitBy)
    where splitBy = splitPut $ \(x,y) z -> do
		ask >>= \vs -> lift $ lift $ suchThat (elements [-x..(z-x)]) (\w -> x+w <= length vs) -- we want to keep the length

exg = [(1, [2, 3]), (2, [1, 3]), (3, [2])]

testGraph = do
	let lns = put2lensM genGraphPut
	generate $ runMaybeT $ put lns exg (5,6)

testSum = do
	generate $ runMaybeT $ runReaderT (put (put2lensM sumEdgesPut) [1,2] 6) [1,2]

anyint :: Gen Int
anyint = choose (minBound,maxBound)

positiveint :: Gen Int
positiveint = choose (0,50)