# MIT 1805 Studio 7 sensorsim(S, t0,1,iter)

sensorsimHelp = function()
{
    cat("# --------------------------------------------------\n")
    cat("sensorsim()\n")
    cat("Simulate an array of sensors for locating a\n")
    cat("signal from a point t0 in the plane.\n")
    cat("\n")
    cat("Syntax: sensorsim(S, t0, uniformPrior, niterations, nsteps, title)\n")
    cat("\n")
    cat("S is a  n x 2 array representing n points in the plane: s_1 ... s_n.\n")
    cat("t0 is a 1 x 2 array giving the point where the signal originates\n")
    cat("uniformPrior is a boolean:\n")
    cat("       0 or FALSE = use normal prior\n")
    cat("       1 or TRUE  = use uniform prior\n")
    cat("niter = number of iterations through signaling and updating per step\n")
    cat("nsteps = number of steps\n")
    cat("  niter = 3 and nsteps = 5: 15 updates, stop every 3 updates and\n")
    cat("  wait until the user presses return\n")
    cat("title = title for plots\n")
    cat("\n")
    cat("Outputs graphs of the posterior density for t0.")
    cat("\n")
    cat("Details:\n")
    cat("If the sensors are at the points s_1, ..., s_n\n")
    cat("then we set the probability that s_i detects the\n")
    cat("signal as p_i = exp(-d_i) where d_i = distance(s_i,t0)\n")
    cat("\n")
    cat("The data due to a signal is X1,...,Xn\n")
    cat("each is a Bernoulli R.V. with prob. p_i.\n")
    cat("Using a prior f(t) we do a Bayesian update\n")
    cat("to find the posterior f(t|X).\n")
    cat("Given t0 we simulate each of the Bernoulli R.V. Xi\n")
    cat("# --------------------------------------------------\n")
}

sensorsim = function(S,t0,useUniformPrior,niter,nsteps,titleMain="")
{

    #Set up plot parameters
    par(mfcol=c(2,1))
    par(mar=c(2,2,2,2)+.1)
    
    skipPosterior = 0;
    if ((niter==0)| (nsteps==0))
    {
        skipPosterior = 1;
    }

    #trueDetectProb = COLUMN of true probabilites sensors will detect 
    trueDetectProb = sensorprob(S,t0[1], t0[2]);
    nsensors = length(trueDetectProb)
    
    # Find plot limits so they contain all the sensors, t0 and the origin:
    # The xlimits are minimum and maximum x values in S, t0 and (0,0)
    # Likewise for the ylimits
    #add or subtract 1 to make sure dots are not right on the plot edge
    xmin = min(S[,1]-1, t0[1]-1, 0) 
    xmax = max(S[,1]+1, t0[1]+1, 0)
    ymin = min(S[,2]-1, t0[2]-1, 0)
    ymax = max(S[,2]+1, t0[2]+1, 0)

    #Our goal is to cover [xmin,xmax] x [ymin,ymax] with a grid of points
    nx = 41 #number of grid points in x direction
    ny = 41 #number of grid points in y direction
    dx = (xmax-xmin)/nx  #step between points in x-direction
    dy = (ymax-ymin)/ny  #step between points in y-direction
    # set up  arrays to hold positions
    xgrd = rep(0,nx)  #one dim. list of x values of grid positions 
    ygrd = rep(0,ny)  #one dim. list of y values of grid positions

    # Need an nx by ny arrays holding value of the probability
    # the true position of the signal sourece is at (xgrd[i], ygrd[j])
    # currProb will be start as the prior and be updated to the posterior.
    currProb = matrix(rep(0,nx*ny), ncol=nx, nrow=ny)

    #give values to arrays, set currProb to the prior
    for (i in c(1:nx))
    {
        x = xmin + (i-1)*dx; #subtract 1 so (x,y) starts at (xmin,ymin)
        xgrd[i] = x
        for (j in c(1:ny))
        {
            y = ymin + (j-1)*dy; #subtract 1 so (x,y) starts at (xmin,ymin)
            ygrd[j] = y

            priorprob = prior(x,y, useUniformPrior);
            currProb[i,j] = priorprob;
        }
    }
    #The discretization will give an unnormalized prob. function
    currProb.sum = sum(as.vector(currProb))
    currProb = currProb/currProb.sum
    if (skipPosterior)
    {
        X = rep(0,nsensors)
        plotresults(S,xgrd,ygrd,currProb,X,t0, titleMain)
    }
    else
    {
        for (ind in c(1:nsteps))
        {
            # Generate random data: 0,1 for detection at each Xi            
            X = rbinom(nsensors, niter, trueDetectProb);
            for (i in c(1:nx))
            {
                for (j in c(1:ny))
                {
                    x = xgrd[i]
                    y = ygrd[j];
                    priorprob = currProb[i,j];
                    lklhd = sensorprob(S,x,y);
                    currProb[i,j] = priorprob*prod((lklhd^X)*((1-lklhd)^(niter-X)));
                }
            }
            currProb.sum = sum(as.vector(currProb))
            currProb = currProb/currProb.sum  
            plotresults(S,xgrd,ygrd,currProb,X,t0,titleMain)
            if (ind != nsteps)
            {
                s = readline(prompt="Hit return for next step, q then return to quit: \n")
                if (s=='q')
                {
                    break
                }
            }
        }
    }
}

plotresults = function(S,xgrd,ygrd,currProb,X,t0,titleMain=titleMain)
{
    #Make two figures. One a heat map in color and one a contour map
    #showing the locations of the sensors and how many times the sensor
    #detected a signal.
    image(xgrd,ygrd,currProb,xlab='',ylab='') #heat map
    points(S[,1],S[,2], pch="o")
    markHits(S,X,'blue')
    title(main=titleMain)
    points(t0[1],t0[2],pch=19, cex=1.3, col='green')

    contour(xgrd,ygrd,currProb)
    points(S[,1],S[,2], pch="o")
    markHits(S,X,'blue')
    points(t0[1],t0[2],pch=19, cex=1.3, col='green')
    #points(t0[1],t0[2],pch='T',col='green')
}

markHits = function(S, X, hitcol)
{
    for (k in c(1:length(X)))
    {
        loc = S[k,]
        #Put X at sensors that are hit. Size by number of hits
        if (X[k] > 0)
        {
            points(loc[1],loc[2], cex=.6 + X[k]/3,pch='X',col=hitcol )
        }
    }
}

prior = function(x,y,useUniformPrior)
{
    if (useUniformPrior)
    {
        p = unifprior(x,y)
    }
    else
    {
        p = normprior(x,y)
    }
    return(p)
}


unifprior = function(x,y)
{
    #NOT NORMALIZED: since we normalize the posterior
    p=1;
    return(p)
}

# Fix normal prior to be bivariate normal centered at (5,5) with variances 1, correlation 0
normprior = function(x,y)
{
    mux=5
    muy=5
    sigx=2
    sigy=2
    p=exp( - ((x-mux)^2)/(2*sigx^2) -((y-muy)^2)/(2*sigy^2))/(2*pi*sigx*sigy)
    return(p)
}

# Change 100 to 2 below to make sensors weaker
sensorprob = function(S,x,y)
{
    #S is a list of positions (each row = (x_i,y_i)
    #returns a list probabilities
    r = sqrt((x-S[,1])^2 + (y-S[,2])^2 );
    p = .5*exp(-r/10);
    #p = (r <= .8);
    return(p)
}

sensorprob1 = function(S,x,y)
{
    r = sqrt((x-S[,1])^2 + (y-S[,2])^2 );
    p = .8*(r<10);
    #p = (r <= .8);
    return(p)
}