import numpy as np
from quantregpy._fortran import rqbr, rqfnb, rqs
from scipy.stats import norm
from scipy.stats import t as studentT
from sklearn.linear_model import LinearRegression
from copy import deepcopy
import pandas as pd
import collections
from patsy import dmatrices, DesignMatrix

Fit = collections.namedtuple(
  'Fit', 
  'na_action formula terms x_design y_design call tau weights residuals rho method model coefficients fitted_values na_message Class')

def bandwidth_rq(p, n, hs = True, alpha = 0.05):
  # Bandwidth selection for sparsity estimation two flavors:
  #	Hall and Sheather(1988, JRSS(B)) rate = O(n^{-1/3})
  #	Bofinger (1975, Aus. J. Stat)  -- rate = O(n^{-1/5})
  # Generally speaking, default method, hs=TRUE is preferred.

  x0 = norm.ppf(p)#qnorm(p)
  f0 = norm.pdf(x0)#dnorm(x0)
  if(hs):
    bandwidth = ( n**(-1./3.) * norm.ppf(1. - alpha/2.)**(2./3.) * ((1.5 * f0**2.)/(2. * x0**2. + 1.))**(1./3.) )
  else:
    bandwidth = n**-0.2 * ((4.5 * f0**4.)/(2. * x0**2. + 1.)**2.)**0.2
  return bandwidth

def print_rq(x : Fit, *args):
  print("Call:")
  print(x.call)
  coef = x.coefficients
  print("\nCoefficients:\n")
  print(coef, args)
  rank = x.rank
  nobs = x.residuals.shape[0]
  p = coef.shape[1] if( len(coef.shape) > 1) else coef.shape[0]
  rdf = nobs - p
  print(f"\nDegrees of freedom:{nobs}total;{rdf}residual\n")
  print(f"{x.na_message}\n")

def print_summary_rq(x, digits = 5, *args):
  print("\nCall: ")
  print(x.call)
  coef = x.coef
  tau = x.tau
  print("\ntau: ")
  print([round(t,digits) for t in tau.tolist()], args)
  print("\nCoefficients:\n")
  print([round(c, digits) for c in coef.tolist()], args)

def failIfMissingData(df):
  if df.isnull().values.any():
    print("Missing data in data")
    raise ValueError

def rq(formula : str, tau : np.array, data : pd.DataFrame, subset = None, weights = None,
       na_action = failIfMissingData,
       method = "br", model = True, contrasts = None, *args):
  tau = np.array([tau]) if type(tau) is float else tau
  expandedArgs = ", ".join(str(arg) for arg in args)
  baseCall = f"rq(formula = {formula}, tau = {tau}, data = {data}, subset = {subset}, weights = {weights}, na_action = {na_action}, method = {method}, model = {model}, constrasts = {contrasts}"
  call = baseCall + ", " + expandedArgs
  mf = baseCall + f", args = {args}"
  m = ",".join([param for param in mf.split(",") if param in ("formula", "data", "subset", "weights", "na_action")])
  Y, X = dmatrices(formula, data)
  mt = X.design_info.term_names
  eps = np.finfo(float).eps**(2/3)
  Rho = lambda u, tau: u * (tau - (u < 0)) 
  if(tau.shape[0]>1):
    if(np.any(tau < 0) or np.any(tau > 1)):
      print("invalid tau:  taus should be >= 0 and <= 1")
      raise ValueError
    tau[tau == 0] = eps
    tau[tau == 1] = 1 - eps
    coef = np.zeros((X.shape[1],tau.shape[0]))
    rho = np.zeros(tau.shape[0])
    fitted = np.zeros((X.shape[0], tau.shape[0]))
    resid = np.zeros((X.shape[0], tau.shape[0]))
    for i in range(tau.shape[0]):
      z = rq_wfit(X, Y, tau[i], weights, method, *args) if not (weights is None) else rq_fit(X, Y, tau[i], method, *args)
      coef[:,i] = z.coefficients
      resid[:,i] = z.residuals
      rho[i] = np.sum(Rho(z.residuals,tau[i]))
      fitted[:,i] = Y - z.residuals
    taulabs = f"tau={np.around(tau,3)}"
    fit = deepcopy(z)
    fit.coefficients = coef
    fit.residuals = resid
    fit.fitted_values = fitted
    if(method == "lasso"): 
      fit.Class = ("lassorqs","rqs")
    elif(method == "scad"):
      fit.Class = ("scadrqs","rqs")
    else:
      fit.Class = "rqs"
  else:
    process = (tau < 0) or (tau > 1)
    if(tau == 0):
      tau = eps
    if(tau == 1):
      tau = 1 - eps
    fit = rq_wfit(X, Y, tau, weights, method, *args) if not (weights is None) else rq_fit(X, Y, tau, method, *args) 
    if(process):
      rho = [fit.sol[1,:],fit.sol[3,:]]
    else:
      rho = np.sum(Rho(fit['residuals'],tau))  
  if(method == "lasso"):
    fit['Class'] = ("lassorq","rq")
  elif(method == "scad"):
    fit['Class'] = ("scadrq","rq")
  else:
    fit['Class'] = "rq.process" if process else "rq"
  fit['na_action'] = na_action
  fit['formula'] = formula
  fit['terms'] = mt
  fit['x_design'] = X.design_info
  fit['y_design'] = Y.design_info
  fit['call'] = call
  fit['tau'] = tau
  fit['weights'] = weights
  fit['rho'] = rho
  fit['method'] = method
  fit['na_message'] = "" # unsure what to do here m.na_message
  if(model):
    fit['model'] = mf
  return fit

