## equivalent class generation

propensityTest <- function(data, treat, nonCfd, method) {     
  ##### values here are from Java Code (starting from 0, so pls plus 1)
  library(MatchIt)
  library(foreign)
  
  # read data from java directly
  data=as.data.frame(data)
  rownames(data) = c(1:nrow(data))

  treat = treat+1

  Rhand <- c()
  if(length(nonCfd) != 0) {
    Rhand <- paste0("data[,",nonCfd+1,"]")
    formula <- as.formula(paste("data[,treat] ~ ", paste(Rhand, collapse= "+")))
  }
  else
    formula <- as.formula("data[,treat] ~ NULL")
  print(formula)

    equivalent = matchit.subclass(formula, data = data, method = "subclass")

  # cat("length",length(equivalent),"\n")
  return(equivalent)
}



matchit.subclass <- function(formula, data = data, method) {
  ###################### matchit function
  m.out = matchit(formula, data = data, method = "subclass")
  QN = summary(m.out)$qn

  ###################### if #subclasses > 6
  colInd <- c()
  if(ncol(QN) > 6) {
    for(i in 1:ncol(QN)) {
      if(QN[3,i] != 0)
        colInd[length(colInd)+1] = i
    }
    if(length(colInd) > 6)
      QN = QN[,colInd[1:6]]
    else {
      s = sample(col(QN)[1,!col(QN)[1,] %in% colInd], (6-length(colInd)))
      QN = QN[,c(colInd,s)]
    }
  }

  equivalent = as.vector(data.matrix(QN))
  equivalent = c(equivalent,-1)    # -1 is used to distinguish QN and sample index
# print(equivalent)
  ###################### get treat & control in the order of increasing subclass index
  m.data.treat <- match.data(m.out,"treat") ## add "control" or "treat" to see each group
  m.data.treat = m.data.treat[order(m.data.treat$subclass),]
  m.treat=as.numeric(rownames(m.data.treat))
  m.data.control <- match.data(m.out,"control")
  m.data.control = m.data.control[order(m.data.control$subclass),]
  m.control=as.numeric(rownames(m.data.control))
  # return(m.treat)

  ##################### QN
  for(i in ncol(QN):1) {
    if(QN[1,i] == 0) {
      rmInd = which(m.data.control$subclass==i)
      if(length(rmInd)!=0) {
        m.control = m.control[-rmInd]
      }
      equivalent[3*(i-1)+2] = 0
      equivalent[3*(i-1)+3] = 0
      QN[1,i] = 0
      QN[3,i] = 0
    }
    if(QN[2,i] == 0) {
      rmInd = which(m.data.treat$subclass==i)
      if(length(rmInd)!=0) {
        m.treat = unlist(m.treat[!m.treat %in% rmInd])
      }
      equivalent[3*(i-1)+1] = 0
      equivalent[3*(i-1)+3] = 0
      QN[1,i] = 0
      QN[3,i] = 0
    }
    if(QN[3,i] == 0) {
      equivalent = equivalent[-c((3*i-2):(3*i))]
      next
    }
  }
#   print(QN)

  ###################### equivalent class generation
  TreatStartInd = 1
  ControlStartInd = 1
  for(i in 1:ncol(QN)) {
    if(QN[3,i] == 0)
      next
    numRecord = min(QN[1,i], QN[2,i])
  
    TreatEndInd = TreatStartInd + QN[1,i] - 1
    equivalent= c(equivalent, m.treat[TreatStartInd:TreatEndInd])
    TreatStartInd = TreatEndInd + 1
  
    ControlEndInd = ControlStartInd + QN[2,i] - 1
    equivalent = c(equivalent, m.control[ControlStartInd:ControlEndInd])
    ControlStartInd = ControlEndInd + 1
  }

  return(equivalent)
}