#
# This code implements the LBFP estimator in
# dimension d=2. The estimator is described
# in Carbon and Duchesne (2023) and in other
# references, such as ...
#

# The data must be stored in a matrix X with
# two columns. The LBFP function requires,
# as arguments
# b: a vector of size 2 with the bandwidths
# (default is standard deviation of each column
# of X times n^(-1/6) )
# Eval.grid: the coordinates of the X1 and X2 values
# used to create a mesh of points at which the LBFP
# will be calculated, in the form of a list with two 
# elements, the first the vector of values for X1
# and the second the vector of values for X2
# X: a matrix with two columns and n rows that
# contains the sample for which the LBFP is desired
#
# Returns a list whose first element is
# adata frame with x1 and x2 coordinates
# as well as LBFP value at those coordinates, and
# whose second element is the b vector.
#
#
# An example with evaluation of the LBFP for simulated
# data is provided below, under the function
#

# LBFP function
LBFP = function(b=c(-1,-1),X,Eval.grid){
  n = nrow(X)
  n.hat = n
  d = 2
  if(b[1]<0){
    b1 = sqrt(var(X[,1]))*n^(-1/6)
    b2 = sqrt(var(X[,2]))*n^(-1/6)
  }
  else{
    b1 = b[1]
    b2 = b[2]
  }
  b = c(b1,b2)
  # I_0(k) cells, in Hjort's notation 
  U = rep(0,d)
  L = rep(0,d)
  N = rep(0,d) 
  Grille = vector(mode = "list", length = d)
  for(s in 1:d){
    U[s] = max(X[,s]+0.000001)
    L[s] = min(X[,s]-0.000001)
    Long.Grille.s = trunc(((U[s]-L[s])/b[s])+2)
    #print(Long.Grille.s)
    Grille.s = L[s] + ((1:(Long.Grille.s+1))-1.5)*b[s]
    Grille[[s]] = Grille.s
    N[s] = length(Grille.s)
  }
  #
  I0k = Grille
  # Calculation of the v_k1,k2 in each cell
  v.mat = matrix(0,ncol=3,nrow=(N[1]-1)*(N[2]-1))
  # col1: coordinate in x1 of lower left corner of the cell
  # col2: coordinate in x2 of lower left corner of the cell
  # col3: number of observations in cell
  v.mat[,1] = rep(Grille[[1]][-N[1]],(N[2]-1))
  v.mat[,2] = rep(Grille[[2]][-N[2]],each=(N[1]-1))
  nl = dim(v.mat)[1]
  # line number in v.mat
  for(i in 1:nl){
    y.bloc = ceiling(i/(N[1]-1))
    x.bloc = i - (y.bloc-1)*(N[1]-1)
    if(y.bloc >= N[2]-1){
      y.sup = I0k[[2]][N[2]]
      y.inf = I0k[[2]][N[2]-1]
    }
    else {
      y.sup = I0k[[2]][y.bloc+1]
      y.inf = I0k[[2]][y.bloc]
    }
    if(x.bloc >= N[1]-1){
      x.sup = I0k[[1]][N[1]]
      x.inf = I0k[[1]][N[1]-1]
    }
    else {
      x.sup = I0k[[1]][x.bloc+1]
      x.inf = I0k[[1]][x.bloc]
    }
    v.i = sum(1*(X[,1]>=x.inf)*(X[,1]<x.sup)*
                (X[,2]>=y.inf)*(X[,2]<y.sup))
    v.mat[i,3] = v.i
  }
  
  # Calculatio of the I(k)
  Ik = vector(mode = "list", length = d)
  for(s in 1:d){
    ns = length(I0k[[s]])
    Ik[[s]] = ((I0k[[s]])[2:ns]+(I0k[[s]])[1:(ns-1)])/2
  }
  
  # Calculation of the 2^d = 4 lines corresponding
  # to neighboring cells
  lignes.x = function(x){
    x1 = max(Ik[[1]][Ik[[1]]<x[1] ])
    x2 = min(Ik[[1]][Ik[[1]]>=x[1] ])
    y1 = max(Ik[[2]][Ik[[2]]<x[2] ])
    y2 = min(Ik[[2]][Ik[[2]]>=x[2] ])
    x1 = x1-b[1]/2
    x2 = x2-b[1]/2
    y1 = y1-b[2]/2
    y2 = y2-b[2]/2
    i = which(
      ( ((v.mat[,1]>(x1-0.0001))&(v.mat[,1]<(x1+0.0001))) |
          ((v.mat[,1]>(x2-0.0001))&(v.mat[,1]<(x2+0.0001))) ) &
        ( ((v.mat[,2]>(y1-0.0001))&(v.mat[,2]<(y1+0.0001))) |
            ((v.mat[,2]>(y2-0.0001))&(v.mat[,2]<(y2+0.0001))) )   
    )
    return(i)
  }
  
  # Calculation of the c weights and of the LBFP
  f.hat = function(xx){
    lignes.xx = lignes.x(xx)
    vmat.xx = as.matrix(v.mat[lignes.xx,])
    vj1j2 = vmat.xx[,3]
    xk = c(mean(unique(vmat.xx[,1])),mean(unique(vmat.xx[,2])))
    u1 = (xx[1]-xk[1])/b[1]*(xx[1]>=xk[1])*(xx[1]<(xk[1]+b[1]))
    u2 = (xx[2]-xk[2])/b[2]*(xx[2]>=xk[2])*(xx[2]<(xk[2]+b[2]))
    j1 = c(1,1,0,0)
    j2 = c(1,0,1,0)
    cj1j2 = u1^j1*(1-u1)^(1-j1)*u2^j2*(1-u2)^(1-j2)
    return(sum(cj1j2*vj1j2)/(n.hat*prod(b)))
  }
  
  x1.grid = Eval.grid[[1]]
  x2.grid = Eval.grid[[2]]
  vals = data.matrix(expand.grid(x1.grid, x2.grid))
  vals.N = dim(vals)[1]
  LBFP.out = cbind(vals,rep(0,vals.N))
  for(i in 1:vals.N){
    ggg = length(lignes.x(c(vals[i,1],vals[i,2])))
    if(ggg!=4) {
      print(i)
      print(c(vals[i,1],vals[i,2]))
      print(lignes.x(c(vals[i,1],vals[i,2])))
    }
    lbfp = f.hat(c(vals[i,1],vals[i,2]))
    LBFP.out[i,3] = lbfp
  }
  
  LBFP.out = as.data.frame(LBFP.out)
  names(LBFP.out) = c("x1","x2","f.hat")
  return(list(LBFP.out,b))
}


