#  Draco Econometrics
#  Copyright (C) 2006-2008 Approximatrix, LLC
#  http://www.approximatrix.com
#  
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.

# Some info taken from:
#   http://en.wikipedia.org/wiki/Least_squares#Weighted_least_squares

import java.lang.System
import com.approximatrix.draco.application.plugin.PluginType

require 'pengine/model/regression_plugin'
include Math

class WeightedLeastSquares < Plugin::RegressionPluginSkeleton
  @@NAME = "Weighted Least Squares"
  
  def initialize(model)
    super(model)
    @covariance = nil
    @x = nil
    @failure = nil
  end

  def has_constant?
    @options.get_option("Include Constant Term")
  end
  
  def initialize_columns
    # simply add the columns here in your override
    @selectioncolumns.add WeightedLeastSquares.createColumn('Independent', true, false)
    @selectioncolumns.add WeightedLeastSquares.createColumn('Dependent', false, true)
    @selectioncolumns.add WeightedLeastSquares.createColumn('Weighting', false, false)
  end
  
  def initialize_options
  	  @options.add "Include Constant Term", "boolean", "true"
      @options.add "Ignore Matrix Conditioning", "boolean", "true"
  end
  
  def run
    # Retrieve variables selected
    independent_vars = selections(0)
    dependent_vars = selections(1)
    weight_vars = selections(2)
    
    # Set up the coef names for results
    generate_coef_names independent_vars
    
    # Need to retrieve all data at once
    all_vars = independent_vars+dependent_vars+weight_vars
    data = @model.to_matrix(all_vars)
    
    # Build the coef. matrix
    a = data.column_vectors[0..-3]
    a.insert(0,Vector.ones(data.m)) if has_constant?
    
    a = Matrix.columns(a)
    b = data.column(-2)
    w = data.column(-1)
    
    # Check if we may proceed
    if(a.m <= a.n)
      @failure = "Data does not supply sufficient degrees of freedom."
      return
    end
    
    # Weight diagonal matrix
    aw = Matrix.diagonal(*(w.to_a))
    
    # Create the matrix to invert (A'A)
    atw = a.transpose*aw
    
    atwa = atw*a
     
    atwy = atw*b

    # Check conditioning
    if(atwa.reciprocal_condition < 1.0E-10 and !@options.get_option("Ignore Matrix Conditioning"))
      @failure = "Matrix is poorly conditioned (RCOND=#{atwa.reciprocal_condition})."
    else
      # Invert the A matrix and solve
      atwa = atwa.inverse
      @x = (atwa)*atwy
      
      # Compute residuals (not using the weights)
      b_hat = a*@x
      @residuals = b_hat-b
      
      # Compute the residual variance
      residual_variance = @residuals.sum_squared/(@residuals.length-@x.length)
      
      # Compute covariance
      @covariance = covariance_matrix(a)
      @covariance = Matrix.zeroes(b.length,b.length) if @covariance.nil?

	  # Get the standard errors
	  @stderr = sqrt(@covariance.diagonal)
	  
	  @tscores = []
	  
	  @x.length.times do |i|
	  	@tscores << @x[i]/@stderr[i]
	  end
	  
	  wresiduals = aw*@residuals
	  
	  # Calculate TSS
	  tss = (b-(b.mean*Vector.ones(b.length))).sum_squared
	  wtss = (aw.transpose*b-((aw.transpose*b).mean*Vector.ones(b.length))).sum_squared
	  
	  # Save some statistics for output
	  statistic "Sum Squared of the Residuals (weighted)", wresiduals.sum_squared
	  statistic "Standard Error of the Fit (weighted)", sqrt(wresiduals.sum_squared/(wresiduals.length-@x.length))
	  statistic "R-Squared Value (weighted)", r_squared(wresiduals,wtss)
	  statistic "Adjusted R-Squared Value (weighted)", 1.0-( (1.0*wresiduals.length-1.0*@x.length)/(1.0*wresiduals.length-1.0)*(1.0-r_squared(wresiduals,wtss)) )
	  
	  
	  statistic "Sum Squared of the Residuals (unweighted)", @residuals.sum_squared
	  statistic "Standard Error of the Fit (unweighted)", sqrt(@residuals.sum_squared/(@residuals.length-@x.length))
	  statistic "R-Squared Value (unweighted)", r_squared(@residuals,tss)
	  statistic "Adjusted R-Squared Value (unweighted)", 1.0-( (1.0*@residuals.length-1.0*@x.length)/(1.0*@residuals.length-1.0)*(1.0-r_squared(@residuals,tss)) )
	  
    end
    
  end
  
  private
  
  def r_squared(res,tss)
    1.0-res.sum_squared/tss
  end
  
  def covariance_matrix(a)
  	ata = a.transpose*a
  	if((ata.reciprocal_condition < 1.0E-10 || @options.get_option("Ignore Matrix Conditioning")) && !@residuals.nil?)
  	  @covariance = ata.inverse
  	  residual_variance = @residuals.sum_squared/(@residuals.length-@x.length) 
  	  @covariance = residual_variance*@covariance
  	  @covariance
  	else
  	  nil
  	end
  end
  
end

register(WeightedLeastSquares)
