{-# LANGUAGE BangPatterns #-}

-- |
-- Module      : Harmonic.Traversal.Probabilistic
-- Description : Gamma distribution sampling for weighted selection
--
-- This module implements probabilistic selection for the generative engine.
-- The key concept is the "entropy" knob: higher gamma shape parameter pushes
-- selection probability deeper into the sorted candidate list, favoring
-- more "unusual" but still valid harmonic choices.
--
-- == Academic Lineage
--
-- /Data Science In The Creative Process/ (South, 2018): the probabilistic
-- traversal strategy (T component of Wiggins' Creative Systems Framework).
-- The original implementation used R's @rgamma@ via inline-r; this module
-- replaces it with a pure Haskell gamma distribution from @mwc-random@.

module Harmonic.Traversal.Probabilistic
  ( -- * Gamma Sampling
    gammaIndex
  , gammaIndexScaled
  , gammaIndexScaledWith
  , gammaSelect
  , gammaSelectFromPool
  , gammaSequence

    -- * Weighted Selection
  , weightedSelect
  , pickWeighted

    -- * Random Utilities
  , withRandomGen
  ) where

import System.Random (randomRIO, randomIO, StdGen, mkStdGen, Random(..))
import System.Random.MWC (GenIO, createSystemRandom, uniformRM)
import qualified System.Random.MWC.Distributions as Dist
import Statistics.Distribution (quantile)
import Statistics.Distribution.Gamma (gammaDistr)
import Data.List (sortBy)
import Data.Ord (Down(..))
import Control.Monad (replicateM)

-------------------------------------------------------------------------------
-- Gamma Sampling
-------------------------------------------------------------------------------

-- |Draw an index from a gamma distribution.
-- 
-- The shape parameter controls "unusualness":
--   shape ≈ 1.0: High probability of index 0 or 1 (common paths)
--   shape > 2.0: Distribution peak moves away from 0 (unusual paths)
--   shape > 4.0: Very likely to pick 3rd, 4th, or deeper indices
--
-- The index is clamped to [0, maxIndex] to ensure valid selection.
-- Alpha is floored at 0.01 to prevent crash when entropy=0.
--
-- Implementation uses Statistics.Distribution.Gamma for the distribution
-- and mwc-random for high-quality random number generation.
gammaIndex :: Double -> Int -> IO Int
gammaIndex :: Double -> Int -> IO Int
gammaIndex Double
shape Int
maxIndex = do
  Gen RealWorld
rng <- IO (Gen RealWorld)
IO GenIO
createSystemRandom
  let safeShape :: Double
safeShape = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0.01 Double
shape  -- Prevent gamma crash at alpha=0
  Double
x <- Double -> Double -> Gen RealWorld -> IO Double
forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
Dist.gamma Double
safeShape Double
1.0 Gen RealWorld
rng  -- Gamma(shape, scale=1.0)
  let idx :: Int
idx = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x
  Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
idx Int
maxIndex)

-- |Draw an index scaled for entropy in [0, 1] range.
--
-- Maps entropy to gamma shape parameter (shape = 1 + entropy * 9):
--   entropy = 0.0 -> shape =  1.0 (strongly peaked at index 0)
--   entropy = 0.5 -> shape =  5.5 (moderate spread)
--   entropy = 1.0 -> shape = 10.0 (spread toward higher indices)
--
-- The index is clamped to [0, poolSize-1].
-- This is the primary selection function for the redesigned algorithm.
gammaIndexScaled :: Double  -- ^ Entropy in [0, 1] range
                 -> Int     -- ^ Pool size (e.g., 30)
                 -> IO Int
gammaIndexScaled :: Double -> Int -> IO Int
gammaIndexScaled Double
entropy Int
poolSize = do
  Gen RealWorld
rng <- IO (Gen RealWorld)
IO GenIO
createSystemRandom
  let clampedEntropy :: Double
clampedEntropy = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0.0 (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
1.0 Double
entropy)
      shape :: Double
shape = Double
1.0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
clampedEntropy Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
9.0  -- Maps [0,1] -> [1,10]
      safeShape :: Double
safeShape = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0.01 Double
shape
  Double
x <- Double -> Double -> Gen RealWorld -> IO Double
forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
Dist.gamma Double
safeShape Double
1.0 Gen RealWorld
rng
  let idx :: Int
idx = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x
      maxIdx :: Int
maxIdx = Int
poolSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
idx Int
maxIdx)

