Random.hs

module WholePixels.Random where

import Relude
import Control.Monad.Random
import WholePixels.Geometry
import WholePixels.Color
import qualified System.Random.Shuffle

disturbedSequence :: MonadRandom m => [Double] -> Double -> m [Double]
disturbedSequence xs amp = do
    dxs <- replicateM (length xs) $ getRandomR (-amp, amp)
    pure $ zipWith (+) xs dxs

filterRandomly :: MonadRandom m => Double -> [a] -> m [a]
filterRandomly probability xs = do
    ps <- replicateM (length xs) $ getRandomR (0, 1)
    pure [x | (x, p) <- zip xs ps, p < probability]

coinToss :: MonadRandom m => m Bool
coinToss = uniform [False, True]

genGrid :: MonadRandom m => Int -> Int -> m a -> m [(Int, Int, a)]
genGrid colCount rowCount genElement =
    genGridWithBoundaries colCount rowCount (const genElement)

genGridWithBoundaries :: MonadRandom m => Int -> Int -> ([Direction] -> m a) -> m [(Int, Int, a)]
genGridWithBoundaries colCount rowCount genElement = do
    let numberGrid = [(x, y) | y <- [0 .. rowCount - 1], x <- [0 .. colCount - 1]]
    forM numberGrid $ \(x, y) -> do
        let dirs = [ R | x == 0] <> [ D | y == 0] <> [ L | x == colCount - 1] <> [ U | y == rowCount - 1]
        (x, y,) <$> genElement dirs

genGrid' :: MonadRandom m => GridSpec -> ([Direction] -> m a) -> m [(Int, Int, a)]
genGrid' (GridSpec gridSpec) genElement = do
    let rowCount = length gridSpec
        colCount = case gridSpec of
            [] -> 0
            (row : _) -> length row
        numberedGrid :: [(Int, Int, CellSpec)]
        numberedGrid = concat $ zipWith
            (\j row -> map (\(i, c) -> (i, j, c)) row)
            [0..rowCount]
            (map
                (zip [0..colCount])
                gridSpec)
    fmap catMaybes . forM numberedGrid $ \(x, y, c) -> do
        let dirs = [ R | x == 0] <> [ D | y == 0] <> [ L | x == colCount - 1] <> [ U | y == rowCount - 1]
        case c of
            N -> pure Nothing
            Y -> (Just . (x, y,)) <$> genElement dirs

probably :: MonadRandom m => Double -> a -> m (Maybe a)
probably probability thing = do
    x <- getRandomR (0, 1)
    pure $
        if x < probability
        then Just thing
        else Nothing

shuffleM :: MonadRandom m => [a] -> m [a]
shuffleM = System.Random.Shuffle.shuffleM

data PaletteStrategy
    = Analogous
    | Complementary
    | SplitComplementary
    | Triangle

genPalette :: MonadRandom m => m Palette
genPalette = do
    strategy <- uniform [Analogous, Complementary, SplitComplementary, Triangle]
    genPaletteWithStrategy strategy

genPaletteWithStrategy :: MonadRandom m => PaletteStrategy -> m Palette
genPaletteWithStrategy strategy = do
    baseHue <- getRandomR (0, 360)
    baseSaturation <- getRandomR (0, 1)
    baseValue <- getRandomR (0, 0.5)
    bgToFgHueDifference <- uniform [180, 0, 30, 120, 240]
    let baseColor = HSV baseHue baseSaturation baseValue
        c = HSV (fixHue $ baseHue + bgToFgHueDifference) 1 1
        colors = case strategy of
            Analogous -> take 4 $ analogous c
            Complementary -> take 3 (analogous c) <> [complementary c]
            SplitComplementary ->
                let (c1, c2) = splitComplementary c
                in [c, c1, c2]
            Triangle ->
                let (c1, c2) = splitTriangle c
                in [c, c1, c2]
    pure $ Palette {..}

genMonochromePaletteForColor :: MonadRandom m => HSV -> m Palette
genMonochromePaletteForColor baseColor@(HSV bh bs bv) = do
    let colors = [HSV bh bs v | v <- [bv + 0.15, bv + 0.2 .. 1]]
    pure $ Palette {..}

genMonochromePalette :: MonadRandom m => Double -> m Palette
genMonochromePalette maxSaturation = do
    hue <- getRandomR (0, 360)
    saturation <- getRandomR (0.0, maxSaturation)
    baseValue <- getRandomR (0, 0.5)
    let baseColor = HSV hue saturation baseValue
        colors = [HSV hue saturation v | v <- [baseValue + 0.15, baseValue + 0.2 .. 1]]
    pure $ Palette {..}

genColor' :: MonadRandom m => Palette -> m HSV
genColor' pal = genColor pal 0.5 1.0

genColor :: MonadRandom m => Palette -> Double -> Double -> m HSV
genColor Palette {..} expected sigma2 = do
    let colorCount = length (baseColor : take 10 colors)
        colorPositions = take colorCount [0.0, (1.0 / fromIntegral colorCount) ..]
        weights = map
            (\p -> toRational $ 1000.0 * exp (- (p - expected) * (p - expected) / sigma2))
            colorPositions
    weighted $ zip (baseColor : colors) weights

genRectSubdivision :: forall m. MonadRandom m => Int -> Rect -> m [Rect]
genRectSubdivision depth r = foldr (>=>) pure (replicate depth step) [r]
    where
    step :: [Rect] -> m [Rect]
    step rs = concat <$> mapM go rs
    go :: Rect -> m [Rect]
    go r' = do
        let hbonusWeight = if rh r' > rw r' then 15 else 0
            vbonusWeight = if rw r' > rh r' then 15 else 0
        f <- weighted
            [ (pure . pure, 3)
            , (pure . hsplit, 1 + hbonusWeight)
            , (pure . hsplitGolden, 3 + hbonusWeight)
            , (pure . hsplitReverseGolden, 3 + hbonusWeight)
            , (pure . vsplit, 1 + vbonusWeight)
            , (pure . vsplitGolden, 3 + vbonusWeight)
            , (pure . vsplitReverseGolden, 3 + vbonusWeight)
            ]
        f r'