def rq_fit(x : np.array, y : np.array, tau = 0.5, method = "br", *args):
  if (method == "fn"): 
    fit = rq_fit_fnb(x, y, tau, *args)
  elif (method == "fnb"):
    fit = rq_fit_fnb(x, y, tau, *args)
  #elif (method == "fnc"):
  #  fit = rq_fit_fnc(x, y, tau, *args)
  #elif (method == "pfn"):
  #	fit = rq_fit_pfn(x, y, tau, *args)
  if (method == "br"):
    fit = rq_fit_br(x, y, tau, *args)
  #elif (method == "lasso"):
  #	fit = rq_fit_lasso(x, y, tau, *args)
  #elif (method == "scad"):
  #	fit = rq_fit_scad(x, y, tau = tau, *args)
  else:
    print(f"rq.fit.{method} not yet implemented")
    raise ValueError

  fit['fitted_values'] = y - fit['residuals']
  return fit

def dropNpColumn(npmat, j):
  if j == 0:
    return npmat[:,1:]
  elif j == npmat.shape[1] - 1:
    return npmat[:,:j]
  else:
    return np.concatenate((npmat[:,:j],npmat[:,:j+1]), axis=1)

def rq_wfit(x, y, tau, weights, method = "br",  *args):
  if(any(weights < 0)):
    raise ValueError("negative weights not allowed")
#	contr <- attr(x, "contrasts")
  if len(weights.shape) != 2:
    weights = weights.reshape((weights.shape[0],1))
  wx = x * weights 
  wy = y * weights

  if (method == "fn"): 
    fit = rq_fit_fnb(x, y, tau, *args)
  elif (method == "fnb"):
    fit = rq_fit_fnb(x, y, tau, *args)
  #elif (method == "fnc"):
  #  fit = rq_fit_fnc(x, y, tau, *args)
  #elif (method == "pfn"):
  #  fit = rq_fit_pfn(x, y, tau, *args)
  if (method == "br"):
    fit = rq_fit_br(wx, wy, tau, *args)
  else:
    print(f"rq.fit.{method} not yet implemented")
    raise ValueError
  if(len(fit.get('sol',[])) > 0):
    fit['fitted_values'] = np.matmul( x , fit['sol'][3:,:])
  else:
    yhat = np.matmul( x , fit['coefficients'])
    ny = 1 if len(y.shape) == 1 else y.shape[1]
    fit['fitted_values'] = yhat.reshape((yhat.shape[0], ny))
  fit['residuals'] = y - fit["fitted_values"]
  fit['weights'] = weights
  return fit