-- |Like 'gammaIndexScaled' but uses an existing random generator instead of
-- creating a new one. Avoids per-call overhead of reading /dev\/urandom.
gammaIndexScaledWith :: GenIO    -- ^ Shared random generator
                     -> Double   -- ^ Entropy in [0, 1] range
                     -> Int      -- ^ Pool size (e.g., 30)
                     -> IO Int
gammaIndexScaledWith :: GenIO -> Double -> Int -> IO Int
gammaIndexScaledWith GenIO
gen Double
entropy Int
poolSize = do
  let clampedEntropy :: Double
clampedEntropy = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0.0 (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
1.0 Double
entropy)
      shape :: Double
shape = Double
1.0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
clampedEntropy Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
9.0  -- Maps [0,1] -> [1,10]
      safeShape :: Double
safeShape = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0.01 Double
shape
  Double
x <- Double -> Double -> Gen RealWorld -> IO Double
forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
Dist.gamma Double
safeShape Double
1.0 Gen RealWorld
GenIO
gen
  let idx :: Int
idx = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x
      maxIdx :: Int
maxIdx = Int
poolSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
idx Int
maxIdx)

-- |Select an element from a candidate pool using scaled gamma sampling.
-- 
-- The pool is expected to be pre-sorted (highest score first).
-- Entropy controls how "adventurous" the selection is:
--   entropy = 0.0 -> almost always picks top candidate
--   entropy = 1.0 -> frequently picks candidates deep in the list
--
-- If the pool is empty, returns Nothing.
gammaSelectFromPool :: Double -> [(a, Double)] -> IO (Maybe a)
gammaSelectFromPool :: forall a. Double -> [(a, Double)] -> IO (Maybe a)
gammaSelectFromPool Double
_ [] = Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
gammaSelectFromPool Double
entropy [(a, Double)]
pool = do
  Int
idx <- Double -> Int -> IO Int
gammaIndexScaled Double
entropy ([(a, Double)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(a, Double)]
pool)
  Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ (a, Double) -> a
forall a b. (a, b) -> a
fst ([(a, Double)]
pool [(a, Double)] -> Int -> (a, Double)
forall a. HasCallStack => [a] -> Int -> a
!! Int
idx)

-- |Select an element from a weighted list using gamma sampling.
-- 
-- The list is first sorted by weight (highest first), then gamma sampling
-- determines which index to pick. Higher shape = deeper index = more unusual.
--
-- If the list is empty, returns Nothing.
gammaSelect :: Double -> [(a, Double)] -> IO (Maybe a)
gammaSelect :: forall a. Double -> [(a, Double)] -> IO (Maybe a)
gammaSelect Double
_ [] = Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
gammaSelect Double
shape [(a, Double)]
candidates = do
  let sorted :: [(a, Double)]
sorted = ((a, Double) -> (a, Double) -> Ordering)
-> [(a, Double)] -> [(a, Double)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (Down Double -> Down Double -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Down Double -> Down Double -> Ordering)
-> ((a, Double) -> Down Double)
-> (a, Double)
-> (a, Double)
-> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (Double -> Down Double
forall a. a -> Down a
Down (Double -> Down Double)
-> ((a, Double) -> Double) -> (a, Double) -> Down Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Double) -> Double
forall a b. (a, b) -> b
snd)) [(a, Double)]
candidates
      maxIdx :: Int
maxIdx = [(a, Double)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(a, Double)]
sorted Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  Int
idx <- Double -> Int -> IO Int
gammaIndex Double
shape Int
maxIdx
  Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ (a, Double) -> a
forall a b. (a, b) -> a
fst ([(a, Double)]
sorted [(a, Double)] -> Int -> (a, Double)
forall a. HasCallStack => [a] -> Int -> a
!! Int
idx)

-- |Generate a sequence of gamma-distributed indices.
-- 
-- Matches legacy gammaGen behavior: generate n indices for n transitions.
-- Each index can be used to select from a sorted candidate list.
gammaSequence :: Double    -- ^ Gamma shape parameter
              -> Double    -- ^ Entropy multiplier (scales shape)
              -> Int       -- ^ Number of indices to generate
              -> IO [Int]
