""" A toy CT-MC example.
  We will sample the expression

   Z = e^{\int_0^\beta w(\tau) d\tau}

   where w(\tau) = alpha * (exp(-\tau) + exp(-\beta+\tau))

   
"""
from scipy import *
from pylab import *
import sys

alpha=10.0  # prefactor in w(t)
b=10.       # beta

def w(t):
    "The function to sample"
    return alpha*(exp(-t)+exp(-b+t))
def Iwr(t1,t2):
    "The integral of w. Iwr(t1,t2) = int_{t_1}^{t_2} w(t)dt "
    return alpha*(exp(-t1)-exp(-t2)+exp(-b+t2)-exp(-b+t1))

dh=0.25

def add_a_kink(tms, wals, Nk):
    "Adds a kink at random time between 0,beta."
    tnew = rand()*b        # new time [0,beta] 
    wnew = w(tnew)         # w(t)
    P = b*wnew/(Nk+1)      # Probability to add
    if P>rand():
        tms[Nk] = tnew     # tms   -- all times
        wals[Nk] = wnew    # wals  -- values of w in the kinks
        Nk += 1            # current perturbation order
    return (tms, wals, Nk)
    
def remove_a_kink(tms, wals, Nk):
    "Removes one of the existing kinks."
    which = int(Nk*rand())     # one of the kinks, selected at random
    P = Nk/(b*wals[which])     # probability to remove a kink
    if P>rand():
        for i in range(which,Nk-1): # all other kinks move down to fill the hole.
            tms[i] = tms[i+1]
            wals[i] = wals[i+1]
        Nk -= 1               # a kink less
    return (tms, wals, Nk)

def Bind(tms, Nk, Ptau):
    """Saves all times at this stage in sampling.
       We will have probability for each time.
    """
    for k in range(Nk):
        Ptau[int(tms[k]/dh)] += 1
    return Ptau

if __name__ == '__main__':

    M = 10000      # all MC steps
    warmup = 1000  # warmup steps
    
    Nmax = int(10*alpha)          # maximum number of allowed times.
    tms = zeros(Nmax,dtype=float) # all times stored in this array
    wals= zeros(Nmax,dtype=float) # all values of w(t) stored in this array
    Nk=0                          # number of kinks

    Nord = zeros(Nmax, dtype=float)      # The histogram of kinks
    Ptau = zeros(int(b/dh), dtype=float) # Probability for a time

    
    for i in range(M):
        if rand()>0.5:
            (tms, wals, Nk) = add_a_kink(tms, wals, Nk)
        else:
            (tms, wals, Nk) = remove_a_kink(tms, wals, Nk)

        if i>warmup:
            # sampling
            Nord[Nk] += 1              # Probability for this order
            Ptau = Bind(tms, Nk, Ptau) # Probability for this tau
        
    
    Nord /= sum(Nord) # probability for order k
    Ptau /= sum(Ptau) # probability for time t

    # Exact probability for order k
    Nord0 = zeros(Nmax, dtype=float)
    Nord0[0] = 1
    Iw = Iwr(0,b)
    for k in range(1,len(Nord0)):
        Nord0[k] = Nord0[k-1]*Iw/k
    Nord0 /= sum(Nord0)


    # Plotting probability for order k
    Nr = int(5*alpha)
    fig1 = figure()
    plot(range(Nr), Nord[:Nr])
    plot(range(Nr), Nord0[:Nr])

    
    
    Ptau0 = zeros(len(Ptau),dtype=float)
    for i in range(len(Ptau)):
        print i*dh, (i+1)*dh, Iwr(i*dh,(i+1)*dh)
        Ptau0[i] = Iwr(i*dh,(i+1)*dh)
    Ptau0 /= sum(Ptau0)
    
    fig2 = figure()
    plot((arange(len(Ptau))+0.5)*dh, Ptau)
    plot((arange(len(Ptau))+0.5)*dh, Ptau0)
    show()