//:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
/** @author  John Miller, Mustafa Nural
 *  @version 1.6
 *  @date    Tue Apr 18 14:24:14 EDT 2017
 *  @see     LICENSE (MIT style license file).
 *
 *  @title   Model: Lasso Regression (L1 Shrinkage/Regularization)
 */

package scalation.analytics

import scala.collection.mutable.Set
import scala.math.{abs, log, pow, sqrt}
import scala.util.control.Breaks.{break, breakable}

import scalation.linalgebra._
import scalation.math.noDouble
import scalation.minima._
import scalation.plot.{Plot, PlotM}
import scalation.stat.Statistic
import scalation.util.banner

import Fit._
import MatrixTransform._
import RegTechnique._

//:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
/** The `LassoRegression` class supports multiple linear regression.  In this case,
 *  'x' is multi-dimensional [1, x_1, ... x_k].  Fit the parameter vector 'b' in
 *  the regression equation
 *  <p>
 *      y  =  b dot x + e  =  b_0 + b_1 * x_1 + ... b_k * x_k + e
 *  <p>
 *  where 'e' represents the residuals (the part not explained by the model).
 *  @see see.stanford.edu/materials/lsoeldsee263/05-ls.pdf
 *  @param x       the data/input m-by-n matrix
 *  @param y       the response/output m-vector
 *  @param fname_  the feature/variable names
 *  @param hparam  the shrinkage hyper-parameter, lambda (0 => OLS) in the penalty term 'lambda * b dot b'
 */
class LassoRegression (x: MatriD, y: VectoD,
                       fname_ : Strings = null, hparam: HyperParameter = LassoRegression.hp)
      extends PredictorMat (x, y, fname_, hparam)
{
    private val DEBUG  = false                                         // debug flag
    private var lambda = hparam ("lambda")                             // weight to put on regularization

    //::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    /** Return the value of the shrinkage parameter 'lambda'.
     */
    def lambda_ : Double = lambda

    //::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    /** Find an optimal value for the shrinkage parameter 'lambda' using Cross Validation
     *  to minimize 'sse_cv'.  The search starts with the low default value for 'lambda'
     *  doubles it with each iteration, returning the minimum 'lambda' and it corresponding
     *  cross-validated 'sse'.
     */
    def findLambda: (Double, Double) =
    {
        var l      = lambda
        var l_best = l
        var sse    = Double.MaxValue
        for (i <- 0 to 20) {
            LassoRegression.hp("lambda") = l
            val rrg = new LassoRegression (x, y)
            val stats = rrg.crossValidate ()
            val sse2 = stats(Fit.index_sse).mean
            banner (s"LassoRegession with lambda = ${rrg.lambda_} has sse = $sse2")
            if (sse2 < sse) { sse = sse2; l_best = l }
            if (DEBUG) showQofStatTable (stats)
            l *= 2
        } // for
        (l_best, sse)
    } // findLambda

    //::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    /** Compute the sum of squares error + lambda * sum of the magnitude  of coefficients.
     *  This is the objective function to be minimized.
     *  @param yy  the response/output vector
     *  @param b   the vector of coefficients/parameters
     * 
    def f (yy: VectoD)(b: VectoD): Double =
    {
        e = yy - x * b                                                 // calculate the residuals/error
        e dot e + lambda * b.norm1 
    } // f
     */

    //::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    /** Train the predictor by fitting the parameter vector (b-vector) in the
     *  multiple regression equation
     *  <p>
     *      y  =  b dot x + e  =  [b_0, ... b_k] dot [1, x_1 , ... x_k] + e
     *  <p>
     *  regularized by the sum of magnitudes of the coefficients.
     *  @see pdfs.semanticscholar.org/969f/077a3a56105a926a3b0c67077a57f3da3ddf.pdf
     *  @see `scalation.minima.LassoAdmm`
     *  @param x_  the training/full data/input matrix
     *  @param y_  the training/full response/output vector
     */
    def train (x_ : MatriD = x, y_ : VectoD = y): LassoRegression =
    {
//      val g       = f(y_) _
//      val optimer = new CoordinateDescent (g)                        // Coordinate Descent optimizer
//      val optimer = new GradientDescent (g)                          // Gradient Descent optimizer
//      val optimer = new ConjugateGradient (g)                        // Conjugate Gradient optimizer
//      val optimer = new QuasiNewton (g)                              // Quasi-Newton optimizer
//      val optimer = new NelderMeadSimplex (g, x_.dim2)               // Nelder-Mead optimizer
//      val b0 = new VectorD (k+1)                                     // initial guess for coefficient vector
//      b = optimer.solve (b0, 0.5)                                    // find an optimal solution for coefficients

        b = LassoAdmm.solve (x_.asInstanceOf [MatrixD], y_, lambda)    // Alternating Direction Method of Multipliers
        this
    } // train

    //::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    /** Build a sub-model that is restricted to the given columns of the data matrix.
     *  @param x_cols  the columns that the new model is restricted to
     */
    def buildModel (x_cols: MatriD): LassoRegression =
    {
        new LassoRegression (x_cols, y, null, hparam)
    } // buildModel

} // LassoRegression class


//:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
/** The `LassoRegression` companion object provides factory methods for the
 *  `LassoRegression` class.
 */
object LassoRegression extends ModelFactory
{
    /** Base hyper-parameter specification for `LassoRegression`
     */
    val hp  = new HyperParameter; hp += ("lambda", 0.01, 0.01)         // L1 regularization/shrinkage parameter
    val drp = (null, hp)                                               // default remaining parameters

    //:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    /** Create a `LassoRegression` object from a combined data matrix.
     *  @param xy      the combined data matrix
     *  @param fname   the feature/variable names (use null for default)
     *  @param hparam  the hyper-parameters (use null for default)
     */
    def apply (xy: MatriD, fname: Strings = null, hparam: HyperParameter = hp): LassoRegression =
    {
        val n = xy.dim2
        if (n < 2) {
            flaw ("apply", s"dim2 = $n of the 'xy' matrix must be at least 2")
            null
        } else {
            val (x, y) = pullResponse (xy)
            new LassoRegression (x, y, fname, hparam)
       } // if
    } // apply

    //::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    /** Create a `LassoRegression` object from a data matrix and a response vector.
     *  This factory function provides data rescaling.
     *  @see `ModelFactory`
     *  @param x       the data/input m-by-n matrix
     *                     (augment with a first column of ones to include intercept in model)
     *  @param y       the response/output m-vector
     *  @param fname   the feature/variable names (use null for default)
     *  @param hparam  the hyper-parameters (use null for default)
     */
    def apply (x: MatriD, y: VectoD, fname: Strings, hparam: HyperParameter): LassoRegression =
    {
        val n = x.dim2
        if (n < 1) {
            flaw ("apply", s"dim2 = $n of the 'x' matrix must be at least 1")
            null
        } else if (rescale) {                                          // normalize the x matrix
            val (mu_x, sig_x) = (x.mean, stddev (x))
            val xn = normalize (x, (mu_x, sig_x))
            new LassoRegression (xn, y, fname, hparam)
        } else {
            new LassoRegression (x, y, fname, hparam)
       } // if
    } // apply

} // LassoRegression object


//:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
/** The `LassoRegressionTest` object tests `LassoRegression` class using the following
 *  regression equation.
 *  <p>
 *      y  =  b dot x  =  b_0 + b_1*x_1 + b_2*x_2.
 *  <p>
 *  It comapres `LassoRegression` to `Regression`.
 *  @see http://statmaster.sdu.dk/courses/st111/module03/index.html
 *  > runMain scalation.analytics.LassoRegressionTest
 */
object LassoRegressionTest extends App
{
    // 5 data points: constant term, x_1 coordinate, x_2 coordinate
    val x = new MatrixD ((5, 3), 1.0, 36.0,  66.0,                 // 5-by-3 matrix
                                 1.0, 37.0,  68.0,
                                 1.0, 47.0,  64.0,
                                 1.0, 32.0,  53.0,
                                 1.0,  1.0, 101.0)
    val y = VectorD (745.0, 895.0, 442.0, 440.0, 1598.0)
    val z = VectorD (1.0, 20.0, 80.0)

//  println ("model: y = b_0 + b_1*x_1 + b_2*x_2")
    println ("model: y = b₀ + b₁*x₁ + b₂*x₂")
    println ("x = " + x)
    println ("y = " + y)

    banner ("Regression")
    val rg = new Regression (x, y)
    println (rg.analyze ().report)
    println (rg.summary)

    var yp = rg.predict (z)                                        // predict y for one point
    println (s"predict ($z) = $yp")

    banner ("LassoRegression")
    val lrg = new LassoRegression (x, y)
    println (lrg.analyze ().report)
    println (lrg.summary)

    yp = lrg.predict (z)                                           // predict y for one point
    println (s"predict ($z) = $yp")

    val yyp = rg.predict (x)                                       // predict y for several points
    println (s"predict (x) = $yyp")

    new Plot (x.col(1), y, yyp)
    new Plot (x.col(2), y, yyp)

} // LassoRegressionTest object


//:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
/** The `LassoRegressionTest2` object tests `LassoRegression` class using the following
 *  regression equation.
 *  <p>
 *      y  =  b dot x  =  b_1*x1 + b_2*x_2.
 *  <p>
 *  Try non-default value for the 'lambda' hyper-parameter.
 *  > runMain scalation.analytics.LassoRegressionTest2
 */
object LassoRegressionTest2 extends App
{
    import LassoRegression.hp

    println (s"hp = $hp")
    val hp2 = hp.updateReturn ("lambda", 1.0)                         // try different values
    println (s"hp2 = $hp2")

