######################################################################
##  Copyright 2017 Carl F. Falk and Katerina M. Marcoulides
##
##  This program is free software: you can redistribute it and/or
##    modify it under the terms of the GNU General Public License as
##    published by the Free Software Foundation, either version 2 of
##    the License, or (at your option) any later version.
##
##    This program is distributed in the hope that it will be useful,
##    but WITHOUT ANY WARRANTY; without even the implied warranty of
##    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
##    GNU General Public License for more details.
##    <http://www.gnu.org/licenses/>
##

# Given a fitted lavaan model (e.g., CFA), prepares a table that contains parameters
# that can be fixed/freed as part of a model specification search.
#
# fitted.model - an object of class "lavaan" that contains the initially fitted model for the search
# loadings - a logical value that indicates whether cross-loadings will be part of the search
# fcov - a logical value indicating whether factor covariances will be part of the search
# errors - a logical value indicating whether error covariances will be part of the search
search.prep<-function(fitted.model,loadings=TRUE,fcov=TRUE,errors=FALSE){
  
  # Parameter table from fitted model
  ptab<-partable(fitted.model)
  
  # Full parameter table
  fulltab<-lavaan:::lav_partable_full(ptab)
  
  # Merge together so we have what's free
  mergetab<-lavaan:::lav_partable_merge(ptab,fulltab, remove.duplicated=TRUE, warn=FALSE)
  
  idx<-vector("numeric")
  null.val<-vector("numeric")
  
  lv.names<-lavaan:::vnames(ptab,type="lv") # extract names of latent vars
  ov.names<-lavaan:::vnames(ptab,type="ov.nox") # indicators
  
  # Loadings
  if(loadings){
    # Index for which rows of table are loadings
    indx<-mergetab$op=="=~"

    # Omit any loadings that are fixed to a non-zero value
    # These may be fixed for identification
    indx[mergetab$free[indx]==0&mergetab$ustart[indx]!=0]<-FALSE
    
    idx<-c(idx,which(indx))
    null.val<-c(null.val,rep(0,sum(indx)))
  }
  
  # Factor cov
  if(fcov){
    # Index for which rows are cov among latent vars
    indx<-mergetab$op=="~~"&mergetab$lhs%in%lv.names&mergetab$lhs!=mergetab$rhs
    idx<-c(idx,which(indx))
    null.val<-c(null.val,rep(0,sum(indx)))
  }
  
  # Error covariances among indicators
  if(errors){
    # Index for which rows are cov among indicators
    indx<-mergetab$op=="~~"&mergetab$lhs%in%ov.names&mergetab$lhs!=mergetab$rhs
    idx<-c(idx,which(indx))
    null.val<-c(null.val,rep(0,sum(indx)))
  }
  
  # Check whether any parameters in the table are involved in any constraints
  flag<-check.const(fitted.model,mergetab[idx,])

  if(flag){
    warning("Some parameters in this table may be involved in (in)equality constraints")
  }

  # Clean up
  mergetab<-mergetab[idx,c("lhs","op","rhs","block","free","label")]
  mergetab$nullval<-null.val
  mergetab$free<-ifelse(mergetab$free==0,0,1)
  
  # Sort table
  lvnames<-lavNames(mergetab,"lv")
  ovnames<-lavNames(mergetab,"ov")
  mergetab<-mergetab[order(mergetab$block, # by group (or block)
                           is.na(match(mergetab$lhs,c(lvnames))), # put latent vars first
                           mergetab$op, # by op w/in lv and ov sections
                           match(mergetab$lhs,c(lvnames,ovnames)), # then by lv and ov names on lhs
                           match(mergetab$rhs,c(lvnames,ovnames))),] # then by same on rhs
  rownames(mergetab)<-1:nrow(mergetab)
  
  ret<-mergetab
  
  return(ret)
}

# Given a fitted lavaan model and a search table, checks whether any of the
# parameters in the search table are involved in (in)equality constraints
#
# fitted.model - fitted lavaan model
# pars - specification table (e.g., prepared by search.prep and optionally further modified by user)
check.const<-function(fitted.model,pars){
  ptab<-parTable(fitted.model)
  const<-ptab[ptab$op=="=="|ptab$op==">"|ptab$op=="<",]
  if(nrow(const)>0){
    plabels<-c(pars$label,pars$plabel)
    plabels<-plabels[plabels!=""]
    flag<-FALSE
    for(j in 1:nrow(const)){
      var.lhs<-all.vars(parse(file="", text=const$lhs[j]))
      var.rhs<-all.vars(parse(file="", text=const$rhs[j]))
      if(any(plabels %in% c(var.lhs,var.rhs) ))
        flag<-TRUE
    }
  } else {
    flag<-FALSE
  }
  return(flag)
}

