import numpy as np
from quantregpy._fortran import rqbr, rqfnb, rqs, rqfnc
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
import logging

logger = logging.getLogger(__name__)
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)
  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)):
      logger.error("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 = dict()
    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

def rq_fit_br(x, y, tau = 0.5, alpha = 0.1, ci = False, iid = True,
  interp = True, tcrit = True):
  """
  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.
  """
  tol = np.finfo(float).eps**(2/3)
  eps = tol
  big = np.finfo(float).max**(2/3)
  p = x.shape[1]
  n = x.shape[0]
  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.flatten(), 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((n,9))
  wn[0:n,0] = (1-tau) #initial value of dual solution
  a = x.T
  wp, nit, info = 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].reshape((p,1))
  residuals = y - np.matmul(x,coefficients).flatten()
  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) * np.sum(x, axis = 0)
  u = np.ones(max(n1,n2)) #upper bound vector and scratch vector
  wn1 = np.zeros((n1, 9) )
  wn1[0:n1] = (1 - tau) #store the values of x1
  wn2 = np.zeros(( n2, 6))
  wn2[0:n2] = 1 #store the values of x2
  _, _, wp, nit, info = rqfnc( x.T, -y, R.T, -r, rhs, u, beta, eps, wn1, wn2)
  if (info != 0):
    raise ValueError(f"Error info = {info} in stepy: singular design")
  coefficients = -wp[0:p,0]
  residuals = y - np.matmul(x,coefficients).flatten()
  it_count = nit
  return dict(coefficients=coefficients, tau=tau, residuals=residuals)

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)