### 
An implementation of the Dirichlet process (DP) using memoization. For an explanation of DPs, see 

Teh et. al. Hierarchical Dirichlet processes. Journal of the Am. Stat. Assoc., 101(476):1566--1581, 2006. 
###

# placeholder for a function that would generate a draw from the beta distribution (so this is a draw from the Beta(1,1) distribution, no matter what a and b are
fun float:beta float:a float:b = ~rand in 

# get a stick, breaking more if necessary 
fun int:pickastick (fun float int):sticks int:j =
    if ~rand < ~sticks j then j else ~pickastick sticks j+1
in

# generic Dirichlet process code
fun (fun int):DP float:alpha (fun int):proc =
    fun float:sticks int:x := ~beta 1.0 alpha in
    fun int:atoms  int:x := ~proc in
    fun int:f = ~atoms ~pickastick sticks 1 in
    f # return f
in

fun (fun (fun int) float):DPmem float:alpha (fun int float):proc =
    fun (fun int):dps float:arg := 
    	fun int:apply = ~proc arg in
    	~DP alpha apply 
    in
    fun (fun int):dp float:arg = ~dps arg in
    dp
in

# this function will create Dirichlet process draws with geometric base distribution
let (fun (fun int) float):geom_dp = ~DPmem 1.0 geom in
 
# this is a DP draw with geometric base distribution with q = .2
let (fun int):mydraw = ~geom_dp .2 in

# use a tail-recursive loop to generate some samples from the Dirichlet Process
fun bool:loop int:i =
    ~print ~mydraw;
    if i > 0 then ~loop i - 1 else true
in
~seed;
~loop 30; ~print_line ~mydraw
    
