Distribution.swift

struct Distribution<A> {
    var get: () -> A?
    
    func sample(n: Int) -> [A] {
        return (1...n).map { x in get()! }
    }
    
    func map<B>(f: A -> B) -> Distribution<B> {
        var d = Distribution<B>(get: {() -> Optional<B> in return nil})
        d.get = {
            (Void) -> B in return f(self.get()!)
        }
        return d
    }
    
    // FlatMaps one Distribution into another
    func flatMap<B>(f: A -> Distribution<B>) -> Distribution<B> {
        var d = Distribution<B>(get: {() -> Optional<B> in return nil})
        d.get = {
            (Void) -> B in return f(self.get()!).get()!
        }
        return d
    }
    
    let N = 10000
    
    // probability of the predicate being true
    func prob(predicate: A -> Bool) -> Double {
        return Double(sample(N).filter(predicate).count) / Double(N)
    }
    
    // TODO: This function doesn't work just yet. It's returning the same number. Needs tail recursion too. Perhaps this is not possible in Swift?
    
    // Samples from the new distribution so that the result matches the predicate
    //    func given(predicate: A -> Bool) -> Distribution<A> {
    //        var d: Distribution<A> = self
    //        let a = d.get()!
    //        d.get = { (Void) -> A in return predicate(a) ? a : d.get()! }
    //        return d
    //    }
    
    func mean() -> Double {
        return sample(N).reduce(0, combine: { $0 + Double(String($1))! }) / Double(N)
    }
    
    func variance() -> Double {
        var sum: Double = 0
        var sqrSum: Double = 0
        
        for x in sample(N) {
            let xx = Double(String(x))!
            sum += xx
            sqrSum += xx * xx
        }
        
        return (sqrSum - sum * sum / Double(N)) / Double(N-1)
    }
    
    func stdDev() -> Double {
        return sqrt(self.variance())
    }
}