#$Id$
#
# Some of the CAKE R modules are based on mkin
# Modifications developed by Tessella for Syngenta: Copyright (C) 2011-2016 Syngenta
# Tessella Project Reference: 6245, 7247, 8361, 7414
 
#    The CAKE R modules are 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 3 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.
# 
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

# Performs an iteratively-reweighted least squares fit on a given CAKE model.
# Remark: this function was originally based on the "mkinfit" function, version 0.1.
#
# cake.model: The model to perform the fit on (as generated by CakeModel.R).
# observed: Observation data to fit to.
# parms.ini: Initial values for the parameters being fitted.
# state.ini: Initial state (i.e. initial values for concentration, the dependent variable being modelled).
# lower: Lower bounds to apply to parameters.
# upper: Upper bound to apply to parameters.
# fixed_parms: A vector of names of parameters that are fixed to their initial values.
# fixed_initials: A vector of compartments with fixed initial concentrations.
# quiet: Whether the internal cost functions should execute more quietly than normal (less output).
# atol: The tolerance to apply to the ODE solver.
# sannMaxIter: The maximum number of iterations to apply to SANN processes.
# control: ...
# useExtraSolver: Whether to use the extra solver for this fit.
CakeIrlsFit <- function (cake.model, 
                         observed, 
                         parms.ini = rep(0.1, length(cake.model$parms)), 
                         state.ini = c(100, rep(0, length(cake.model$diffs) - 1)), 
                         lower = 0, 
                         upper = Inf, 
                         fixed_parms = NULL, 
                         fixed_initials = names(cake.model$diffs)[-1], 
                         quiet = FALSE, 
                         atol=1e-6, 
                         sannMaxIter = 10000, 
                         control=list(),
                         useExtraSolver = FALSE,
                         ...) 
{
    NAind <-which(is.na(observed$value))
    mod_vars <- names(cake.model$diffs)
    observed <- subset(observed, name %in% names(cake.model$map))
    ERR <- rep(1,nrow(observed))
    observed <- cbind(observed,err=ERR)
    fitted_with_extra_solver <- 0
    
    obs_vars = unique(as.character(observed$name))
    
    if (is.null(names(parms.ini))) {
        names(parms.ini) <- cake.model$parms
    }
        
    mkindiff <- function(t, state, parms) {
        time <- t
        diffs <- vector()
        for (box in mod_vars) {
            diffname <- paste("d", box, sep = "_")
            diffs[diffname] <- with(as.list(c(time, state, parms)), 
                                    eval(parse(text = cake.model$diffs[[box]])))
        }
        return(list(c(diffs)))
    }
    
    if (is.null(names(state.ini))) {
        names(state.ini) <- mod_vars
    }
        
    parms.fixed <- parms.ini[fixed_parms]
    optim_parms <- setdiff(names(parms.ini), fixed_parms)
    parms.optim <- parms.ini[optim_parms]
    state.ini.fixed <- state.ini[fixed_initials]
    optim_initials <- setdiff(names(state.ini), fixed_initials)
    state.ini.optim <- state.ini[optim_initials]
    state.ini.optim.boxnames <- names(state.ini.optim)
    
    if (length(state.ini.optim) > 0) {
        names(state.ini.optim) <- paste(names(state.ini.optim), 
            "0", sep = "_")
    }
    
    costFunctions <- CakeInternalCostFunctions(cake.model, state.ini.optim, state.ini.optim.boxnames, 
            state.ini.fixed, parms.fixed, observed, mkindiff, quiet, atol=atol)
    
    ############### Iteratively Reweighted Least Squares#############
    ## Start with no weighting
    
    ## a prefitting step since this is usually the most effective method
    if(useExtraSolver)
    {
        pnames=names(c(state.ini.optim, parms.optim))
        fn <- function(P){
            names(P) <- pnames
            FF<<-costFunctions$cost(P)
            return(FF$model)}
        a <- try(fit <- solnp(c(state.ini.optim, parms.optim),fun=fn,LB=lower,UB=upper,control=control),silent=TRUE)

        fitted_with_extra_solver <- 1
        if(class(a) == "try-error")
        {
            print('solnp fails, try PORT or other algorithm by users choice, might take longer time. Do something else!')
            warning('solnp fails, switch to  PORT or other algorithm by users choice')
            ## now using submethod already
            a <- try(fit <- modFit(costFunctions$cost, c(state.ini.optim, parms.optim), lower = lower, upper = upper, method='Port',control=control))
            fitted_with_extra_solver <- 0
            if(class(a) == "try-error")
            {
                fit <- modFit(costFunctions$cost, c(state.ini.optim, parms.optim), lower = lower, upper = upper, method='L-BFGS-B',control=control)
            }
        }
    }else{
        # modFit parameter transformations can explode if you put in parameters that are equal to a bound, so we move them away by a tiny amount.
        all.optim <- ShiftAwayFromBoundaries(c(state.ini.optim, parms.optim), lower, upper)
        
        fit <- modFit(costFunctions$cost, all.optim, lower = lower, 
                          upper = upper,control=control,...)
    }
    
    if(length(control)==0)
      {
        irls.control <- list(maxIter=10,tol=1e-05)
        control <- list(irls.control=irls.control)
      }else{
        if(is.null(control$irls.control))
          {
            irls.control <- list(maxIter=10,tol=1e-05)
            control <- list(irls.control=irls.control)
          }
      }
      
    irls.control <- control$irls.control
    maxIter <- irls.control$maxIter
    tol <- irls.control$tol
    ####
    
    if(length(cake.model$map)==1){
        ## there is only one parent just do one iteration:
        maxIter <- 0
        
        if(fitted_with_extra_solver==1)## managed to fit with extra solver
        {
            fit$ssr <- fit$values[length(fit$values)]
            fit$residuals <-FF$residual$res
            ## mean square per variable
            if (class(FF) == "modCost") {
                names(fit$residuals)  <- FF$residuals$name
                fit$var_ms            <- FF$var$SSR/FF$var$N
                fit$var_ms_unscaled   <- FF$var$SSR.unscaled/FF$var$N
                fit$var_ms_unweighted <- FF$var$SSR.unweighted/FF$var$N

                names(fit$var_ms_unweighted) <- names(fit$var_ms_unscaled) <-
                    names(fit$var_ms) <- FF$var$name
            } else fit$var_ms <- fit$var_ms_unweighted <- fit$var_ms_unscaled <- NA
        }
        err1 <- sqrt(fit$var_ms_unweighted)
        ERR <- err1[as.character(observed$name)]
        observed$err <-ERR
    }
    
    niter <- 1
    ## insure one IRLS iteration
    diffsigma <- 100
    olderr <- rep(1,length(cake.model$map))
    while(diffsigma>tol & niter<=maxIter)
    {      
        if(fitted_with_extra_solver==1 && useExtraSolver)## managed to fit with extra solver
        {
            fit$ssr <- fit$values[length(fit$values)]
            fit$residuals <-FF$residual$res
            ## mean square per variable
            if (class(FF) == "modCost") {
                names(fit$residuals)  <- FF$residuals$name
                fit$var_ms            <- FF$var$SSR/FF$var$N
                fit$var_ms_unscaled   <- FF$var$SSR.unscaled/FF$var$N
                fit$var_ms_unweighted <- FF$var$SSR.unweighted/FF$var$N

                names(fit$var_ms_unweighted) <- names(fit$var_ms_unscaled) <-
                    names(fit$var_ms) <- FF$var$name
            } else fit$var_ms <- fit$var_ms_unweighted <- fit$var_ms_unscaled <- NA
        }

        err <- sqrt(fit$var_ms_unweighted)
        ERR <- err[as.character(observed$name)]
        costFunctions$set.error(ERR)
        diffsigma <- sum((err-olderr)^2)
        cat("IRLS iteration at",niter, "; Diff in error variance ", diffsigma,"\n")
        olderr <- err
        
        if(useExtraSolver)
        {
            fitted_with_extra_solver <- 1
            a <- try(fit <- solnp(fit$par,fun=fn,LB=lower,UB=upper,control=control),silent=TRUE)

            if(class(a) == "try-error")
            {
                fitted_with_extra_solver <- 0
                print('solnp fails during IRLS iteration, try PORT or other algorithm by users choice.This may takes a while. Do something else!') ## NOTE: because in kingui we switch off the warnings, we need to print out the message instead.
                warning('solnp fails during IRLS iteration, switch to  PORT or other algorithm by users choice')
            
                fit <- modFit(costFunctions$cost, fit$par, lower = lower, upper = upper, method='Port',control=list())
            }
        }else{
            # modFit parameter transformations can explode if you put in parameters that are equal to a bound, so we move them away by a tiny amount.
            fit$par <- ShiftAwayFromBoundaries(fit$par, lower, upper)
            
            fit <- modFit(costFunctions$cost, fit$par, lower = lower, upper = upper, control=control, ...)
        }
                
        niter <- niter+1       
       
        ### If not converged, reweight and fit                
    }
      
    if(fitted_with_extra_solver==1 && useExtraSolver){
        ## solnp used
        optimmethod <- 'solnp'
        fit$ssr <- fit$values[length(fit$values)]
        fit$residuals <-FF$residual$res
        ## mean square per varaible
        if (class(FF) == "modCost") {
            names(fit$residuals)  <- FF$residuals$name
            fit$var_ms            <- FF$var$SSR/FF$var$N
            fit$var_ms_unscaled   <- FF$var$SSR.unscaled/FF$var$N
            fit$var_ms_unweighted <- FF$var$SSR.unweighted/FF$var$N

            names(fit$var_ms_unweighted) <- names(fit$var_ms_unscaled) <-
                names(fit$var_ms) <- FF$var$name
        } else fit$var_ms <- fit$var_ms_unweighted <- fit$var_ms_unscaled <- NA
        np <- length(c(state.ini.optim, parms.optim))
        fit$rank <- np
        fit$df.residual <- length(fit$residuals) - fit$rank
        
        # solnp can return an incorrect Hessian, so we use another fitting method at the optimised point to determine the Hessian
        fitForHessian <- modFit(costFunctions$cost, fit$par, lower=lower, upper=upper, method='L-BFGS-B', control=list())
        
        fit$solnpHessian <- fit$hessian
        fit$hessian <- fitForHessian$hessian
    }
      
    ###########################################
    fit$mkindiff <- mkindiff
    fit$map <- cake.model$map
    fit$diffs <- cake.model$diffs
    
    out_predicted <- costFunctions$get.predicted()
    
    predicted_long <- wide_to_long(out_predicted, time = "time")
    fit$predicted <- out_predicted
    fit$start <- data.frame(initial = c(state.ini.optim, parms.optim))
    fit$start$type = c(rep("state", length(state.ini.optim)), 
    rep("deparm", length(parms.optim)))
    fit$start$lower <- lower
    fit$start$upper <- upper
    fit$fixed <- data.frame(value = c(state.ini.fixed, parms.fixed))
    fit$fixed$type = c(rep("state", length(state.ini.fixed)), 
    rep("deparm", length(parms.fixed)))

    fit$errmin <- CakeChi2(cake.model, observed, predicted_long, obs_vars, parms.optim, state.ini.optim, state.ini, parms.ini, fit$fixed)
    parms.all = c(fit$par, parms.fixed, state.ini)
    fit$penalties <- CakePenaltiesLong(parms.all, out_predicted, observed)
    fit$distimes <- data.frame(DT50 = rep(NA, length(obs_vars)), 
        DT90 = rep(NA, length(obs_vars)), row.names = obs_vars)
    fit$extraDT50<- data.frame(k1 = rep(NA, length(names(cake.model$map))), k2 = rep(NA, length(names(cake.model$map))), row.names = names(cake.model$map))   
        
    for (compartment.name in names(cake.model$map)) {
      type = names(cake.model$map[[compartment.name]])[1]
      fit$distimes[compartment.name, ] = CakeDT(type,compartment.name,parms.all,sannMaxIter)
      fit$extraDT50[compartment.name, ] = CakeExtraDT(type, compartment.name, parms.all)
    }

    fit$ioreRepDT = CakeIORERepresentativeDT("Parent", parms.all)
    fit$fomcRepDT = CakeFOMCBackCalculatedDT(parms.all)
    
    data <- merge(observed, predicted_long, by = c("time", "name"))
    
    names(data) <- c("time", "variable", "observed","err-var", "predicted")
    data$residual <- data$observed - data$predicted
    data$variable <- ordered(data$variable, levels = obs_vars)
    fit$data <- data[order(data$variable, data$time), ]
    
    fit$costFunctions <- costFunctions
    fit$upper <- upper
    fit$lower <- lower
    
    fit$atol <- atol
    
    class(fit) <- c("CakeFit", "mkinfit", "modFit") 
    return(fit)
}