# Function to compute regression quantiles using original simplex approach
# of Barrodale-Roberts/Koenker-d'Orey.  There are several options.
# The options are somewhat different than those available for the Frisch-
# Newton version of the algorithm, reflecting the different natures of the
# problems typically solved.  Succintly BR for "small" problems, FN for
# "large" ones.  Obviously, these terms are conditioned by available hardware.
#
# Basically there are two modes of use:
# 1.  For Single Quantiles:
#
#       if tau is between 0 and 1 then only one quantile solution is computed.
#
#       if ci = FALSE  then just the point estimate and residuals are returned
#		If the column dimension of x is 1 then ci is set to FALSE since
#		since the rank inversion method has no proper null model.
#       if ci = TRUE  then there are two options for confidence intervals:
#
#               1.  if iid = TRUE we get the original version of the rank
#                       inversion intervals as in Koenker (1994)
#               2.  if iid = FALSE we get the new version of the rank inversion
#                       intervals which accounts for heterogeneity across
#                       observations in the conditional density of the response.
#                       The theory of this is described in Koenker-Machado(1999)
#               Both approaches involve solving a parametric linear programming
#               problem, the difference is only in the factor qn which
#               determines how far the PP goes.  In either case one can
#               specify two other options:
#                       1. interp = FALSE returns two intervals an upper and a
#                               lower corresponding to a level slightly
#                               above and slightly below the one specified
#                               by the parameter alpha and dictated by the
#                               essential discreteness in the test statistic.
#				interp = TRUE  returns a single interval based on
#                               linear interpolation of the two intervals
#                               returned:  c.values and p.values which give
#                               the critical values and p.values of the
#                               upper and lower intervals. Default: interp = TRUE.
#                       2.  tcrit = TRUE uses Student t critical values while
#                               tcrit = FALSE uses normal theory ones.
# 2. For Multiple Quantiles:
#
#       if tau < 0 or tau >1 then it is presumed that the user wants to find
#       all of the rq solutions in tau, and the program computes the whole
#	quantile regression solution as a process in tau, the resulting arrays
#	containing the primal and dual solutions, betahat(tau), ahat(tau)
#       are called sol and dsol.  These arrays aren't printed by the default
#       print function but they are available as attributes.
#       It should be emphasized that this form of the solution can be
#	both memory and cpu quite intensive.  On typical machines it is
#	not recommended for problems with n > 10,000.
#	In large problems a grid of solutions is probably sufficient.
#
def rq_fit_br(x, y, tau = 0.5, alpha = 0.1, ci = False, iid = True,
  interp = True, tcrit = True):
    tol = np.finfo(float).eps**(2/3)
    eps = tol
    big = np.finfo(float).max**(2/3)
    p = x.shape[1]
    n = x.shape[0]
    ny = y.shape[1]
    nsol = 2
    ndsol = 2
    # Check for Singularity of X since br fortran isn't very reliable about this
    #storage.mode(y) <- "double"
    if (np.linalg.matrix_rank(x) < p):
        raise ValueError("Singular design matrix")
    if (tau < 0) or (tau > 1):
      nsol = 3 * n
      ndsol = 3 * n
      lci1 = False
      qn = np.array([0] * p)
      cutoff = 0
      tau = -1
    else:
        if (p == 1):
            ci = False
        if (ci):
          lci1 = True
          if (tcrit):
            cutoff = studentT.ppf(1 - alpha/2, n - p)
          else: 
            cutoff = norm.ppf(1 - alpha/2.)
          if (not iid):
            h = bandwidth_rq(tau, n, hs = True)
            bhi = rq_fit_br(x, y, tau + h, ci = False)
            bhi = bhi['coefficients']
            blo = rq_fit_br(x, y, tau - h, ci = False)
            blo = blo['coefficients']
            dyhat = np.matmul(x, (bhi - blo))
            if (np.any(dyhat <= 0)):
              pfis = (100 * np.sum(dyhat <= 0))/n
              print(f"{pfis}percent fis <=0")
            f = np.maximum(eps, (2 * h)/(dyhat - eps))
            qn = np.array([0]*p)
            for j in range(p):
              tempX = dropNpColumn(x, j)
              tempY = x[:,j]
              lr = LinearRegression().fit(tempX, tempY, sample_weight=f)
              qnj = lr.predict(tempX) - tempY
              qn[j] <- np.sum(qnj * qnj)
          else:
            qn = 1./np.diagonal(np.linalg.inv(np.matmul(x.T,x)))
        else:
            lci1 = False
            qn = np.array([0]*p)
            cutoff = 0
    sFor,waFor,wbFor,nsolFor,ndsFor= np.zeros([n]), np.zeros([(n + 5), (p + 4)]), np.zeros(n), nsol,ndsol
    tnmat = np.zeros([4,p])
    flag,coef,resid,sol,dsol,lsol, h, qn, cutoff, ci, tnmat = rqbr(p+3,x,y,tau,tol,sFor,waFor,wbFor,nsolFor,ndsFor,tnmat, big, lci1)
    if (flag != 0):
        if flag == 1:
          print("Solution may be nonunique")
        else:
          print("Premature end - possible conditioning problem in x")
    if (tau < 0) or (tau > 1):
        sol = sol[1:((p + 3) * lsol)]
        dsol = dsol[1:(n * lsol)]
        return({"sol" : sol, "dsol" : dsol})
    if (not np.any(ci)):
        dual = dsol.T.flatten()[0:n]
        yhatCols = 1 if len(coef.shape) < 2 else coef.shape[1]
        yhat = np.matmul(x, coef).reshape((x.shape[0], yhatCols))
        return(dict(coefficients = coef, x = x, y = y, residuals = y - yhat, dual = dual))
    if (interp):
        Tn = tnmat
        Tci = ci
        Tci[3, :] = Tci[3, :] + (np.abs(Tci[4, :] - Tci[3, :]) * (cutoff -
            np.abs(Tn[3, :])))/np.abs(Tn[4,: ] - Tn[3, :])
        Tci[2, :] = Tci[2, :] - (np.abs(Tci[1,: ] - Tci[2, :]) * (cutoff -
            np.abs(Tn[2, :])))/np.abs(Tn[1, :] - Tn[2, :])
        Tci[2, np.isnan(Tci[2,:]) ] = -big
        Tci[3, np.isnan(Tci[3,:]) ] = big
        coefficients = np.concatinate((coef,Tci[2:4, : ].T), axis = 1)
        residuals = y - np.matmul(x, coef)
        return(dict(coefficients = coefficients, residuals = residuals))
    else:
        Tci = ci
        coefficients = np.concatenate([coef, Tci.T], axis=1)
        residuals = y - np.matmul(x , coef)
        c_values = tnmat.T
        c_values = np.fliplr(c_values) 
        p_values = studentT.cdf(c_values, n - p) if (tcrit) else norm.cdf(c_values)
        return dict(coefficients = coefficients, residuals = residuals,
            c_values = c_values, p_values = p_values)

