/*
 *  Draco Econometrics
 *  Copyright (C) 2006-2007 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/>.
 */

package com.approximatrix.draco.fit;


import java.util.List;

import org.netlib.blas.DAXPY;
import org.netlib.blas.DGEMV;
import org.netlib.blas.DSCAL;

import com.approximatrix.draco.datamodel.StatTableModel;
import com.approximatrix.draco.fit.statistics.ExtendedStatistic;
import com.approximatrix.draco.io.XMLio.AnalysisLoadStore;
import com.approximatrix.draco.math.MatrixSupport;
import com.approximatrix.draco.math.VectorSupport;
import com.approximatrix.draco.utility.ResultsUtilities;

/** Least-Squares with Autocorrelation Solver
 * 
 * An iterative class for performing autocorrelation-corrected least-squares
 * fits of data.  Routine relies heavily on the standard RunLS class for most
 * calculations.  
 * 
 * @author Jeffrey Armstrong
 *
 */
public class RunLSAutocorrelation implements LSIterative {
	
	public static final int CO_METHOD = 1;
	public static final int HL_METHOD = 2;
	
	private int ac_method = 1;
	private int ac_order = 1;
	
	private double rho = 0.0;
	private double last_rho = 0.0;
	
	private String[] independent_vars = null;
	private String dependent_var = null;

	private boolean calculate_constant = true;	
	private int variance_type = RunLS.VAR_OLS;

	// For iterative procedures
	private int iteration_count = 0;
	private double relaxation = 1.0;
	
	private StatTableModel original = null;
	
	private RunLS first_iteration = null;
	
	// Results variables
	private double[] coef = null;
	private double[] residuals = null;
	private double[][] covariance_matrix = null;
	
	// Data that wasn't actually built here
	private double[][] X_matrix = null;
	
	private String results = null;
	private boolean success = true;
	
	/** The list of extended statistics to compute */
	private List<ExtendedStatistic> estats = null;

	/** Basic object constructor that actually performs assignments of private
	 * variables, but does not initiate any calculations.  This constructor
	 * should only be called by the more complex constructors provided.
	 * 
	 * @param indep vector of Strings representing the independent variable names
	 * @param dep String containing the name of the independent variable
	 * @param passed_model the StatTableModel from which all data is taken
	 * @param in_variance_type the type of variance to use (see RunLS)
	 * @param calc_constant boolean flage indicating whether the constant term should be included in the fit
	 * @param method the method of autocorrelation-solving to use
	 * @param order the order of autocorrelation
	 */
	public RunLSAutocorrelation(String[] indep, String dep,
	                            StatTableModel passed_model,
	                            int in_variance_type,
	                            boolean calc_constant,
	                            int method, int order) {

		original = passed_model;
		independent_vars = indep;
		dependent_var = dep;
		variance_type = in_variance_type;
		
		ac_method = method;
		ac_order = order;
		calculate_constant = calc_constant;
		rho = 0;
		last_rho = Double.MAX_VALUE;
	}
	
	/** Constructor for starting a new AC-corrected least-squares fitting
	 * procedure.  This routine takes the input parameters and passes them
	 * to a simple assignment constructor before proceeding to the actual
	 * fitting.   
	 * 
	 * @param indep vector of Strings representing the independent variable names
	 * @param dep String containing the name of the independent variable
	 * @param passed_model the StatTableModel from which all data is taken
	 * @param in_variance_type the type of variance to use (see RunLS)
	 * @param calc_constant boolean flage indicating whether the constant term should be included in the fit
	 * @param method the method of autocorrelation-solving to use
	 * @param order the order of autocorrelation
	 * @param relaxation the relaxation factor used during updates (actually not impleented right now...)
	 */ 
	public RunLSAutocorrelation(String[] indep, String dep, 
		     StatTableModel passed_model,
			 int in_variance_type,
			 boolean calc_constant,
			 int method, int order,
			 double relaxation) {
		
		this(indep,dep,passed_model,in_variance_type,calc_constant,
		     method, order);
		
		this.relaxation = relaxation;
		
	}
	
