{-# LANGUAGE MultiParamTypeClasses, AllowAmbiguousTypes #-}

module Control.Metaheuristics(
  HasQuality,quality,
  -- * Termination and post processing
  takeWhileProperty,
  -- * Execution of search processes
  openStreamMapArrow,
  executeMetaheuristic,
  -- * Simple lifting and injection
  mapArrow,
  inject,
  -- * General stream arrow combinators
  chunk,
  dechunk,
  stretch,
  doMany,
  delay,
  window,
  fanAll,
  newArrow,
  -- * Selection and filtering
  improvingFilter,
  improving,
  selectBest,
  selectRandom,
  diffFilter,
  -- * Escape Strategies
  escapeStrategy,
  replace,
  restart,
  uniformChoice,
  
  -- * SA cooling strategies
  logCooling,
  linCooling,
  geoCooling,
  
  -- * Metaheuristic algorithms
  iterativeImprover,
  stochasticIterativeImprover,
  simpleTABU,
  tabu,
  mergeStochasticAndTemperature,
  simulatedAnnealing,
  breed1,ga,
  populationPattern,degrade,aco
  ) where

import Control.Arrow
import Data.Stream (Stream(..),(<:>))
import qualified Data.Stream as S
import Control.Arrow.Transformer.Stream (StreamArrow(..),StreamMap(..))
import System.Random
import qualified Data.List as L

class Num v=>HasQuality f v where
  quality :: f->v

takeWhileProperty :: ([a]->Bool)->Int ->Stream a->[a]
takeWhileProperty term sz input = twi' input
  where
    -- twi' :: Stream a->[a]
    twi' xs = if term (S.take sz xs)
                then S.take sz xs
                else (S.head xs) : twi' (S.tail xs)


openStreamMapArrow :: StreamMap a b->Stream a->Stream b
openStreamMapArrow (StreamArrow f) = f

executeMetaheuristic :: StreamArrow (->) b b->[b]->Stream b 
executeMetaheuristic a x = let xs = S.prefix x (openStreamMapArrow a xs) in xs

newArrow :: (Stream a->Stream b)->StreamMap a b
newArrow = StreamArrow

mapArrow :: (a->b)->StreamMap a b
mapArrow f = StreamArrow (S.map f)

inject :: Stream a->(a->b->c)->StreamMap b c
inject s f = StreamArrow (S.zipWith f s)

chunk :: Int->StreamMap a [a]
chunk n = StreamArrow (f ) 
  where f ys = S.take n ys `S.Cons` (f $ S.drop n ys)


dechunk :: StreamMap [a] a
dechunk = StreamArrow (S.fromList . concat. S.toList )


stretch:: Int->StreamMap sol sol
stretch n = StreamArrow 
  (S.fromList  . concatMap (replicate n) .    S.toList )


doMany :: Int->StreamMap sol sol'->StreamMap sol sol'
doMany n f = stretch n >>> f

delay::x->StreamMap x x
delay x=StreamArrow(S.Cons x)

window :: Int->StreamMap sol [sol]
window n=StreamArrow(\xs->S.scan f [] xs)
  where f qu x=if length qu==n then f (tail qu) x else qu++ [x]

fanAll :: [StreamMap a b]->StreamMap a [b]
fanAll fs = StreamArrow (\xs->S.fromList $ L.transpose [S.toList (f xs) | f<-map openStreamMapArrow fs       ])

improvingFilter :: Ord a=>StreamMap ([a],a) [a] 
improvingFilter = StreamArrow $ S.map (\(xs,x)->filter (<x) xs)


improving :: Ord a=>StreamMap a [a]->StreamMap a [a]
improving nF = nF &&& (StreamArrow id)  >>> improvingFilter

historyFilter :: Eq a => StreamMap ([a],[a]) [a]
historyFilter = mapArrow (\(n,h)->n L.\\ h)

selectBest :: Ord a=>StreamMap  [a] a 
selectBest = StreamArrow $ S.map minimum

selectRandom :: RandomGen g=>g->StreamMap [a] a
selectRandom seedG = StreamArrow (f seedG)
  where f g xss = let xs = S.head xss
                      (i,g') = randomR (0,length xs-1) g
                  in xs !! i <:> f g' (S.tail xss)


replace :: Stream sol ->StreamMap sol sol
replace supply = StreamArrow (const supply)

escapeStrategy :: StreamMap sol (Either sol sol)->StreamMap sol sol->StreamMap sol sol->StreamMap sol sol
escapeStrategy escapeTest s1 s2 = escapeTest >>> (s1 ||| s2)

restart :: StreamMap sol (Either sol sol)->StreamMap sol sol->Stream sol->StreamMap sol sol
restart restartTest s1 solSupply = escapeStrategy restartTest s1 (replace solSupply)

iterativeImprover :: Ord sol=>StreamMap sol [sol]->StreamMap sol sol
iterativeImprover neighbourhoodArrow = improving neighbourhoodArrow >>> selectBest

stochasticIterativeImprover :: (Ord sol,RandomGen g)=>g->StreamMap sol [sol]->StreamMap sol sol
stochasticIterativeImprover g neighbourhoodArrow = improving neighbourhoodArrow >>> selectRandom g

simpleTABU :: Ord sol=>StreamMap sol [sol]->Int->StreamMap sol sol 
simpleTABU nF n = nF &&& (window n) >>> historyFilter >>> selectBest

tabu :: Ord sol=>StreamMap sol [sol]->Int->StreamMap sol sol 
tabu nF n = fanAll [ improving nF,
                     nF &&& window n >>> historyFilter,
                     nF] >>> mapArrow concat >>> mapArrow head

mergeStochasticAndTemperature:: Floating v=>S.Stream v -> S.Stream v -> S.Stream v
mergeStochasticAndTemperature = S.zipWith (\r t ->1.0/t - log r)

--diffChoice :: (Num a, Ord a, HasQuality f) => a -> (f, f) -> f
--diffChoice r (a,b) = if r >= quality a - quality b then b else a

--sa :: HasQuality sol => StreamMap sol sol->S.Stream Double->S.Stream Double -> StreamMap sol sol
--sa mutate rs ts = (StreamArrow id)&&&mutate >>> inject (mergeStochasticAndTemperature rs ts) diffChoice

diffFilter :: (Num a, Ord a, HasQuality f a) => a -> (f, [f]) -> [f]
diffFilter r (s,xs) = [x | x<-xs, r >= quality s - quality x]

simulatedAnnealing :: (Floating v,Num v,Ord v, Ord sol, HasQuality sol v) =>StreamMap sol [sol] -> S.Stream v -> S.Stream v -> StreamMap sol sol
simulatedAnnealing neighbourhood rs ts = fanAll [ improving neighbourhood,
                                                  (StreamArrow id)&&&neighbourhood >>> inject (mergeStochasticAndTemperature rs ts) diffFilter,
                                                  mapArrow (:[]) ] >>> mapArrow concat >>> mapArrow head

{-| A logarithmic cooling strategy intended for use within simulated annealing. Broadly the first value is 
    the starting temperature and the second a value between 0 and 1. -}
logCooling :: Floating b=>b->b->[b]
logCooling c d = map (\t->c / (log (t + d))) (iterate (+1) 1)

{-| The most common cooling strategy for simulated annealing, geometric. The first value is the starting temperature, 
    the second a value between 0 and 1, the cooling rate.  -}
geoCooling :: Floating b=>b->b->[b]
geoCooling startTemp tempChange = iterate (* tempChange) startTemp

{-| Included for completeness, this is a cooling strategy for simulated annealing that is usually not very effective,
    a linear changing strategy. The first value is the starting temperature the second is the value to increase it by 
    at each step. In order to have it reduce at each step, pass a negative value. 
-}
linCooling :: Floating b=>b->b->[b]
linCooling startTemp tempChange= iterate (+ tempChange) startTemp

breed1 :: RandomGen g=>StreamMap [sol] sol->g->StreamMap [sol] sol
breed1 recombine g1 = doMany 2 (selectRandom g1) >>> chunk 2 >>> recombine

ga :: Int ->StreamMap sol (Either sol sol)->StreamMap sol sol ->StreamMap [sol] sol->StreamMap sol sol
ga popSize mutateTest mutate breed = chunk popSize >>> doMany popSize breed >>> mutateTest >>> (StreamArrow id ||| mutate)


uniformChoice :: RandomGen g=>Float->g->StreamMap sol (Either sol sol)
uniformChoice chance g = StreamArrow (S.fromList . f (randoms g) . S.toList)
  where f (x:xs) (s:sols) | x<chance = Left s:f xs sols
                          | otherwise = Right s : f xs sols

populationPattern::Int->(sol->pat)->StreamMap [pat] pat-> StreamMap sol pat 
populationPattern popSize createPattern mergePattern=mapArrow createPattern>>>chunk popSize>>>mergePattern>>>stretch popSize

degrade :: StreamMap a a -> StreamMap (a,a) a->StreamMap a a 
degrade df mf = loop (second df >>> mf >>> (StreamArrow id &&& StreamArrow id))

aco :: pat->Int->(sol -> pat)->StreamMap [pat] pat ->StreamMap (pat, pat) pat->StreamMap pat pat ->StreamMap pat sol->StreamMap sol sol
aco defaultPattern popsize cfp mergePopulation mergePair df createsolution = mapArrow cfp >>> chunk popsize >>> mergePopulation >>> degrade (delay defaultPattern >>> df ) mergePair >>> stretch popsize >>> createsolution