# An example of simulated data, LBFP
# evaluation and plots
#
#X = matrix(rnorm(d*samsize),nrow=samsize)
library(mvtnorm)
VV = matrix(c(1,0.5*1.5,0.5*1.5,1.5^2),ncol=2)
XX = rmvnorm(n=100000, mean = rep(0, 2), sigma = VV)

bb = c(0.13,0.2)
n.grid = 50
eval.grid = list(seq(-3, 3, length = n.grid),
                 seq(-3, 3, length = n.grid) )

# With 100000 observations, this might take a minute to run ...
lbfp = LBFP(b=bb,X=XX,Eval.grid=eval.grid)
# to use default b : 
#lbfp = LBFP(X=XX,Eval.grid=eval.grid)

# To visualize the LBFP
# Contour plot
plot1 = levelplot(f.hat ~ x1*x2, data=lbfp[[1]],
                  xlim=c(-4,4),ylim=c(-4,4),main="LBFP")
grid.arrange(plot1,ncol=1)

# Perspective plot
LBFP.z = matrix(lbfp[[1]]$f.hat,nrow=n.grid,byrow=T)
par(mfrow=c(1,1))
persp(x = eval.grid[[1]],
      y = eval.grid[[2]], 
      z = LBFP.z, 
      phi = 35, theta = 70, d = 5,
      main="LBFP",
      xlab="x1",ylab="x2",zlab="LBFP")