	/** Constructor for subsequent iterations that builds the new object from
	 * the previous and automatically increments the iteration counter upon
	 * creation.  This constructor calls the simple assignment constructor
	 * to set all variables based on the previous iteration.  The fit is 
	 * performed by this constructor.
	 * 
	 * @param last_iteration the last RunLSAutocorrelation iteration performed
	 */
	public RunLSAutocorrelation(RunLSAutocorrelation last_iteration) {
		this(last_iteration.getIndependentVariables(),
		     last_iteration.getDependentVariable(),
		     last_iteration.getModel(),
		     last_iteration.getVarianceType(),
		     last_iteration.getIncludeConstant(),
		     last_iteration.getACMethod(),
		     last_iteration.getACOrder());
		
		iteration_count = last_iteration.getCurrentIteration() + 1;
		residuals = last_iteration.getResiduals();
		last_rho = last_iteration.getRho();
		first_iteration = last_iteration.getFirstIteration();
		this.setExtendedStatistics(last_iteration.getExtendedStatisticsWithoutCompute());
	}
	
	/** Generates an estimate of the AC-coefficient based on residuals from the
	 * previous fitting iteration.  Uses a Cochran-Orcutt estimate of rho.
	 * 
	 * @param residuals a vector containing the latest known residuals
	 * @return the AC-coefficient, rho
	 */
	private double calculate_rho_estimate(double[] residuals) {
		double numerator = 0.0;
		double denominator = 0.0;
		
		for(int i=ac_order;i<residuals.length;i++) {
			numerator = numerator + residuals[i]*residuals[i-ac_order];
			denominator = denominator + residuals[i-ac_order]*residuals[i-ac_order];
		}
		
		return numerator/denominator;
	}
	
	/** Performs the actual fitting procedure using the Cochran-Orcutt method.
	 * The first iteration merely performs a least squares fit.  Subsequent
	 * iterations use an estimate of the AC coefficient, rho, to generate a new
	 * model containing AC-corrected variables, performing a least squares fit
	 * on the new model.  The success flag is set internally in the routine;
	 * it returns nothing.
	 *
	 */
	private void perform_co_fit() {

		// For first passes
		if(iteration_count == 0) {
			// Generate the first pass results
			first_iteration = new RunLS(independent_vars, dependent_var,
					original,variance_type,calculate_constant);
			first_iteration.perform();
			
			success = first_iteration.isSuccessful();

			coef = first_iteration.getCoef();

			// Independent data for building X matrices
			X_matrix = first_iteration.getXMatrix();
			covariance_matrix = first_iteration.getCovarianceMatrix();
			
			// Retrieve the residuals
			residuals = first_iteration.getResiduals();

		} else {
			
			// Generate temporary column names
			String dependent_name = dependent_var+"_AC";
			String[] independent_names = new String[independent_vars.length];
			for(int i=0;i<independent_vars.length;i++) {
				independent_names[i] = independent_vars[i]+"_AC";
			}

			// Calculate rho
			rho = calculate_rho_estimate(residuals);			

			// Set up new columns
			StatTableModel iteration = new StatTableModel(original);

			// Build a new equation
			String dependent_equation = build_equation(dependent_var,ac_order,rho);
			String[] independent_equations = new String[independent_vars.length];
			for(int i=0;i<independent_vars.length;i++) {
				independent_equations[i] = build_equation(independent_vars[i],ac_order,rho);
			}

			// Add the new columns
			iteration.addColumn(dependent_equation,dependent_name);
			for(int i=0;i<independent_vars.length;i++) {
				iteration.addColumn(independent_equations[i],independent_names[i]);
			}

			// Evaluate a least-squares fit
			RunLS acfit = new RunLS(independent_names,
					dependent_name, 
					iteration,
					variance_type, 
					calculate_constant);

			acfit.perform();
			
			success = acfit.isSuccessful();

			coef = acfit.getCoef();

			// Independent data for building X matrices
			X_matrix = acfit.getXMatrix();
			covariance_matrix = acfit.getCovarianceMatrix();

			// Compute new residuals
			residuals = generate_residuals(first_iteration);

		}
		
	}
	
	/** Internal function used to construct an equation describing new AC-corrected
	 * columns in the model to be used for fitting.
	 * 
	 * @param variable name of the variable of interest
	 * @param order the order of AC
	 * @param rho the current estimate of AC coefficient
	 * @return an equation ready to compute a new AC-corrected variable
	 */
	private String build_equation(String variable, int order, double rho) {
		StringBuffer eqn = new StringBuffer();
		
		eqn.append(variable);
		eqn.append("-");
		eqn.append(new Double(rho).toString());
		eqn.append("*");
		eqn.append("back(");
		eqn.append(variable);
		eqn.append(",");
		eqn.append(new Integer(order).intValue());
		eqn.append(")");
		
		return eqn.toString();
	}