def rq_fit_fnb (x, y, tau = 0.5, beta = 0.99995, eps = 1e-06):
  n = y.shape[0]
  p = 1 if len(x.shape) == 1 else x.shape[1]
  if(n != x.shape[0]):
    raise ValueError("x and y don't match n")
  if (tau < eps) or (tau > 1 - eps):
    raise ValueError("No parametric Frisch-Newton method.  Set tau in (0,1)")
  rhs = (1 - tau) * np.sum(x, axis = 0)
  d   = np.ones(n)
  u   = np.ones(n)
  wn = np.zeros(10*n)
  wn[0:n] = (1-tau) #initial value of dual solution
  a = x.T
  info, wp = rqfnb(a,-y,rhs,d,u,beta,eps,wn)
  if (info != 0):
    raise ValueError(f"Error info = {info} in stepy: singular design")
  coefficients = wp[0:p]
  residuals = y - np.matmul(x,coefficients)
  return dict(coefficients=coefficients, tau=tau, residuals=residuals)

def rq_fit_fnc(x, y, R, r, tau = 0.5, beta = 0.9995, eps = 1e-06):
  n1 = y.shape[0]
  n2 = r.shape[0]
  p = 1 if len(x.shape) == 1 else x.shape[1]
  if (n1 != x.shape[0]):
    raise ValueError("x and y don't match n1")
  if (n2 != R.shape[0]):
    raise ValueError("R and r don't match n2")
  if (p != ( 1 if len(R.shape) == 1 else R.shape[1] ) ):
    raise ValueError("R and x don't match p")
  if (tau < eps) or (tau > 1 - eps):
    raise ValueError("No parametric Frisch-Newton method.  Set tau in (0,1)")
#    rhs <- (1 - tau) * apply(x, 2, sum)
#    u <- rep(1, max(n1,n2)) #upper bound vector and scratch vector
#    wn1 <- rep(0, 9 * n1)
#    wn1[1:n1] <- (1 - tau) #store the values of x1
#    wn2 <- rep(0, 6 * n2)
#    wn2[1:n2] <- 1 #store the values of x2
#    z <- .Fortran("rqfnc", as.integer(n1), as.integer(n2), as.integer(p),
#	a1 = as.double(t(as.matrix(x))), c1 = as.double(-y),
#	a2 = as.double(t(as.matrix(R))), c2 = as.double(-r),
#	rhs = as.double(rhs), d1 = double(n1), d2 = double(n2),
#	as.double(u), beta = as.double(beta), eps = as.double(eps),
#        wn1 = as.double(wn1), wn2 = as.double(wn2), wp = double((p + 3) * p),
#        it.count = integer(3), info = integer(1), PACKAGE = "quantreg")
#    if (z$info != 0)
#        stop(paste("Error info = ", z$info, "in stepy2: singular design"))
#    coefficients <- -z$wp[1:p]
#    names(coefficients) <- dimnames(x)[[2]]
#    residuals <- y - x %*% coefficients
#    it.count <- z$it.count
#    list(coefficients=coefficients, tau=tau, residuals=residuals, it = it.count)
#}
def rqs_fit(x, y, tau = 0.5, tol = 0.0001):
  """ 
  function to compute rq fits for multiple y's
  """
  p = x.shape[1]
  n = x.shape[0]
  m = y.shape[1]

  flag, coef, e = rqs(
                          x,
                          y,
                          tau,
                          tol,
                          np.zeros([n]),
                          np.zeros([(n + 5) , (p + 2)]),
                          np.zeros(n))
  if(np.sum(flag)>0):
    if(np.any(flag==2)):
      print(f"{np.sum(flag==2)} out of {m} BS replications have near singular design")
    if(np.any(flag==1)):
      print(f"{np.sum(flag==1)} out of {m} may be nonunique")
      
  return(coef.T)