{-# LANGUAGE NamedFieldPuns #-}
module Language.Synthesis.Synthesis (
    Mutation, synthesizeMhList, runningBest, Problem(..)
) where

import           Control.Monad.Random

import           Language.Synthesis.Distribution (Distr)
import qualified Language.Synthesis.Distribution as Distr
import           Language.Synthesis.MCMC
import           Language.Synthesis.Mutations    (Mutation)

-- | This type specifies which program to synthesize. It comes with a
-- specification, which is a program that already works, some inputs
-- and a distance function.
data Problem program = Problem { score :: program -> Double
                               , prior :: Distr program
                               , jump  :: Mutation program
                               }

-- |Given a prior distribution, score function, mutation distribution, generate
-- a list of (program, score) values through MH sampling.
synthesizeMhList :: RandomGen gen => Problem program -> Rand gen [(program, Double)]
synthesizeMhList Problem {prior, score, jump} = do
    first <- Distr.sample prior
    let density prog = (sc, sc + Distr.logProbability prior prog)
            where sc = score prog
    list <- mhList first density jump
    return [(prog, sc) | (prog, sc, _) <- list]


scanl' :: (a -> b -> a) -> a -> [b] -> [a]
scanl' f q xs = q : (case xs of
                        [] -> []
                        first:rest -> next `seq` scanl f next rest
                            where next = f q first)

-- |Given (value, score) pairs, return a running list of the best pair so far.
runningBest []           = []
runningBest (first:rest) = scanl' maxScore first rest
    where maxScore (p, ps) (q, qs) | qs >= ps = (q, qs)
                                   | otherwise = (p, ps)

-- runningBest :: [(a, Double)] -> [(a, Double)]
-- runningBest [] = []
-- runningBest [only] = [only]
-- runningBest ((p,ps):(q,qs):rest) | qs >= ps = (p, ps) : runningBest ((q,qs):rest)
--                                  | otherwise = (p, ps) : runningBest ((p,ps):rest)