content: module AI.Learning.LogisticRegression where
  
  import Numeric.LinearAlgebra
  import Numeric.LinearAlgebra.Util
  
  import AI.Learning.Core
  import AI.Util.Matrix
  
  -- |Multivariate logistic regression. Given a vector /y/ of target variables and
  --  a design matrix /x/, this function fits the parameters /theta/ of a
  --  logistic model, that is,
  --
  --  y = sigmoid( x_1 * theta_1 + ... + x_n * theta_n )
  --
  --  Typically the values in the vector /y/ are either boolean (i.e. 0/1) or they
  --  represent frequency of observations, i.e. they are values between 0.0 and
  --  1.0.
  -- 
  --  The function fits /theta/ by numerically maximizing the likelihood function.
  --  It may be subject to overfit or non-convergence in the case where the number
  --  of observations is small or the predictors are highly correlated.
  lr :: Vector Double    -- targets (y)
     -> Matrix Double    -- design matrix (x)
     -> Vector Double    -- coefficient vector (theta)
  lr y x = lrHelper (lrLogLikelihood y x) theta0
      where theta0  = constant 0 (cols x)
  
  -- |Regularized logistic regression with quadratic penalty on the coefficients.
  --  You should standardize the coefficients of the design matrix before using
  --  this function, as the regularization procedure is sensitive to the scale
  --  of the predictors. See also 'lr'.
  lrRegularized :: Vector Double  -- targets (y)
                -> Matrix Double  -- design matrix (x)
                -> Bool           -- first column of design matrix is all ones?
                -> Double         -- regularization constant (lambda)
                -> Vector Double  -- coefficient vector (theta)
  lrRegularized y x useConst lambda = lrHelper costfun theta0
      where
          costfun = lrLogLikRegularized y x useConst lambda
          theta0  = constant 0 (cols x)
  
  -- |Helper function for logistic regressions. The first argument is a function
  --  that returns the cost and gradient for a given vector of parameters, and
  --  the second is the initial set of parameters to use.
  lrHelper :: (Vector Double -> (Double, Vector Double)) -> Vector Double -> Vector Double
  lrHelper fun theta0 = minimizeS cost grad theta0
      where cost    = negate . fst . fun -- negate because we call minimize
            grad    = negate . snd . fun
  
  -- |Cost function and derivative for logistic regression. This is maximized when
  --  fitting parameters for the regression.
  lrLogLikelihood :: Vector Double             -- targets (y)
                  -> Matrix Double             -- design matrix (x)
                  -> Vector Double             -- coefficients (theta)
                  -> (Double, Vector Double)   -- (value, derivative)
  lrLogLikelihood y x theta = (cost, grad)
      where
          m    = fromIntegral (rows x)    -- For computing average
          h    = sigmoid (x <> theta)     -- Predictions for y
          cost = sumVector (y * log h + (1-y) * log (1-h)) / m
          grad = (1/m) `scale` (y - h) <> x
  
  -- |Cost function and derivative for regularized logistic regression. This is
  --  maximized when fitting parameters for the regression.
  lrLogLikRegularized :: Vector Double            -- targets (y)
                      -> Matrix Double            -- design matrix (x)
                      -> Bool                     -- is first col all ones?
                      -> Double                   -- regularization const (lambda)
                      -> Vector Double            -- coefficients (theta)
                      -> (Double, Vector Double)  -- (value, derivative)
  lrLogLikRegularized y x useConst lambda theta = (cost, grad)
      where
          m      = fromIntegral (rows x)
          (c,g)  = lrLogLikelihood y x theta
          theta' = if useConst then join [0, dropVector 1 theta] else theta
          cost   = c - (lambda / (2 * m)) * norm theta' ^ 2
          grad   = g - (lambda / m) `scale` theta'
  
  -------------
  -- Testing --
  -------------
  
  test n k lambda = do
      x <- randn n k                  -- design matrix
      e <- flatten `fmap` randn n 1   -- errors
      let theta = fromList $ 1 : replicate (k-1) 0
          h = sigmoid $ x <> theta + e
          y = (\i -> if i > 0.5 then 1 else 0) `mapVector` h
          theta_est1 = lr y x
          theta_est2 = lrRegularized y x False lambda
      --putStrLn $ "[y, h, x]"
      --disp 3 $ takeRows 10 $ fromColumns [y, h] ! x
      --putStrLn $ "[y, h, ypred]"
      --disp 3 $ takeRows 10 $ fromColumns [y, h, sigmoid $ x <> theta_est1]
      putStrLn $ "         Number of observations: " ++ show n
      putStrLn $ "           Number of regressors: " ++ show k
      putStrLn $ "                   Actual theta: " ++ show theta
      putStrLn $ "Estimated theta (unregularized): " ++ show theta_est1
      putStrLn $ "  Estimated theta (regularized): " ++ show theta_est2
  
  