# Adds a parameter to the given search table. Checks whether parameter
# is involved in any (in)equality constraints in a fitted lavaan model
#
# fitted.model - fitted lavaan model
# ptab - search table
# syntax - model.syntax specifying the parameter to add to the current table
# nullval - optional numeric value specifying what the parameter should be fixed to (when fixed)
# free - optional logical value specifying whether the parameter should initially be set free (or not)
# block - optional numeric value specifying the group number to which the parameter corresponds
add.param<-function(fitted.model, ptab, syntax, nullval=NULL, free=NULL, block=NULL) {
  
  newpar<-lavaanify(syntax)

  if(!is.null(block)){
    newpar$block<-block
  }
  if(!is.null(free)){
    newpar$free<-ifelse(free,1,0)    
  }
  if(!is.null(nullval)){
    newpar$nullval<-nullval    
  } else {
    newpar$nullval<-newpar$ustart
  }
  
  newpar<-newpar[,c("lhs","op","rhs","block","free","label","nullval")]
  
  # Check whether newpar includes more than one row
  if(nrow(newpar)>1){
    newpar<-newpar[1,]
    warning("lavaanify() results in more than one row of parameters, adding only the first row")
  }
  
  # Check whether parameter already exists in table
  matches<-par.matches(ptab,paste(newpar$lhs,newpar$op,newpar$rhs,sep=""),block=block)
  if(any(matches)){
    stop("Parameter already found in table")
  }
  
  # Check to see if parameter is part of constraint in existing model
  flag<-check.const(fitted.model, newpar)
  if(flag){
    warning("Some parameters added to this table may be involved in (in)equality constraints")
  }
  
  # Add newpar to existing table
  ptab<-rbind(ptab,newpar)
  
  return(ptab)
}

# Internal function used to match parameters in label against
# those in a search table (ptab)
par.matches<-function(ptab,label,block=NULL){
  
  # Collapse whitespace in label (if any)
  label<-gsub("[[:space:]]","", label)
  
  # Look through lhsoprhs
  lhsoprhs<-paste(ptab$lhs,ptab$op,ptab$rhs,sep="")
  
  # Look through labels
  matches<-label==lhsoprhs | label==ptab$label
  
  # Check against block
  if(!is.null(block)){
    matches <- matches & ptab$block==block
  }
  return(matches)
}

# Internal function that modifies search table
manip.ptab<-function(ptab,label,task,nullval=NULL,block=NULL){
  
  if(!task %in% c("remove","free","fix","nullval")){
    stop()
  }

  matches<-par.matches(ptab,label,block=block)
  
  nmatches<-sum(matches)
  if(nmatches>1){
    warning("More than one match for parameter found")
  } else if (nmatches<1){
    warning("No matches for parameter found")
  }
  
  if(task=="remove"){
    # Remove from table
    ptab<-ptab[!matches,]
    
  } else if (task=="free"){
    # Set free
    ptab$free[matches]<-1
    
  } else if (task=="fix"){
    # Fix
    ptab$free[matches]<-0
    
    if(!is.null(nullval)){
      ptab$nullval[matches]<-nullval
    }
  } else if (task=="nullval"){
    # Change nullval
    ptab$nullval[matches]<-nullval
  }
  
  return(ptab)
}

# Given a search table and a label, sets the parameter free in the search table
#
# ptab - search table
# label - combination of lhs, op, and rhs as would appear in lavaan model.syntax
#    or a parameter label from the label column of a lavaan parameter table
# block - optional numeric value specifying the group number to which the parameter corresponds
free.param<-function(ptab, label, block=NULL){
  ptab<-manip.ptab(ptab,label,"free",block=block)
  return(ptab)
}

# Given a search table and a label, sets the parameter fixed in the search table
#
# ptab - search table
# label - combination of lhs, op, and rhs as would appear in lavaan model.syntax
#    or a parameter label from the label column of a lavaan parameter table
# nullval - optional numeric value specifying what the parameter should be set to when fixed (existing value in ptab is otherwise used)
# block - optional numeric value specifying the group number to which the parameter corresponds
fix.param<-function(ptab, label, nullval=NULL, block=NULL){
  ptab<-manip.ptab(ptab,label,"fix",nullval=nullval,block=block)
  return(ptab)
}

# Given a search table and a label, removes the parameter from the search table
#
# ptab - search table
# label - combination of lhs, op, and rhs as would appear in lavaan model.syntax
#    or a parameter label from the label column of a lavaan parameter table
# block - optional numeric value specifying the group number to which the parameter corresponds
rm.param<-function(ptab, label, block=NULL){
  ptab<-manip.ptab(ptab,label,"remove",block=block)
  return(ptab)
}

# Given a search table and a label, sets the value that the parameter would be
# fixed to if the parameter is fixed
#
# ptab - search table
# label - combination of lhs, op, and rhs as would appear in lavaan model.syntax
#    or a parameter label from the label column of a lavaan parameter table
# nullval - numeric value specifying what the parameter should be set to when fixed
# block - optional numeric value specifying the group number to which the parameter corresponds
nullval.param<-function(ptab, label, nullval, block=NULL){
  ptab<-manip.ptab(ptab,label,"nullval",nullval=nullval,block=block)
  return(ptab)
}

