# tol-mc.py
# Monte-Carlo Simulation of a basic muon time-of-flight experiment
# @author: bryan newbold
# @email: bnewbold@mit.edu
# @date: oct 2006

from pylab import *;
from scipy import *;

# default number of tries to run...
max_tries = 10**2

# default geometry
h = 100.0 # cm
h_err = 1.0 # cm

Tw = 15.5*2.54 # inches->cm
Tw_err = .5*2.54
Tl = 23.5*2.54
Tl_err = 1.*2.54
Tpmtl = 19.*2.54
Tpmtl_err = 1.*2.54

Bw = 42.
Bw_err = 4./10 # mm->cm
Bl = 53.
Bl_err = 1.
Bpmtl = 20.
Bpmtl_err = 1.

n_scint = 1.5 # index of refraction of scintillators
C_ = 2.99792458*10**10 # speed of light in cm

# shortcuts to random number generators
ur = lambda: rand();
gr = lambda: randn();

aTx = lambda: ur()*Tl
aTy = lambda: ur()*Tw
#aphi = lambda: arccos( sqrt(ur()))
aomega = lambda: ur()*2*pi
aBx = lambda x,p,o,h: x+h*sin(o)*tan(p)
aBy = lambda y,p,o,h: y+h*cos(o)*tan(p)
#aP = lambda p: 100.*cos(p)
aP = lambda p,mc: ((mc/2.99) + gr() * 0.02) * C_
aTime = lambda x,y,pmt,w: n_scint/(C_)*sqrt((x+pmt)**2 + (w/2 - y)**2) 

if not vars().has_key('phi_table'):
    phi_table = zeros((1000,1),typecode='f')
    for i in frange(0., pi/2,0.00001):
        m = (2*i + sin(2*i))/pi
        phi_table[int(1000*(m - (m%.001)))] = i
    phi_table[-1] = pi/2

def aphi():
    n = ur();
    return phi_table[int(1000*(n-(n%.001)))][0];

def run(tries=max_tries, h=h, h_err=h_err,printstuff=True,mean_c=2.93):
    num_events=0;
    if(printstuff): print "Going to make %d tries..." % tries
    events=zeros((tries,7),typecode='f')
    while tries>0:
        tries = tries-1;
        Tx,Ty,phi,omega=aTx(),aTy(),aphi(),aomega();
        Bx = aBx(Tx,phi,omega,h)
        if((Bx < 0.) or (Bx > Bl)):
            continue;
        By = aBy(Ty + abs(Tw-Bw)/2,phi,omega,h)
        if((By < 0.) or (By > Bw)):
            continue;
        P = aP(phi,mean_c)
        dL = (h/cos(phi)); 
        dT = -1 *aTime(Tx,Ty + abs(Tw-Bw)/2,Tpmtl,Tw) \
             + aTime(Bx,By,Bpmtl,Bw) \
             + dL/P; 
        events[num_events] = (Tx,Ty,Bx,By,P,dT,dL);
        num_events = num_events +1;
    if num_events>0:
        if(printstuff): print "... %d valid events!" % num_events;
        return events[0:num_events];
    else:
        if(printstuff): print "No valid events!"
        return [];

def plot_events(events, xres=16,yres=16):
    if(size(events[0]) != 7):
        print "Passed matrix is no good...";
        return;
    figure();
    # TOP part
    subplot(211);
    x = arange(0.,Tl,Tl/xres);
    y = arange(0.,Tw,Tw/yres);
    X,Y = meshgrid(x,y);
    Z = zeros((xres,yres));
    for i in events:
        Z[int(i[1] * yres/Tw),int(i[0] * xres/Tl)] +=1;
    contourf(X,Y,Z);
    title("Distribution on Top Paddle (counts per %.3g cm^2)" \
        % (Tl/xres * Tw/yres));
    xlabel("Paddle Length (%g cm +/- %g cm)" % (Tl,Tl_err));
    ylabel("Paddle Width (%g cm +/- %g cm)" % (Tw,Tw_err));
    colorbar();

    # BOTTOM part
    subplot(212);
    x = arange(0.,Bl,Bl/xres);
    y = arange(0.,Bw,Bw/yres);
    X,Y = meshgrid(x,y);
    Z = zeros((xres,yres));
    for i in events:
        Z[int(i[3] * yres/Bw),int(i[2] * xres/Bl)] +=1;
    contourf(X,Y,Z);
    title("Distribution on Bottom Paddle (counts per %.3g cm^2)" \
        % (Bl/xres * Tw/yres));
    xlabel("Paddle Length (%g cm +/- %g)" % (Bl,Bl_err));
    ylabel("Paddle Width (%g cm +/- %g cm)" % (Bw,Bw_err));
    colorbar();

