{-# LANGUAGE DeriveGeneric #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Generics.Putlenses.Examples.Examples
-- Copyright   :  (C) 2013 Hugo Pacheco
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Hugo Pacheco <hpacheco@nii.ac.jp>
-- Stability   :  provisional
--
-- Collection of put-based programming examples
-- 
--
--
----------------------------------------------------------------------------

module Generics.Putlenses.Examples.Examples where

import Generics.Putlenses.Putlens
import Generics.Putlenses.Language
import Generics.Putlenses.TH
import GHC.Generics
import Data.Char as Char
import Data.Maybe as Maybe
import Data.List as Map
import Data.Map (Map(..))
import qualified Data.Map as Map

-- ** Lists

mapPut :: Eq a => Putlens st e b a -> Putlens st e [b] [a]
mapPut f = innPut .< (idPut -|-< f ><< mapPut f) .< outPut

unfoldrPut :: (Eq a,Eq b) => Putlens st e (b,a) a -> a -> Putlens st e [b] a
unfoldrPut f x = innPut .< (ignorePut x -|-< (idPut ><< unfoldrPut f x) .< f) .< injPut (\st e v -> v==x)

unfoldrSPut :: (Eq a,Eq b) => Putlens st e (b,a) a -> a -> Putlens st e [b] a
unfoldrSPut f x = innPut .< (ignorePut x -|-< (idPut ><< unfoldrSPut f x) .< f) .< injSPut

-- *** Constructor/destructor putlenses for lists
nilPut :: Putlens st e [a] ()
nilPut = innPut .< injlPut

consPut :: Putlens st e [a] (a,[a])
consPut = innPut .< injrPut

unnilPut :: Putlens st e () [a]
unnilPut = uninjlPut .< outPut

unconsPut :: Putlens st e (a,[a]) [a]
unconsPut = uninjrPut .< outPut

unheadPut :: Eq a => Putlens st e [a] a
unheadPut = consPut .< keepsndPut

untailPut :: Eq a => Putlens st e [a] [a]
untailPut = consPut .< keepfstPut

wrapPut :: Eq a => Putlens st e [a] a
wrapPut = consPut .< (idPut ><< nilPut) .< addsndOnePut

unwrapPut :: Eq a => Putlens st e a [a]
unwrapPut = remsndOnePut .< (idPut ><< unnilPut) .< unconsPut

-- ** example of automatically deriving constructor/destructor putlenses for binary trees
data Tree a = Empty | Node a (Tree a) (Tree a) deriving Generic
$( makePutlensConstructors ''Tree)

-- ** List concatenation

-- | List concatenation (positional)
catPut :: Eq a => Putlens st e ([a],[a]) [a]
catPut = (innPut ><< idPut) .< undistlPut .< (addfstOnePut -|-< assoclPut .< (idPut ><< catPut) .< unconsPut) .< choiceNil
	where choiceNil = ifVthenelsePut null injlPut injSPut

exCatPut1 = get (put2lens catPut) ([1,2],[3,4])
exCatPut2 = put (put2lens catPut) ([1,2],[3,4]) [0,1,2,3,4]
exCatPut3 = put (put2lens catPut) ([1,2,3,4],[]) [0,1,2]
exCatPut4 = put (put2lens catPut) ([1,2],[3,4]) [0,1]


-- | List concatenation (split the view list at position n-1)
catPutN :: Eq a => (Maybe ([a],[a]) -> [a] -> Int) -> Putlens st e ([a],[a]) [a]
catPutN f = withMbS $ initSt (\_ s v -> f s v) catPutN'
catPutN' :: Eq a => Putlens Int e ([a],[a]) [a]
catPutN' = (innPut ><< idPut) .< undistlPut .< (addfstOnePut -|-< rec) .< (injPut (\st _ _ -> st == 0))
	where rec = modifySt (\st _ _ -> pred st) (assoclPut .< (idPut ><< catPutN') .< unconsPut)

-- | List concatenation (split the view list in half)
catPut2 = catPutN (\s v -> length v `div` 2)

exCatPut21 = get (put2lens catPut2) ([1,2],[3,4])
exCatPut22 = put (put2lens catPut2) ([1,2],[3,4]) [-2,-1,0,1,2,3,4]
exCatPut23 = put (put2lens catPut2) ([1,2,3,4],[]) [0,1,2]
exCatPut24 = put (put2lens catPut2) ([1,2],[3,4]) [0,1]

-- | List concatenation (puts elements to the left while satisfying a predicate)
catPutP :: Eq a => (e -> a -> Bool) -> Putlens st e ([a],[a]) [a]
catPutP p = (innPut ><< idPut) .< undistlPut .< (addfstOnePut -|-< rec) .< (injPut (\_ m vs -> null vs || not (p m $ head vs)))
	where rec = (assoclPut .< (idPut ><< catPutP p) .< unconsPut)
	
-- | List concatenation (puts elements to the left while being equal)
catPutSame :: Eq a => Putlens st e ([a],[a]) [a]
catPutSame = withV' $ catPutP p
    where p [] v = False
          p (x:xs) v = x==v

exCatPutSame1 = get (put2lens catPutSame) ([1,2],[3,4])
exCatPutSame2 = put (put2lens catPutSame) ([1,2],[3,4]) [5,5,5,5,6,7,8]
exCatPutSame3 = put (put2lens catPutSame) ([1,2],[3,4]) [0,1,2]

-- | List concatenation (puts elements to the left while smaller than a particular value)
catPutPred :: Ord a => a -> Putlens st e ([a],[a]) [a]
catPutPred x = catPutP (\i v -> (v <= x))

exCatPutPred1 = get (put2lens $ catPutPred 1) ([1,2],[3,4])
exCatPutPred2 = put (put2lens $ catPutPred 1) ([1,2],[3,4]) [-2,-1,0,1,2,3,4]
exCatPutPred3 = put (put2lens $ catPutPred 1) ([1,2,3,4],[]) [0,1,2]
exCatPutPred4 = put (put2lens $ catPutPred 1) ([1,2],[3,4]) [0,1]

-- ** Integers

succPut :: Putlens st e Int Int
succPut = customPut (\st s -> succ) pred

predPut :: Putlens st e Int Int
predPut = customPut (\st s -> pred) succ

-- ** Naturals

data Nat = ZeroN | SuccN Nat deriving (Eq,Show,Generic)
$( makePutlensConstructors ''Nat)

natPut :: Putlens st e Nat Int
natPut = ifVthenelsePut (==0) (zeroNPut .< ignorePut 0) (succNPut .< natPut .< predPut)

int = get (put2lens natPut)
nat = put (put2lens natPut) ZeroN

-- ** List length

-- | Length as a natural number
lengthNatPut :: Eq a => (Int -> a) -> Putlens st e [a] Nat
lengthNatPut f = innPut .< (idPut -|-< (idPut ><< lengthNatPut f) .< withMbS (addfstPut restoreVal)) .< outPut
    where restoreVal st Nothing n = f (int n)
          restoreVal st (Just (x,_)) _ = x

exLengthNatPut1 = get (put2lens $ lengthNatPut $ \i -> chr (ord 'A' + i)) "abcd"
exLengthNatPut2 = put (put2lens $ lengthNatPut $ \i -> chr (ord 'A' + i)) "abcd" (nat 10)
exLengthNatPut3 = put (put2lens $ lengthNatPut $ \i -> chr (ord 'A' + i)) "abcd" (nat 2)

-- ** List lookup

-- | Embeds a value at a fixed position in a list
embedAtPut :: Eq a => Int -> Putlens st e [a] a
embedAtPut i = fixedi .< embedAtPut'
    where fixedi = remfstPut (\x -> ()) .< (ignorePut i ><< idPut)

embedAtPut' :: Eq a => Putlens st e (Int,[a]) a
embedAtPut' = ifSthenelsePut (\(i,l) -> i == 0) stop it
    where stop = addfstPut (\st m v -> 0) .< unheadPut
          it = (succPut ><< untailPut) .< embedAtPut'

-- | Embeds a value at a fixed position in a list (supports extending the length original list)
embedAtPut1 :: Eq a => Int -> Putlens st e [a] a
embedAtPut1 i = fixedi .< embedAtPut1'
    where fixedi = remfstPut (\x -> ()) .< (ignorePut i ><< idPut)

embedAtPut1' :: Eq a => Putlens st e (Int,[a]) a
embedAtPut1' = initSt (\st e v' -> v') (ifSthenelsePut (\(i,l) -> i == 0) stop it)
    where stop = addfstPut (\st m v -> 0) .< unheadPut'
          it = (succPut ><< untailPut') .< embedAtPut1'
          unheadPut' = consPut .< keepsndOrPut (\x e v -> [])
          untailPut' = consPut .< keepfstOrPut (\x e v -> x)

exEmbedAtPut1 = get (put2lens $ embedAtPut 2) "abcd"
exEmbedAtPut2 = put (put2lens $ embedAtPut 2) "abcd" 'x'
exEmbedAtPut3 = put (put2lens $ embedAtPut 2) "a" 'x'
exEmbedAtPut4 = get (put2lens $ embedAtPut1 2) "a"
exEmbedAtPut5 = put (put2lens $ embedAtPut1 2) "a" 'x'

-- | Embeds a value at a fixed position in a list (splitAt approach)
embedAtPut2 :: Eq a => Int -> Putlens st e [a] a
embedAtPut2 i = fixedi .< splitAtPut .< keepfstPut .< unheadPut
    where fixedi = remfstPut (\x -> ()) .< (ignorePut i ><< idPut)

splitAtPut :: Eq a => Putlens st e (Int,[a]) ([a],[a])
splitAtPut = (stop .\/< rec) .< distlPut .< (outPut ><< idPut)
    where stop = (phiPut (==0) .< newPut 0 ><< idPut)
          rec = (succPut ><< consPut) .< subrPut .< (idPut ><< splitAtPut) .< assocrPut

exEmbedAtPut21 = get (put2lens $ embedAtPut2 3) "abcde"
exEmbedAtPut22 = put (put2lens $ embedAtPut2 3) "abcde" 'x'
exEmbedAtPut23 = put (put2lens $ embedAtPut2 3) "a" 'x'

-- ** List summation

-- | Splits a view number into two summands (by adding an offset to the original first value)
splitPut :: Integral a => (st -> a -> a -> a) -> Putlens st e (a,a) a
splitPut offset = customPut put' (uncurry (+))
    where put' st (Just (x,y)) z = let x' = x + offset st (x + y) z in (x',z - x')
          put' st Nothing z = (z,0)

-- | Updates the sum of a list (preserves the original source and appends the difference to the end)
summandsPut1 :: Integral a => Putlens st e [a] a
summandsPut1 = unfoldrPut splitBy 0
    where splitBy = splitPut (\st e v -> 0)

-- | Updates the sum of a list (distributes the difference by dividing it by two at each recursive step)
-- half of the difference is added to the first element of the source, a quarter to the second, and so on until the remainder is 0
summandsPut2 :: Integral a => Putlens st e [a] a
summandsPut2 = unfoldrPut splitBy 0
    where splitBy = withV (splitPut (\st v v' -> div (v' - v) 2))

-- | Updates the sum of a list (distributes the difference by dividing it by the length of the original list)
-- distributes the difference evenly among original list numbers
summandsPut3 :: Integral a => Putlens st e [a] a
summandsPut3 = withS $ initSt (\st s v -> toEnum (length s)) summandsPut3'
summandsPut3' :: Integral a => Putlens a e [a] a
summandsPut3' = unfoldrPut (updateSt (\st s v -> pred st) splitBy) 0
    where splitBy = withV (splitPut (\st v v' -> div (v' - v) st))
	
-- | Updates the sum of a list (distributes the difference by dividing it by the length of the original list, always preserving the size f the original list even when the view is zero)
summandsPut4 :: Integral a => Putlens st e [a] a
summandsPut4 = withS $ initSt (\_ s v -> toEnum (length s)) summandsPut4'
summandsPut4' :: Integral a => Putlens a e [a] a
summandsPut4' = unfoldrSPut (updateSt (\st s v -> pred st) splitBy) 0
    where splitBy = withV (splitPut (\st v v' -> div (v' - v) st))
	
exSummandsPut1 = get (put2lens summandsPut1) [1,2,3,4]
exSummandsPut2 = put (put2lens summandsPut1) [1,2,3,4] 15
exSummandsPut3 = put (put2lens summandsPut2) [1,2,3,4] 15
exSummandsPut4 = put (put2lens summandsPut3) [1,2,3,4] 15
exSummandsPut5 = put (put2lens summandsPut3) [1,2,3,4] 0
exSummandsPut6 = put (put2lens summandsPut4) [1,2,3,4] 15
exSummandsPut7 = put (put2lens summandsPut4) [1,2,3,4] 0

-- ** Replicate

-- | Replicate
replicatePut :: Eq a => Putlens st e (a,Int) [a]
replicatePut = (zero .\/< (one .\/< rec) .< distrPut .< (idPut ><< outPut)) .< outPut
    where zero = keepfstPut .< phiPut (==0) .< newPut 0
          one = idPut ><< phiPut (==1) .< newPut 1
          rec = (idPut ><< succPut) .< remfstPut fst .< (idPut ><< replicatePut .< consPut)

exReplicatePut1 = put (put2lens replicatePut) ('b',2) "aaaa"
exReplicatePut2 = put (put2lens replicatePut) ('b',2) ""

-- | Replicates a list of elements into a sequence of replicated elements
replicateListPut :: Eq a => Putlens st e [(a,Int)] [a]
replicateListPut = recoverzerosPut .< mapPut replicatePut .< splitListPut

recoverzerosPut :: Eq a => Putlens st e [(a,Int)] [(a,Int)]
recoverzerosPut = withMbS $ ifthenelsePut (\st s v -> isJust s && p (fromJust s)) zero rec
    where p ((a,0):xs) = True
          p _ = False
          zero = phiPut p .< consPut .< (idPut ><< recoverzerosPut) .< addfstPut (\st (Just (x:xs)) v -> x)
          rec = innPut .< (idPut -|-< idPut ><< recoverzerosPut) .< outPut

splitListPut :: Eq a => Putlens st e [[a]] [a]
splitListPut = innPut .< (idPut -|-< ifVthenelsePut p eq neq .< (idPut ><< splitListPut)) .< outPut
    where p (x,(x':xs):xss) = x == x'
          p (x,xss) = False
          eq = (consPut ><< idPut) .< assoclPut .< (idPut ><< unconsPut)
          neq = wrapPut ><< idPut

exReplicateListPut1 = get (put2lens replicateListPut) [('a',4),('b',0),('c',3)]
exReplicateListPut2 = put (put2lens replicateListPut) [('a',4),('b',0),('c',3)] "xxyyz"

-- ** Halve

-- | Takes the first half of a list (with a default empty element)
halvePut :: Eq a => a -> Putlens st e [a] [a]
halvePut x = remfstPut half .< splitAtPut .< withS (addsndPut right)
    where half xs = let d = length xs
	                in if even d then d `div` 2 else (d `div` 2)+1
          right st (l,r) v = let s = take (length v) r
                             in s ++ replicate (length v - length s - 1) x

exHalvePut1 = put (put2lens $ halvePut ' ') "abcde" "xy"
exHalvePut2 = put (put2lens $ halvePut '0') "xyz" "abcde"

-- | Takes the first half of a list (using an increasing counter with each consumed element in the forward direction)
halvePut2 :: Eq a => a -> Putlens st e [a] [a]
halvePut2 x = remsndPut (const 0) .< halvePut2' x .< addsndPut (\st e -> length)
	
halvePut2' :: Eq a => a -> Putlens st e ([a],Int) ([a],Int)
halvePut2' x = withMbS $ (phiPut ifcond .< empty .\/< rec) .< distlPut .< (outPut ><< idPut)
    where ifcond (xs,k) = k >= length xs
          empty = addfstPut appendRight .< remfstOnePut
          appendRight st (Just (rs,_)) k = if k >= length rs then take (pred k) (rs++repeat x) else take k rs
          appendRight st Nothing k = take (pred k) (repeat x)
          rec = (consPut ><< predPut) .< assoclPut .< (idPut ><< halvePut2' x) .< assocrPut
          
exHalvePut21 = put (put2lens $ halvePut2 ' ') "abcde" "xy"
exHalvePut22 = put (put2lens $ halvePut2 '0') "xyz" "abcde"
exHalvePut23 = put (put2lens $ halvePut2 ' ') "abc" "xy"

-- | Takes the first half of a list (using a decreasing counter with the size of the input list, decreased 2 by 2, in the forward direction)
halvePut3 :: Eq a => a -> Putlens st e [a] [a]
halvePut3 x = initSt (\_ _ _ -> 0) $ remsndPut length .< halvePut3' x

halvePut3' :: Eq a => a -> Putlens Int e ([a],Int) [a]
halvePut3' x = withMbS $ ifVthenelsePut null (nil .< unnilPut) (cons .< unconsPut)
    where nil = phiPut (\(xs,k) -> k <= 0) .< pntPut appendRight
          appendRight k (Just (rs,_)) = if k >= length rs then (take (pred k) (rs++repeat x),-1) else (take k rs,0)
          appendRight k Nothing = (take (pred k) (repeat x),-1)
          cons = (consPut ><< succPut .< succPut) .< assoclPut .< (modifySt (\i _ _ -> succ i) idPut ><< halvePut3' x)

exHalvePut31 = put (put2lens $ halvePut3 ' ') "abcde" "xy"
exHalvePut32 = put (put2lens $ halvePut3 '0') "xyz" "abcde"
exHalvePut33 = put (put2lens $ halvePut3 ' ') "abc" "xy"

-- | Incremental summation
isumPut :: Putlens st e [Int] [Int]
isumPut = innPut .< (idPut -|-< (idPut ><< isumPut) .< paramfstPut dec) .< outPut
	where dec i = mapPut (subtractPut i)

subtractPut :: Int -> Putlens st e Int Int
subtractPut i = customPut (\st s x -> x-i) (+i)

exIsumPut1 = get (put2lens isumPut) [1,2,3,4]
exIsumPut2 = put (put2lens isumPut) [1,2,3,4] [3,4]

-- ** Sorting

-- | Insertion sort (put according to the original relative source order)
iunsortPut1 :: Ord a => Putlens st e [a] [a]
iunsortPut1 = ifVthenelsePut null (nilPut .< unnilPut) it
	where it = consPut .< (idPut ><< iunsortPut1) .< delPut1

delPut1 :: Ord a => Putlens st e (a,[a]) [a]
delPut1 = ifVthenelsePut (null . snd) idPut disarrange .< unconsPut
  where  disarrange = ifSthenelsePut p idPut it
         p (x,ys) = x <= head ys
         it = (idPut ><< consPut) .< subrPut .< (idPut ><< delPut1)

-- | Insertion sort (identity backward transformation)
iunsortPut2 :: Ord a => Putlens st e [a] [a]
iunsortPut2 = ifVthenelsePut null (nilPut .< unnilPut) it
	where it = consPut .< (idPut ><< iunsortPut2) .< delPut2

delPut2 :: Ord a => Putlens st e (a,[a]) [a]
delPut2 = ifVthenelsePut (null . snd) idPut disarrange .< unconsPut
  where  disarrange = eitherSPut p idPut it .< injPut (\st e v -> True)
         p (x,ys) = x <= head ys
         it = (idPut ><< consPut) .< subrPut .< (idPut ><< delPut2)

exIunsortPut1 = put (put2lens iunsortPut1) [4,1,3,2] [5,6,7,8,9,10]
exIunsortPut2 = put (put2lens iunsortPut2) [4,1,3,2] [5,6,7,8,9,10]
exIunsortPut3 = put (put2lens iunsortPut2) [4,1,3,2] [1,2,3,4]

-- | Quicksort (put according to the original relative source order)
qsortPut :: Ord a => Putlens st e [a] [a]
qsortPut = innPut .< (unnilPut -|-< partitionPut .< (idPut ><< (qsortPut ><< qsortPut)) .< subrPut .< (idPut ><< unconsPut) .< catPutNonEmptyRight) .< (injPut (\_ _ -> null))

-- | Partition a list into smaller and bigger elements than a given element
partitionPut :: Ord a => Putlens st e (a,[a]) (a,([a],[a]))
partitionPut = customPut (\st s (y,(xs,ys)) -> (y,xs++ys)) (\(x,xs) -> (x,partition (<x) xs))

-- List concatenation (put guarantees that the right side has always one element)
catPutNonEmptyRight :: Eq a => Putlens st e ([a],[a]) [a]
catPutNonEmptyRight = catPutN f
	where f Nothing v = 0
	      f (Just (xs,y)) v = if length xs > length v then 0 else length xs

exQsortPut1 = get (put2lens qsortPut) [4,3,1,2]
exQsortPut2 = put (put2lens qsortPut) [4,1,3,2] [-4,-3,-2,-1,0,7]
exQsortPut3 = put (put2lens qsortPut) [4,1,3,2] [5,6,7,8,9,10]

-- ** Maximum segment sum

type Pos = Int

segments :: [a] -> [[a]]
segments = concat . map inits . tails

-- | Updating maximum segment sum
-- when the sum increases, update only the largest segment
-- when the sum decreases, update all segments that surpass the new maximum, from smallest to largest
mssPut :: Putlens st e [Int] Int
mssPut = withV (ifKthenelsePut (\st v v' -> v < v') (mssPut' True) (mssPut' False))
mssPut' b = positionsPut .< remsndPut (sortedsegments b) .< (fromListPut ><< idPut) .< maxSumSegsPut
  where sortedsegments b = sortBy (comparelength b) . filter (/=[]) . segments . map fst
        comparelength b x y = if b then compare (length y) (length x) else compare (length x) (length y)

positionsPut :: Eq a => Putlens st e [a] [(Pos,a)]                                              
positionsPut = remfstPut (const 0) .< withV' positionsPut' 
positionsPut' :: Eq a => Putlens st [(Pos,a)] (Int,[a]) [(Pos,a)]                              
positionsPut' = (idPut ><< innPut) .< undistrPut .< (empty -|-< it) .< outPut          
  where  empty = addfstPut (\st m v -> length m)                                       
         it = (remsndPut succ ><< idPut) .< distpPut .< (idPut ><< positionsPut')  
                                                                                     
maxSumSegsPut :: Putlens st e ((Map Pos Int,[[Pos]])) Int                             
maxSumSegsPut = initSt (\st e m' -> False) maxSumSegsPut'                           
maxSumSegsPut' = (idPut ><< innPut) .< undistrPut .< (empty -|-< it) .< injSPut          
  where  empty = keepfstPut .< modifyV' changeV (ignorePut 0)                                       
         changeV st e v = if st then 0 else v
         it = unforkPut ((idPut ><< keepsndPut) .< lookupSegPut .< summandsPut4) ((idPut ><< keepfstPut) .< maxSumSegsPut') .< maxPut

lookupSegPut :: Eq a => Putlens st e ((Map Int a,[Int])) [a]
lookupSegPut = (idPut ><< innPut) .< undistrPut .< (empty -|-< it) .< outPut
    where empty = keepfstPut
          it = unforkPut ((idPut ><< keepsndPut) .< swapPut .< lookupIntMapPut) ((idPut ><< keepfstPut) .< lookupSegPut)

lookupIntMapPut :: Eq a => Putlens st e ((Int,Map Int a)) a
lookupIntMapPut = (idPut ><< toListPut) .< embedAtPut' .< keepfstPut

maxPut :: (Num a,Ord a) => Putlens Bool e ((a,a)) a
maxPut = withS (ifthenelsePut (\st (x,y) m -> st && x < m) left right)                  
  where  left = phiPut (\(x,y) -> x < y) .< addfstPut (\st (x,y) m -> x) 
         right = updateSt (\st e v -> True) (addsndPut (\st e m -> m))

fromListPut :: (Eq a,Ord k) => ((Putlens st e [(k,a)] (Map k a)))
fromListPut = customPut (\st s -> Map.toList) (Map.fromList)

toListPut :: (Eq a,Ord k) => ((Putlens st e (Map k a) [(k,a)]))
toListPut = customPut (\st s -> Map.fromList) (Map.toList)

exMssPut1 = get (put2lens mssPut) [1,2,-3,4]
exMssPut2 = put (put2lens mssPut) [1,2,-3,4] 6
exMssPut3 = put (put2lens mssPut) [1,2,-3,4] 2
exMssPut4 = put (put2lens mssPut) [1,2,-3,4] (-4)
exMssPut5 = put (put2lens mssPut) [1,2,-3,4] 0