-----------------------------------------------------------------------------
-- |
-- Module      :  Data.SBV.Tools.ExpectedValue
-- Copyright   :  (c) Levent Erkok
-- License     :  BSD3
-- Maintainer  :  erkokl@gmail.com
-- Stability   :  experimental
-- Portability :  portable
--
-- Computing the expected value of a symbolic variable
-----------------------------------------------------------------------------

{-# LANGUAGE PatternGuards #-}
module Data.SBV.Tools.ExpectedValue (expectedValue, expectedValueWith) where

import Control.DeepSeq (rnf)
import System.Random   (newStdGen, StdGen)
import Numeric

import Data.SBV.BitVectors.Data

-- | Generalized version of 'expectedValue', allowing the user to specify the
-- warm-up count and the convergence factor. Maximum iteration count can also
-- be specified, at which point convergence won't be sought. The boolean controls verbosity.
expectedValueWith :: Outputtable a => Bool -> Int -> Maybe Int -> Double -> Symbolic a -> IO [Double]
expectedValueWith verbose warmupCount mbMaxIter epsilon m
  | warmupCount < 0 || epsilon < 0
  = error $ "SBV.expectedValue: warmup count and epsilon both must be non-negative, received: " ++ show (warmupCount, epsilon)
  | True
  = warmup warmupCount (repeat 0) >>= go warmupCount
  where progress s | not verbose = return ()
                   | True        = putStr $ "\r*** " ++ s
        warmup :: Int -> [Integer] -> IO [Integer]
        warmup 0 v = do progress $ "Warmup complete, performed " ++ show warmupCount ++ " rounds.\n"
                        return v
        warmup n v = do progress $ "Performing warmup, round: " ++ show (warmupCount - n)
                        g <- newStdGen
                        t <- runOnce g
                        let v' = zipWith (+) v t
                        rnf v' `seq` warmup (n-1) v'
        runOnce :: StdGen -> IO [Integer]
        runOnce g = do (_, Result _ _ _ _ cs _ _ _ _ _ cstrs os) <- runSymbolic' (Concrete g) (m >>= output)
                       let cval o = case o `lookup` cs of
                                      Nothing -> error "SBV.expectedValue: Cannot compute expected-values in the presence of uninterpreted constants!"
                                      Just cw -> case (cwSigned cw, cwSize cw) of
                                                   (True,  Size Nothing ) -> error "Cannot compute expected-values for unbounded integer results."
                                                   (False, Size (Just 1)) -> if cwToBool cw then 1 else 0
                                                   _                      -> cwVal cw
                       if all ((== 1) . cval) cstrs
                          then return $ map cval os
                          else runOnce g -- constraint not satisfied try again with the same set of constraints
        go :: Int -> [Integer] -> IO [Double]
        go cases curSums
         | Just n <- mbMaxIter, n < curRound
         = do progress "\n"
              progress "Maximum iteration count reached, stopping.\n"
              return curEVs
         | True
         = do g <- newStdGen
              t <- runOnce g
              let newSums  = zipWith (+) curSums t
                  newEVs = map ev' newSums
                  diffs  = zipWith (\x y -> abs (x - y)) newEVs curEVs
              if all (< epsilon) diffs
                 then do progress $ "Converges with epsilon " ++ show epsilon ++ " after " ++ show curRound ++ " rounds.\n"
                         return newEVs
                 else do progress $ "Tuning, round: " ++ show curRound ++ " (margin: " ++ showFFloat (Just 6) (maximum (0:diffs)) "" ++ ")"
                         go newCases newSums
         where curRound = cases - warmupCount
               newCases = cases + 1
               ev, ev' :: Integer -> Double
               ev  x  = fromIntegral x / fromIntegral cases
               ev' x  = fromIntegral x / fromIntegral newCases
               curEVs = map ev curSums

-- | Given a symbolic computation that produces a value, compute the
-- expected value that value would take if this computation is run
-- with its free variables drawn from uniform distributions of its
-- respective values, satisfying the given constraints specified by
-- 'constrain' and 'pConstrain' calls. This is equivalent to calling
-- 'expectedValueWith' the following parameters: verbose, warm-up
-- round count of @10000@, no maximum iteration count, and with
-- convergence margin @0.0001@.
expectedValue :: Outputtable a => Symbolic a -> IO [Double]
expectedValue = expectedValueWith True 10000 Nothing 0.0001