def plot_events_3d(events):
    if(size(events[0] != 6)):
        print "Passed matrix is no good...";
        return;
    figure();
    return;

def plot_deltaT_h(heights=frange(33.,233,50.),num_events=1000,err=3,c=2.93):
    
    points = zeros((size(heights),6),typecode='f');
    points[:,0] = list(heights[:]);
    times = frange((int(err)));
    dL = frange((int(err)));
    actP = frange((int(err)));
    for i in points:
        for r in range(0,size(times)):
            events = run(int(num_events),h=i[0],printstuff=False,mean_c=c);
            times[int(r)] = mean(events[:,5]);
            dL[int(r)] = mean(events[:,6]);
            actP[int(r)] = mean(events[:,4]);
        i[1] = mean(times);
        i[2] = std(times);
        i[3] = mean(dL);
        i[4] = mean(actP)
        i[5] = std(actP)
    
    stuff = fittodata_linear(points[:,0],points[:,1],[1.,1.],err=points[:,2],ret_all=True)
    plotfit(stuff)
    v = stuff['params']
    ax = gca()
    title("Change in delta-T with average h");
    xlabel("Average h in cm");
    ylabel("Average delta-T");
    txt = text(0.65,.22,"delta-T/h: %g \noffset: %g \nmean muon speed: %g cm/s\nmonte carlo input speed:%g +/- %g cm/s" 
            % (v[0],v[1],1/v[0],mean(points[:,4]),mean(points[:,5])), 
            horizontalalignment='center',
            verticalalignment='center',
            transform = ax.transAxes)
    

    stuff = fittodata_linear(points[:,3],points[:,1],[1.,1.],err=points[:,2],ret_all=True)
    plotfit(stuff)
    vc = stuff['params']
    ax=gca()
    title('Change in mean delta-T with change in mean L')
    xlabel('Mean Length of flight')
    ylabel('Mean time gap between events')
    txt = text(0.65,.22,"delta-T/h: %g \noffset: %g \nmean muon speed: %g cm/s\nmean muon speed corrected: %g cm/s\nmonte carlo input speed:%g +/- %g cm/s" 
            % (v[0],v[1],1/v[0],1/vc[0],mean(points[:,4]),mean(points[:,5])), 
            horizontalalignment='center',
            verticalalignment='center',
            transform = ax.transAxes)

def plot_L_vs_h(heights=r_[18.5:218.5:8.],num_events=400,err=2,c=2.93):
    points = zeros((size(heights),5),typecode='f');
    
    for i in range(0,size(heights)):
        points[i,0] = heights[i];

    times = r_[0.:int(err)];
    dL = r_[0.:int(err)];
    for i in range(0,size(heights)):
        for r in range(0,size(times)):
            events = run(int(num_events),h=points[i,0],printstuff=False,mean_c=c);
            times[int(r)] = mean(events[:,5]);
            dL[int(r)] = mean(events[:,6]);
        points[i,1] = mean(times);
        points[i,2] = std(times);
        points[i,3] = mean(dL);
        points[i,4] = std(dL);

    x = points[:,0].tolist(); x = array(x)
    y = points[:,3].tolist(); y = array(y)
    err = points[:,4].tolist(); err = array(err)

    plotfit(fittodata_linear(x,y,[3.,-1.],err=err,ret_all=True))
    
    title("Change in mean L with h");
    xlabel("h in cm");
    ylabel("Mean L");
    figure();
    plot(x,y-x);
    return (x,y,err);