## note: draws heavily from https://m-clark.github.io/models-by-example/tobit.html


const_tobit_ll <- function(par, X, y, ul = -Inf, ll = Inf) {
	
	# this function only takes a lower OR upper limit
	
	# parameters
	sigma = exp(par[2]) 
	const  = par[1]
	beta = c(const, 1)
	
	# create indicator depending on chosen limit
	if (!is.infinite(ll)) {
		limit = ll
		indicator = y > ll
	} else {
		limit = ul
		indicator = y < ul
	}
	
	# linear predictor
	lp = X %*% beta
	
	# log likelihood
	ll = sum(indicator * log((1/sigma)*dnorm((y-lp)/sigma)) ) + 
		sum((1-indicator) * log(pnorm((lp-limit)/sigma, lower = is.infinite(ll))))
	
	-ll
}

fit_const = function(formula, data, thresh = 66) {
	
	initmod = lm(formula, data = data)
	
	outcome.var = as.character(formula)[2]
  y = data %>% pull(outcome.var)
  
	X = model.matrix(initmod)
	init = c(coef(initmod)[1], log_sigma = log(summary(initmod)$sigma))
	
	param_est = optim(
		par = init,
		const_tobit_ll,
		y  = y,
		X  = X,
		ul = thresh,
		method  = 'BFGS',
		control = list(maxit = 2000, reltol = 1e-15)
	)
	
	ret = list(
		data = data,
		thresh = thresh,
		const_est = param_est$par[1]
	)
	
	class(ret) = "fit_const"
	
	ret
}


predict.fit_const = function(object, data.new = NULL) {
	
	if (is.null(data.new)) {data.new = object$data}
	
	pred = object$const_est + data.new$x
	pred = ifelse(pred < object$thresh, pred, object$thresh)
	
	pred
}

## Examples: 
## simulated data:
#
# df = 
# 	tibble(
# 		x = rnorm(100, mean = 30, sd = 10),
# 		y_obs = x + 30 + rnorm(100, 0, 3)
# 	) %>% 
# 	mutate(y_obs = ifelse(y_obs > 66, 66, y_obs))
# 
# df %>% 
# 	ggplot(aes(x = x, y = y_obs)) + 
# 	geom_point()
# 
# fit_const(y_obs ~ x, data = df)
# 
# df %>% 
# 	modelr::add_predictions(fit) %>% 
# 	ggplot(aes(x = x, y = y_obs)) + 
# 	geom_point() + 
# 	geom_point(aes(y = pred), color = "red")
# 
#
## data from winters
#
# pred_df = 
# 	real_dfs %>% 
# 	mutate(fmii = 66 - x) %>% 
# 	filter(severe == FALSE) %>% 
# 	rename(y_obs = y, delta_obs = delta)
# 
# df =
# 	pred_df %>%
# 	filter(name == "Stinear & Byblow")
# 
# df %>%
# 	ggplot(aes(x = x, y = y_obs)) +
# 	geom_point()
# 
# fit = fit_const(y_obs ~ x, data = df)
# 
# df %>%
# 	modelr::add_predictions(fit) %>%
# 	ggplot(aes(x = x, y = y_obs)) +
# 	geom_point() +
# 	geom_point(aes(y = pred), color = "red")
