-- Monada za verjetnostno računanje.
--
-- Izračun predstavlja končno verjetnostno porazdelitev nad poljubnim tipom.
-- Porazdelitev elemtov tipa a predstavimo s seznamom parov (izid, verjetnost):
-- seznam [(x_1, p_1), ..., (x_n, p_n)] pomeni, da izid x_i nastopi z
-- verjetnostjo p_i. Vsota verjetnosti je največ 1; če je strogo manjša,
-- manjkajoča verjetnost ustreza neuspelim ali zavrženim izračunom.

module Verjetnost where

-- Iz standardne knjižnice uvozimo nekaj funkcij za delo s seznami,
-- ki jih bomo potrebovali v funkciji `zdruzi`.
import Data.List (sortBy, groupBy)
import Data.Ord (comparing)
import Data.Function (on)

-- Definicija tipa za verjetnostno porazdelitev.
--
-- Ključna beseda `newtype` v Haskellu definira nov tip, ki ima natanko
-- en konstruktor in eno polje, in pri tem ne uvede dodatne ovojnice v
-- pomnilniku. Tip `Verjetnost a` je v bistvu `[(a, Float)]`,
-- vendar ga Haskell obravnava kot ločen tip. Tako ga lahko opremimo z
-- lastnimi instancami razredov tipov (Functor, Applicative, Monad).
--
-- Zapis `{ izidi :: [(a, Float)] }` poleg konstruktorja `Verjetnost`
-- definira še konstruktor
--     izidi :: Verjetnost a -> [(a, Float)]
-- s katerim iz vrednosti tipa `Verjetnost a` dobimo nazaj seznam parov.
newtype Verjetnost a = Verjetnost { izidi :: [(a, Float)] }

-- Instanca razreda `Show` določa, kako naj se vrednost tipa
-- `Verjetnost a` izpiše v GHCi. Vsak izid postavimo v svojo vrstico,
-- ločeno s tabulatorjem od njegove verjetnosti.
instance Show a => Show (Verjetnost a) where
    show (Verjetnost xs) =
        concat [show x ++ ":\t" ++ show p ++ "\n" | (x, p) <- xs]

-- Da bi tip `Verjetnost` postal monada v smislu Haskella, moramo podati
-- tri instance: `Functor`, `Applicative` in `Monad`. Razred `Monad`
-- namreč podeduje od `Applicative`, ta pa od `Functor`.

-- Functor: na vsakem izidu uporabimo čisto funkcijo `f`, verjetnost pa
-- ostane nespremenjena. Tako npr. `fmap (+10) kocka` pretvori izide
-- 1..6 v 11..16, vsak še vedno z verjetnostjo 1/6.
instance Functor Verjetnost where
    fmap f (Verjetnost xs) = Verjetnost [(f x, p) | (x, p) <- xs]

-- Applicative:
--   * `pure x` je porazdelitev, ki z verjetnostjo 1 vrne izid `x`.
--     To je čisti izračun, ki vedno vrne isto vrednost.
--   * `<*>` izvede dva neodvisna izračuna in zmnoži njuni verjetnosti,
--     saj je verjetnost neodvisnih dogodkov produkt njihovih verjetnosti.
instance Applicative Verjetnost where
    pure x = Verjetnost [(x, 1)]
    (Verjetnost fs) <*> (Verjetnost xs) =
        Verjetnost [(f x, p * q) | (f, p) <- fs, (x, q) <- xs]

-- Monad: ključna operacija `>>=` (beri "veži") sestavi dva izračuna
-- v zaporedje. Za vsak izid `x` z verjetnostjo `p` v prvem izračunu
-- poženemo nadaljnji izračun `f x`, ki vrne svojo porazdelitev. Vsako
-- verjetnost v tej drugi porazdelitvi pomnožimo s `p`. Tako dobimo
-- skupno verjetnost zaporedja dveh dogodkov.
instance Monad Verjetnost where
    return = pure
    (Verjetnost xs) >>= f =
        Verjetnost [(y, p * q) | (x, p) <- xs, (y, q) <- izidi (f x)]