    // 5 data points:            one   x_1    x_2
    val x = new MatrixD ((5, 3), 1.0, 36.0,  66.0,                    // 5-by-3 matrix
                                 1.0, 37.0,  68.0,
                                 1.0, 47.0,  64.0,
                                 1.0, 32.0,  53.0,
                                 1.0,  1.0, 101.0)
    val y = VectorD (745.0, 895.0, 442.0, 440.0, 1598.0)
    val z = VectorD (1.0, 20.0, 80.0)

    println ("x = " + x + "\ny = " + y + "\nz = " + z)

    banner ("LassoRegression")
    val lrg = new LassoRegression (x, y, hparam = hp2)
    println (lrg.analyze ().report)
    println (lrg.summary)

    val yp  = lrg.predict (z)                                         // predict response for z
    println (s"predict ($z) = $yp")

} // LassoRegressionTest2 object


//:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
/** The `LassoRegressionTest3` object tests `LassoRegression` class using the following
 *  regression equation.
 *  <p>
 *      y  =  b dot x  =  b_1*x1 + b_2*x_2.
 *  <p>
 *  Test regression, forward selection and backward elimination.
 *  > runMain scalation.analytics.LassoRegressionTest3
 */
object LassoRegressionTest3 extends App
{
    // 5 data points:            one   x_1    x_2
    val x = new MatrixD ((5, 3), 1.0, 36.0,  66.0,                    // 5-by-3 matrix
                                 1.0, 37.0,  68.0,
                                 1.0, 47.0,  64.0,
                                 1.0, 32.0,  53.0,
                                 1.0,  1.0, 101.0)
    val y = VectorD (745.0, 895.0, 442.0, 440.0, 1598.0)
    val z = VectorD (1.0, 20.0, 80.0)

    println ("x = " + x + "\ny = " + y + "\nz = " + z)

    banner ("LassoRegression")
    val lrg = new LassoRegression (x, y)
    println (lrg.analyze ().report)

    banner ("Forward Selection Test")
    lrg.forwardSelAll (cross = false)

    banner ("Backward Elimination Test")
    lrg.backwardElimAll (cross = false)

} // LassoRegressionTest3 object


//:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
/** The `LassoRegressionTest4` object tests the `LassoRegression` class using the AutoMPG
 *  dataset.  It illustrates using the `Relation` class for reading the data
 *  from a .csv file "auto-mpg.csv".  Assumes no missing values.
 *  It also combines feature selection with cross-validation and plots
 *  R^2, R^2 Bar and R^2 cv vs. the instance index.
 *  > runMain scalation.analytics.LassoRegressionTest4
 */
object LassoRegressionTest4 extends App
{
    import scalation.columnar_db.Relation

    banner ("auto_mpg relation")
    val auto_tab = Relation (BASE_DIR + "auto-mpg.csv", "auto_mpg", null, -1)
    auto_tab.show ()

    banner ("auto_mpg (x, y) dataset")
    val (x, y) = auto_tab.toMatriDD (1 to 6, 0)
    println (s"x = $x")
    println (s"y = $y")

    banner ("auto_mpg regression")
    val lrg = new LassoRegression (x, y)
    println (lrg.analyze ().report)
    println (lrg.summary)
    val n = x.dim2                                                 // number of parameters/variables

    banner ("Forward Selection Test")
    val (cols, rSq) = lrg.forwardSelAll ()                         // R^2, R^2 Bar, R^2 cv
    val k = cols.size
    val t = VectorD.range (1, k)                                   // instance index
    new PlotM (t, rSq.t, Array ("R^2", "R^2 bar", "R^2 cv"),
               "R^2 vs n for LassoRegression", lines = true)
    println (s"rSq = $rSq")

} // LassoRegressionTest4 object


//:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
/** The `LassoRegressionTest5` object tests the `LassoRegression` class using the AutoMPG
 *  dataset.  It illustrates using the `Relation` class for reading the data
 *  from a .csv file "auto-mpg.csv".  Assumes no missing values.
 *  It also uses the 'findLambda' method to search for a shrinkage parameter
 *  that roughly mininizes 'sse_cv'.
 *  > runMain scalation.analytics.LassoRegressionTest5
 */
object LassoRegressionTest5 extends App
{
    import scalation.columnar_db.Relation

    banner ("auto_mpg relation")
    val auto_tab = Relation (BASE_DIR + "auto-mpg.csv", "auto_mpg", null, -1)
    auto_tab.show ()

    banner ("auto_mpg (x, y) dataset")
    val (x, y) = auto_tab.toMatriDD (1 to 6, 0)

    banner ("auto_mpg Lasso Regression")
    val lrg = new LassoRegression (x, y)
    println (lrg.analyze ().report)
    println (lrg.summary)

    println (s"best (lambda, sse) = ${lrg.findLambda}")

} // LassoRegressionTest5 object