	/** Generates the residuals for this iteration from the original model and
	 * the freshly computed coefficients.
	 * 
	 * @param original
	 * @return a numeric array of residuals
	 */
	private double[] generate_residuals(RunLS original) {
		double[] return_val = (double[])original.getDependentData().clone();
		DGEMV.DGEMV("N",original.getDependentData().length,
		            coef.length,-1.0,original.getXMatrix(),
		            coef,1,1.0,return_val,1);
		return return_val;
	}
	
	/** Provides results of the fit in String form for display
	 *
	 * @return text outlining the results of the fit
	 */
	public String getResults() {
		if(results == null) {
			StringBuffer text = null;
			if(!isSuccessful()) {
				text = LSFitResultsGen.GenerateFailureResults(this,null,"Results of the Multiple Regression Model\nAutocorrelation present\n");
			} else {
				text = LSFitResultsGen.GenerateResults(this,"Results of the Multiple Regression Model\nAutocorrelation present\n");
				text.append("Autocorrelation Coeficient (");
				text.append(ac_order);
				text.append(" Order): ");
				text.append(ResultsUtilities.OutNumber(rho));
				text.append("\n");
			}
			results = text.toString();
		}
		return results;
	}
	
	/** Returns the Covariance matrix resulting from the fit
	 *
	 * @return a square array of doubles containing the covariance matrix
	 */
	public double[][] getCovarianceMatrix() {
		return covariance_matrix;
	}

	/** Returns the residuals from the fit
	 *
	 * @return a vector containing the residuals
	 */
	public double[] getResiduals() {
		return residuals;
	}
	
	/** Returns the sum of the squares of the residuals
	 *
	 * @return the sum of the squares of the residuals
	 */
	public double getResidualSumSqared() {
		return VectorSupport.VectorElementSum(VectorSupport.MultiplyElements(getResiduals(),getResiduals()));
	}
	
	/** Returns the total sum of the squares
	 *
	 * @return the sum of the squares
	 */
	public double getTotalSumSqared() {
		double mean = VectorSupport.VectorElementSum(first_iteration.getDependentData())/
	       (new Double(first_iteration.getDependentData().length).doubleValue());
		double[] ones = VectorSupport.VectorOfOnes(first_iteration.getDependentData().length);

		DSCAL.DSCAL(ones.length,-1.0*mean,ones,1);
		DAXPY.DAXPY(ones.length,1.0,first_iteration.getDependentData(),1,ones,1);
		ones = VectorSupport.MultiplyElements(ones,ones);
		return VectorSupport.VectorElementSum(ones);
	}
	
	/** Returns the coefficients from the fit
	 *
	 * @return a vector containing the coefficients from the fit
	 */
	public double[] getCoef() {
		return coef;
	}
	
	/** Returns the labels to be used for generating results
	 * 
	 * @return a vector of Strings representing the labels for results
	 */
	public String[] getCoefLabels() {
		String[] vars = getIndependentVariables();
		String[] return_val = null;
		int start = 0;
		if(calculate_constant) {
			return_val = new String[vars.length + 1];
			return_val[0] = "Constant";
			start = 1;
		} else {
			return_val = new String[vars.length];
		}
		
		for(int i=0;i<vars.length;i++) {
			StringBuffer sb = new StringBuffer("Coef ");
			sb.append(i);
			sb.append(" (");
			sb.append(vars[i]);
			sb.append(")");
			return_val[i+start] = sb.toString();
		}
		
		return return_val;
	}
	
	/** Returns the variance of each coefficient
	 * 
	 * @return a vector containing the variances of the coefficients
	 */
	public double[] getCoefVariances() {
		return MatrixSupport.GetDiagonal(covariance_matrix);
	}
	
	/** Returns the names of the independent variables
	 * 
	 * @return an array of Strings contianing the variable names
	 */
	public String[] getIndependentVariables() {
		return independent_vars;
	}
	
	/** Returns the name of the dependent variable
	 *
	 * @return a String containing the name of the dependent variable
	 */
	public String getDependentVariable() {
		return dependent_var;
	}
	
	/** Returns true if fitting was successful on this iteration
	 * 
	 * @return true if iteration was successful, false otherwise
	 */
	public boolean isSuccessful() {
		return success;
	}