-- Pomožna funkcija `zdruzi`: sešteje verjetnosti enakih izidov in
-- vrne urejen seznam različnih izidov.
--
-- OPOMBA: matematično bi to združevanje sodilo neposredno v operacije
-- `>>=` in `<*>`, saj je sumiranje verjetnosti enakih izidov del same
-- definicije sestavljanja porazdelitev. V Haskellu pa tega ne moremo
-- doseči, ker razreda `Monad` in `Applicative` ne dovolita dodatnih
-- omejitev na tipovne parametre: `>>=` mora delovati za poljuben tip
-- izidov `a`, ne le za tiste z `Eq a` ali `Ord a`. Brez primerjanja
-- pa ne moremo prepoznati enakih izidov. Združevanje zato naknadno
-- opravimo z `zdruzi`, ki potrebuje omejitev `Ord a`.
zdruzi :: Ord a => Verjetnost a -> Verjetnost a
zdruzi (Verjetnost xs) =
    Verjetnost [ (k, sum (map snd gs))
               | gs@((k, _) : _) <- groupBy ((==) `on` fst) (sortBy (comparing fst) xs) ]

-- Pomožna funkcija `pogoj`: zavrže veje izračuna, ki ne ustrezajo pogoju.
--   * Če pogoj velja, vrnemo enotni izid `()` z verjetnostjo 1;
--     `>>=` nato z verjetnostjo 1 nadaljuje s preostankom izračuna.
--   * Če pogoj ne velja, vrnemo prazno porazdelitev. Ker `>>=` množi
--     verjetnosti, se s tem tista veja izračuna "ubije" in njena
--     verjetnost je izgubljena.
-- Uporabljamo jo v `do` blokih za filtriranje izidov.
preveri :: Bool -> Verjetnost ()
preveri True  = Verjetnost [((), 1)]
preveri False = Verjetnost []

-- Primer monade. Met poštene šestostrane kocke: vsak izmed izidov
-- 1, 2, ..., 6 nastopi z verjetnostjo 1/6.
kocka :: Verjetnost Int
kocka = Verjetnost [(n, 1/6) | n <- [1..6]]

-- Primer 1: porazdelitev vsote pri metu dveh kock.
--
-- V `do` zapisu vrstica `x <- kocka` pomeni "vrzi kocko in njen izid
-- poimenuj `x`"; podobno za `y`. Ker je `>>=` desugarizacija `do`,
-- se v ozadju za vsako izbiro `x` razveji vseh šest možnih izbir za
-- `y`. Skupaj dobimo 36 enako verjetnih parov (x, y), vsak z
-- verjetnostjo (1/6)*(1/6) = 1/36. Z `return (x + y)` vsak par
-- pretvorimo v njegovo vsoto.
--
-- Naključne vsote se ponovijo: 7 dobimo kot 1+6, 2+5, 3+4, 4+3, 5+2,
-- 6+1, torej z verjetnostjo 6/36 = 1/6. Brez `zdruzi` bi izpis imel
-- 36 vrstic, s klicem `zdruzi` pa 11 vrstic za različne vsote 2..12.
vsotaDveh :: Verjetnost Int
vsotaDveh = zdruzi $ do
    x <- kocka
    y <- kocka
    return (x + y)

-- Primer 2: porazdelitev vsote dveh metov, pri čemer met zavržemo,
-- če je vsota soda.
--
-- Polovica vseh 36 parov ima sodo vsoto, zato vsota verjetnosti
-- preostalih izidov ni 1, temveč 1/2. Dobimo torej "pod-porazdelitev",
-- v kateri manjkajoča verjetnost ustreza zavrženim metom.
lihaVsotaDveh :: Verjetnost Int
lihaVsotaDveh = zdruzi $ do
    x <- kocka
    y <- kocka
    preveri (odd (x + y))
    return (x + y)