# Given a fitted lavaan model and a search table, refits the model using the search
# table as specifying what changes should be done (parmeters fixed/freed)
# 
# fitted.model - fitted model of class lavaan
# ptab - search table
refit.model<-function(fitted.model, ptab){

  # Extract parameter table
  tab<-parTable(fitted.model)
  
  # Expand to get full table
  fulltab<-lavaan:::lav_partable_full(tab)
  mergetab<-lavaan:::lav_partable_merge(tab,fulltab, remove.duplicated=TRUE, warn=FALSE)
    
  # Modify parameter table based on new input matrix
    
  # Obtain indices for matching
  mergetab$labels<-paste(mergetab$lhs,mergetab$op,mergetab$rhs,sep="")
  ptab$labels<-paste(ptab$lhs,ptab$op,ptab$rhs,sep="")
  midx<-vector("numeric")
  nidx<-vector("numeric")
  for(j in 1:nrow(ptab)){
    idx<-which(ptab$labels[j]==mergetab$labels & ptab$block[j]==mergetab$block)
    if(length(idx)>0){
      midx<-c(midx,idx)
      nidx<-c(nidx,j)
    }
  }
  
  # Replace free parameters
  mergetab$free[midx]<-ptab$free[nidx]

  # If some parameters are now fixed; what are they fixed to?<-function
  mergetab$ustart[mergetab$free!=0]<-NA
  mergetab[midx,]$ustart[mergetab[midx,]$free==0]<-ptab[nidx,]$nullval[ptab[nidx,]$free==0]
  
  # Get rid of old parameter estimates and starting values
  mergetab$est<-NULL
  mergetab$se<-NULL
  mergetab$start<-NULL
  mergetab$labels<-NULL
  
  # Re-order values in mergetab
  lvnames<-lavNames(mergetab,"lv")
  ovnames<-lavNames(mergetab,"ov")
  mergetab<-mergetab[order(mergetab$block, # by group (or block)
                           is.na(match(mergetab$lhs,c(lvnames))), # put latent vars first
                           mergetab$op, # by op w/in lv and ov sections
                           match(mergetab$lhs,c(lvnames,ovnames)), # then by lv and ov names on lhs
                           match(mergetab$rhs,c(lvnames,ovnames))),] # then by same on rhs
  
  # Refit model
  prevmodel<-as.list(fitted.model@call)
  prevmodel$model<-mergetab
  newmod<-try(do.call("lavaan",prevmodel[-1]),silent=TRUE)
  
  return(newmod)
}

# Given a fitted lavaan model, a search table, and an objective criterion,
# performs a Tabu model specification search. Currently only supports
# neighbors that are 1 move away from the current model.
# 
# fitted.model - fitted model of class lavaan
# ptab - search table (e.g., created by search.prep) that lists candidate parameters
#   that can be modified as part of the search and how the parameters can be modified (fixed to what values)
# obj - objective function to be MINIMIZED. Any function that takes a lavaan object as the sole argument and returns
#   a numeric value can be used.
# niter - number of Tabu iterations to perform
# tabu.size - size of Tabu list
tabu.sem<-function(init.model,ptab,obj,niter=30,tabu.size=5){
  
  # Initialize objective function and best model
  best.obj<-current.obj<-obj(init.model)
  best.model<-current.model<-init.model
  best.binvec<-current.binvec<-ptab
  
  tabu.list<-vector("numeric")
  
  # Do iterations
  for(it in 1:niter){
    
    # Loop through all neighbors
    tmp.obj<-vector("numeric")
    tmp.mod<-list()
    tmp.vec<-list()
    for(j in 1:nrow(current.binvec)){
      tmp.binvec<-current.binvec
      bin<- 1-tmp.binvec$free[j]
      tmp.binvec$free[j]<-bin
      fitmodel<-refit.model(init.model,tmp.binvec)
      
      if(fitmodel@Fit@converged&!any(is.na(fitmodel@Fit@se))){
        fit.val<-obj(fitmodel)
      } else {
        fit.val<-NA
      }
      
      tmp.obj<-c(tmp.obj,fit.val)
      tmp.mod[[j]]<-fitmodel
      tmp.vec[[j]]<-tmp.binvec
    }
    
    # Check which indices result in a valid objective function
    valid<-which(!is.na(tmp.obj))
    
    # Get just models not on Tabu list
    valid<-valid[!(valid %in% tabu.list)]
    
    # Out of valid models, pick model with best objective function value
    indx<-which.min(tmp.obj[valid])
    
    # Move current state to next model
    current.obj<-(tmp.obj[valid])[indx]
    current.mod<-(tmp.mod[valid])[[indx]]
    current.binvec<-(tmp.vec[valid])[[indx]]

    # Update Tabu list
    tabu.list<-c(valid[indx],tabu.list)
    if(length(tabu.list)>tabu.size){
      tabu.list<-tabu.list[1:tabu.size]
    }
    
    # Update if the current model is better than the best model
    if(current.obj<best.obj){
      best.obj<-current.obj
      best.mod<-current.mod
      best.binvec<-current.binvec
      tabu.list<-vector("numeric") # Clear Tabu list
    }
  }
  
  ret<-list()
  ret$best.obj<-best.obj
  ret$best.mod<-best.mod
  ret$best.binvec<-best.binvec
  
  return(ret)
}