	/** Returns a vector of Number objects representing estimates of the
	 * dependent variable based on the results of the curve fit suitable for
	 * inclusion in a StatTableModel
	 * 
	 * @returnan array of Numbers containing dependent estimates, where applicable
	 */
	public Number[] getFittedData() {
		if(original != null) {
			return LSFitResultsGen.GetEstimates(this, original);
		} else {
			return null;
		}
	}
	
	/** Returns a vector of Double objects representing residuals suitable for
	 * inclusion into a StatTableModel
	 * 
	 * @return an array of Doubles containing residuals, where applicable
	 */
	public Double[] getResidualData() {
		if(original != null) {
			return LSFitResultsGen.GetResidualsFromEstimates(this,original);
		} else {
			return null;
		}
	}
	
	/** Returns the original model used in the fit
	 * 
	 * @return the unmodified StatTableModel
	 */
	public StatTableModel getModel() {
		return original;
	}
	
	// LSIterative methods
	/** Returns the current iteration number
	 * 
	 * @return the iterative step
	 */
	public int getCurrentIteration() {
		return iteration_count;
	}
	
	/** Returns a measure of change for a single iteration that can
	 * be used for tolerance limiting
	 * 
	 * @return a measure of the amount of change ocmputed in this step
	 */
	public double getIterationChange() {
		return Math.abs(rho-last_rho);
	}
	
	/** Returns a new LSIterative object representing the next iteration
	 * 
	 * @return the next iteration
	 */
	public LSIterative getNextIteration() {
		return new RunLSAutocorrelation(this);
	}
	
	/** Returns the type of variance to calculate
	 * 
	 * @return the variance type based on constants in RunLS
	 */
	protected int getVarianceType() {
		return variance_type;
	}
	
	/** Returns whether or not to include the constant during the fit
	 * 
	 * @return true if a constant should be estimated
	 */
	protected boolean getIncludeConstant() {
		return calculate_constant;
	}
	
	/** Returns the method of fitting to use based on the static constants in
	 * this class
	 * 
	 * @return the method to use
	 */
	protected int getACMethod() {
		return ac_method;
	}
	
	/** Returns the order of autocorrelation being used
	 * 
	 * @return the order of autocorrelation
	 */
	protected int getACOrder() {
		return ac_order;
	}
	
	/** Returns the autocorrellation coefficient used in this iteration
	 * 
	 * @return the rho used in this iteration
	 */
	protected double getRho() {
		return rho;
	}

	/** Returns the first estimate (without AC) least squares fit, which
	 * is used to store basic data needed for calculations
	 * 
	 * @return RunLS object representing the first estimate of solutions
	 */
	protected RunLS getFirstIteration() {
		return first_iteration;
	}

	/** Sets all data in the fit to the information loaded from
	 * a file
	 * 
	 * @param data a valid AnalysisLoadStore object that has been properly filled
	 * @param model the model with which this fit is associated
	 */
	public void setLoadedAnalysis(AnalysisLoadStore data, StatTableModel model) {
		this.original = model;
		
		this.coef = data.getCoef();
		this.covariance_matrix = data.getCovarianceMatrix();
		this.dependent_var = data.getDependentVariable();
		this.independent_vars = data.getIndependentVariables();
		this.residuals = data.getResiduals();
		this.results = data.getResults();
		
		// Include a constant?
		this.calculate_constant = (this.covariance_matrix.length > this.independent_vars.length);
		
		this.success = true;
	}

	public void perform() {
		this.perform_co_fit();
	}
	
	/** This method computes and returns the extended statistics
	 * 
	 * @see com.approximatrix.draco.fit.LSFit#getExtendedStatistics()
	 */
	public List<ExtendedStatistic> getExtendedStatistics() {
		if(estats == null) return null;
		for(ExtendedStatistic stat : estats) {
			stat.setAnalysis(this);
			stat.compute();
		}
		return estats;
	}

	/** Sets the list of extended statistics
	 * @see com.approximatrix.draco.fit.LSFit#setExtendedStatistics(java.util.List)
	 */
	public void setExtendedStatistics(List<ExtendedStatistic> statistics) {
		estats = statistics;
	}

	/** Returns the current list of extended statistics without performing computations
	 * @see com.approximatrix.draco.fit.LSIterative#getExtendedStatisticsWithoutCompute()
	 */
	public List<ExtendedStatistic> getExtendedStatisticsWithoutCompute() {
		return estats;
	}
}