gammaSequence :: Double -> Double -> Int -> IO [Int]
gammaSequence Double
baseShape Double
entropy Int
count = do
  Gen RealWorld
rng <- IO (Gen RealWorld)
IO GenIO
createSystemRandom
  let effectiveShape :: Double
effectiveShape = Double
baseShape Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
entropy)
  Int -> IO Int -> IO [Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
count (IO Int -> IO [Int]) -> IO Int -> IO [Int]
forall a b. (a -> b) -> a -> b
$ do
    Double
x <- Double -> Double -> Gen RealWorld -> IO Double
forall g (m :: * -> *).
StatefulGen g m =>
Double -> Double -> g -> m Double
Dist.gamma Double
effectiveShape Double
1.0 Gen RealWorld
rng
    Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x

-------------------------------------------------------------------------------
-- Weighted Selection (Non-Gamma)
-------------------------------------------------------------------------------

-- |Select from a weighted list using standard roulette wheel selection.
-- 
-- Each element's selection probability is proportional to its weight.
-- This is used when gamma "unusualness" is not desired.
weightedSelect :: [(a, Double)] -> IO (Maybe a)
weightedSelect :: forall a. [(a, Double)] -> IO (Maybe a)
weightedSelect [] = Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
weightedSelect [(a, Double)]
candidates = do
  let total :: Double
total = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((a, Double) -> Double) -> [(a, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (a, Double) -> Double
forall a b. (a, b) -> b
snd [(a, Double)]
candidates)
  if Double
total Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0
    then Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ (a, Double) -> a
forall a b. (a, b) -> a
fst ([(a, Double)] -> (a, Double)
forall a. HasCallStack => [a] -> a
head [(a, Double)]
candidates)
    else do
      Double
r <- (Double, Double) -> IO Double
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Double
0, Double
total)
      Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ Double -> [(a, Double)] -> a
forall {t} {a}. (Ord t, Num t) => t -> [(a, t)] -> a
pick Double
r [(a, Double)]
candidates
  where
    pick :: t -> [(a, t)] -> a
pick t
_ [(a
x, t
_)] = a
x
    pick t
r ((a
x, t
w):[(a, t)]
rest)
      | t
r t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
w    = a
x
      | Bool
otherwise = t -> [(a, t)] -> a
pick (t
r t -> t -> t
forall a. Num a => a -> a -> a
- t
w) [(a, t)]
rest
    pick t
_ [] = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"weightedSelect: empty list"

-- |Pure version of weighted selection given a random value in [0, total].
-- Useful for deterministic testing.
pickWeighted :: [(a, Double)] -> Double -> a
pickWeighted :: forall a. [(a, Double)] -> Double -> a
pickWeighted [] Double
_ = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"pickWeighted: empty list"
pickWeighted [(a, Double)]
candidates Double
r = Double -> [(a, Double)] -> a
forall {t} {a}. (Ord t, Num t) => t -> [(a, t)] -> a
go Double
r [(a, Double)]
candidates
  where
    go :: t -> [(a, t)] -> a
go t
_ [(a
x, t
_)] = a
x
    go t
remaining ((a
x, t
w):[(a, t)]
rest)
      | t
remaining t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
w = a
x
      | Bool
otherwise      = t -> [(a, t)] -> a
go (t
remaining t -> t -> t
forall a. Num a => a -> a -> a
- t
w) [(a, t)]
rest
    go t
_ [] = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"pickWeighted: exhausted list"

-------------------------------------------------------------------------------
-- Random Utilities
-------------------------------------------------------------------------------

-- |Execute an action with a fresh random generator.
-- Provides a seeded StdGen for reproducible randomness.
withRandomGen :: Int -> (StdGen -> a) -> a
withRandomGen :: forall a. Int -> (StdGen -> a) -> a
withRandomGen Int
seed StdGen -> a
f = StdGen -> a
f (Int -> StdGen
mkStdGen Int
seed)

-- Utility for comparison in sorting
on :: (b -> b -> c) -> (a -> b) -> a -> a -> c
on :: forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
on b -> b -> c
cmp a -> b
f a
x a
y = b -> b -> c
cmp (a -> b
f a
x) (a -> b
f a
y)