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


// For equation parsing
import com.approximatrix.draco.datamodel.StatTableModel;
import com.approximatrix.draco.fit.estimation.ArbitraryNonLinearEstimationProblem;
import com.approximatrix.draco.fit.statistics.ExtendedStatistic;
import com.approximatrix.draco.io.XMLio.AnalysisLoadStore;
import com.approximatrix.draco.math.MatrixSupport;
import com.graphbuilder.math.ExpressionParseException;

// LAPACK/BLAS stuff
import org.apache.commons.math.estimation.EstimatedParameter;
import org.apache.commons.math.estimation.EstimationException;
import org.apache.commons.math.estimation.LevenbergMarquardtEstimator;
import org.netlib.blas.DAXPY;
import org.netlib.blas.DDOT;

public class RunNLLS implements LSFit {
	
	private StatTableModel model = null;
	String dependent_var = null;
	String[] parameters = null;
	double[] estimates = null;
	
	String equation = null;
	String[] derivative_eqns = null;
	
	String results = null;
	boolean success = true;
	
	int[] use_row = null;
	
	double relaxation = 1.0;
	
	String failure_message = null;
	
	double[][] covariance_matrix = null;
	
	private int iteration = 0;
	private double change = Double.POSITIVE_INFINITY;
	
	/** The list of extended statistics to compute */
	private List<ExtendedStatistic> estats = null;
	
	/** The estimation problem */
	ArbitraryNonLinearEstimationProblem eProblem = null;
	
	/** The L-M estimator */
	LevenbergMarquardtEstimator estimator = null;
	
	/** Dummy constructor */
	public RunNLLS() {
		success = false;
		
		failure_message = "No data provided...";
	}
	
	/** Constructor for a RunNLLS iteration that accepts pre-computed derivatives.
	 * The constructor is useful for special cases of NLLS where the derivatives
	 * are analytically pre-determined, such as probit and logit fits.
	 * 
	 * @param model the StatTableModel containing all data to use
	 * @param dependent the dependent variable name
	 * @param parser the pre-existing DJep parser
	 * @param equation the equation of the base function
	 * @param derivatives the precomputed derivatives wrt the paramters
	 * @param parameters the parameter names
	 * @param estimates current estimate of the parameters
	 */
	public RunNLLS(StatTableModel model, 
			       String dependent, 
			       String equation,
			       String[] derivatives,
			       String[] parameters, 
			       double[] estimates) {
		
		this.model = model;
		this.parameters = parameters;
		this.estimates = estimates;
		this.dependent_var = dependent;
		
		this.equation = equation;
		this.derivative_eqns = derivatives;
		
		// Determine the rows to use
		determine_allowable_rows();
		if(use_row == null) {
			success = false;
			failure_message = "Unable to determine usable rows from the data model.";
		} else if(use_row.length<parameters.length) { 
			success = false;
			failure_message = "Too few degrees of freedom to continue with regression.";
		}
		
	}
	
	/** Wrapper constructor that allows the relaxation factor to be specified
	 * for an initial RunNLLS computation.  This constructor should be used
	 * on the first iteration.  Future iterations should use another
	 * constructor.
	 * 
	 * @param model the data model to use
	 * @param dependent the dependent variable for fitting 
	 * @param equation the equation to which to fit
	 * @param parameters parameters in the equation
	 * @param estimates estimates corresponding to each parameter
	 * @param relaxation the factor to multiply the estimate updates by before the update
	 */
	public RunNLLS(StatTableModel model, String dependent, String equation, String[] parameters, double[] estimates, double relaxation) {
		this(model, dependent, equation, parameters, estimates);
		this.relaxation = relaxation;
	}

	
	/** Constructor for starting the initial RunNLLS computation.  This should
	 * be used on the first iteration.  Future iterations would be better off
	 * with another constructor.
	 * 
	 * @param model the data model to use
	 * @param dependent the dependent variable for fitting 
	 * @param equation the equation to which to fit
	 * @param parameters parameters in the equation
	 * @param estimates estimates corresponding to each parameter
	 */
	public RunNLLS(StatTableModel model, String dependent, String equation, String[] parameters, double[] estimates) {
		this.model = model;
		this.parameters = parameters;
		this.estimates = estimates;
		this.dependent_var = dependent;
		
		this.equation = equation;
		
		// Determine the rows to use
		determine_allowable_rows();
		if(use_row == null) {
			success = false;
			failure_message = "Unable to determine usable rows from the data model.";
		} else if(use_row.length<parameters.length) { 
			success = false;
			failure_message = "Too few degrees of freedom to continue with regression.";
		}
	}
	
	/** Constructor that builds a RunNLLS object from a previous iteration.  This
	 * constructor is provided to simplify the iteration process needed to properly
	 * perform the non-linear least squares fit
	 * 
	 * @param last_iteration a valid RunNLLS object
	 */
	public RunNLLS(RunNLLS last_iteration) {
	
		// Get the last iteration's data
		model = last_iteration.getModel();
		parameters = last_iteration.getCoefLabels();
		estimates = last_iteration.getCoef();
		equation = last_iteration.getEquation();
		dependent_var = last_iteration.getDependentVariable();
		
		// Parser and derivatives
		derivative_eqns = last_iteration.getDerivatives();
		
		// Get the last step's success rating
		success = last_iteration.isSuccessful();
		failure_message = last_iteration.getFailureMessage();
		
		relaxation = last_iteration.getRelaxationFactor();
		
		// Update the iteration
		iteration = last_iteration.getCurrentIteration()+1;
		
		// Determine the rows to use
		determine_allowable_rows();
		if(use_row == null) {
			success = false;
			failure_message = "Unable to determine usable rows from the data model.";
		} else if(use_row.length<parameters.length) { 
			success = false;
			failure_message = "Too few degrees of freedom to continue with regression.";
		}
		
		this.setExtendedStatistics(last_iteration.getExtendedStatisticsWithoutCompute());

	}
	
	public String getResults() {
		if(results == null) {
			StringBuffer text = null;
			if(!success) {
				text = LSFitResultsGen.GenerateFailureResults(this,getFailureMessage(),"Results of the Non-Linear Least Squares Regression Model\n");
			} else {
				text = LSFitResultsGen.GenerateResults(this,"Results of the Non-Linear Least Squares Regression Model\n");
				//print_special_options(text);
				//checkExtendedOutput(text);
			}
			results = text.toString();
		}
		return results;
	}

	/** Computes the covariance matrix if not yet done and returns the
	 * matrix as an array of doubles
	 * 
	 * @return an array containing the covariance matrix
	 */
	public double[][] getCovarianceMatrix() {
		if(!success) return null;
		return covariance_matrix;
	}

	/** Returns the difference between the data and the estimated
	 * dependent values based on the current value of the parameter
	 * estimates.
	 * 
	 * @return a vector containing the residuals
	 */
	public double[] getResiduals() {
		if(!success) return null;
		
		// Get the current y_hat values
		double[] b = evaluate_eqn();
		if(!success) {
			failure_message = "Could not evaluate base expression.";
			return null;
		}
		
		// calculate db
		double[] db = get_allowable_dependents();
		DAXPY.DAXPY(b.length,-1.0,b,1,db,1);
		
		return db;
	}

	/** Returns the sum squared of the residuals
	 * 
	 * @return a double value representing R^2
	 */
	public double getResidualSumSqared() {
		double[] residuals = getResiduals();
		if(residuals == null) return 0;
		
		return DDOT.DDOT(residuals.length,residuals,1,residuals,1);
	}

	/** Returns the total sum squared
	 * 
	 * @return a double value representing the TSS
	 */
	public double getTotalSumSqared() {
		double[] dep = get_allowable_dependents();
		if(dep == null) return 0;
		
		return DDOT.DDOT(dep.length,dep,1,dep,1);
	}

	/** Returns the estimates for the parameters
	 * 
	 * @return an array of estimates
	 */
	public double[] getCoef() {
		return estimates;
	}
	
	/** Returns the labels of each parameter
	 * 
	 * @return an array of strings
	 */
	public String[] getCoefLabels() {
		return parameters;
	}
	
	/** Returns the derivative nodes already computed by JEP
	 * 
	 * @return a vector of Nodes 
	 */
	public String[] getDerivatives() {
		return derivative_eqns;
	}
	
	/** Returns the equation being fitted
	 * 
	 * @return the equation in String form
	 */
	public String getEquation() {
		return equation;
	}
	
	/** Returns the variances of the computed parameters
	 * 
	 * @return a vector containing the variance of each estimated parameter
	 */
	public double[] getCoefVariances() {
		return MatrixSupport.GetDiagonal(getCovarianceMatrix());
	}

	/** Returns a list of independent variables.  For the time being,
	 * this list includes ALL variables in the StatTableModel with
	 * the exception of the dependent variable.
	 * 
	 * @return a vector of Strings containing the names of variables
	 */
	public String[] getIndependentVariables() {
		String[] vars = model.getColumnNames();
		String[] return_val = new String[model.getColumnCount()-1];
		int j=0;
		for(int i=0;i<vars.length;i++) {
			if(vars[i].compareTo(dependent_var)!=0) {
				return_val[j] = vars[i];
				j++;
			}
		}
		
		return return_val;
	}
	
	/** Returns the relaxation factor used when updating the estimates
	 * on a single iteration
	 * 
	 * @return the relaxation factor
	 */
	public double getRelaxationFactor() {
		return relaxation;
	}

	/** Returns the dependent variable
	 * 
	 * @return the name of the dependent variable
	 */
	public String getDependentVariable() {
		return dependent_var;
	}

	/** Returns whether the iteration step was successful
	 * 
	 * @return true if completed successfully, false otherwise
	 */
	public boolean isSuccessful() {
		return success;
	}
	
	/** Returns estimates of the dependent data generated by the model
	 * appropriate for inclusion in a StatTableModel
	 * 
	 * @return a vector of Doubles containing the estimated values or null
	 */
	public Double[] getFittedData() {
		Double[] return_val = new Double[model.getRowCount()];
		double[] y_hat = evaluate_eqn();
		for(int i=0;i<use_row.length;i++) {
			return_val[use_row[i]] = new Double(y_hat[i]);
		}
		return return_val;
	}

	/** Returns a vector of Double values appropriate for a StatTableModel
	 * containing all the residuals.
	 * 
	 * @return a vector of Doubles containing residual values or null
	 */
	public Double[] getResidualData() {
		Double[] return_val = new Double[model.getRowCount()];
		double[] residuals = getResiduals();
		for(int i=0;i<use_row.length;i++) {
			return_val[use_row[i]] = new Double(residuals[i]);
		}
		return return_val;
	}

	/** Returns the current data model being stored in the fitting object
	 * 
	 * @return a StatTableModel
	 */
	public StatTableModel getModel() {
		return model;
	}

	private void constructProblem() {
		double[][] indep_data = this.getModel().getDataDouble(this.use_row, this.getIndependentVariables());
		
		try {
			eProblem = 
				new ArbitraryNonLinearEstimationProblem(this.equation,
						indep_data,
						this.getIndependentVariables(),
						this.get_allowable_dependents(),
						this.getDependentVariable(),
						this.getModel().getConstants());
			eProblem.addParameters(this.getCoefLabels(), this.estimates);
		} catch(ExpressionParseException epe) {
			success = false;
			failure_message = "Expression failed to parse properly: " + epe.getDescription();
		}
		
	}
	
	/** Performs a single NLLS regression step and updates all estimates accordingly
	 * 
	 * @return true if successful, false otherwise
	 */
	private boolean regress() {
		
		this.constructProblem();
		
		if(!success)
			return false;
		
		estimator = new LevenbergMarquardtEstimator();
		
		try {
			estimator.estimate(eProblem);
			covariance_matrix = estimator.getCovariances(this.eProblem);
			covariance_matrix = MatrixSupport.ScalarTimesMatrix(getResidualVariance(),covariance_matrix,parameters.length,parameters.length);
			success = true;
		} catch(EstimationException ee) {
			failure_message = "Error executing the Levenberg-Marquardt Estimator - Convergence Failed\n"+
							  ee.getLocalizedMessage();
			success = false;
		}
		
		if(success)
			this.assignEstimates();
		
		return success;
	}
	
	/** Assigns estimates from the eProblem back into this local estimate vector */
	private void assignEstimates() {
		EstimatedParameter[] ep = eProblem.getAllParameters();
		
		for(int i=0;i<Math.min(ep.length,this.estimates.length);i++)
			this.estimates[i] = ep[i].getEstimate();
	}
	
	/** Evaluates the base equation using the variables, constants, and current
	 * estimates of parameters, returning a vector of estimates of the dependent
	 * variable.  If the function fails, the success flag is set to false.
	 * 
	 * @return a vector of estimates of the current equation, or null if failed
	 */
	private double[] evaluate_eqn() {
		
		if(!success) 
			return null;
		else {
			if(eProblem == null)
				this.constructProblem();
			try {
				return eProblem.evaluateAll();
			} catch(NullPointerException npe) {
				success = false;
				failure_message = "An internal error has occured...";
				return null;
			}
		}
	}
	
	/** Simplistic method for determining which rows are acceptable for use in 
	 * the fitting.  
	 *
	 */ 
	private void determine_allowable_rows() {
		String[] to_use = new String[this.getIndependentVariables().length+1];
		for(int i=0;i<to_use.length-1;i++) to_use[i] = this.getIndependentVariables()[i];
		to_use[to_use.length-1] = this.dependent_var;
		use_row = model.getNonNullRows(to_use);
	}
	
	/** Returns a vector of doubles representing the values of the dependent variable
	 * on rows corresponding to the entries in the use_row vector.
	 * 
	 * @return a vector of dependent values
	 */
	private double[] get_allowable_dependents() {
		double[] base_values = model.getColumnDouble(model.getColumnNumber(dependent_var));
		double[] return_val = new double[use_row.length];
		for(int i=0;i<use_row.length;i++) {
			return_val[i] = base_values[use_row[i]];
		}
		return return_val;
	}
	
	/** Returns the variance of the residuals
	 * 
	 * @return double representing the residual variance
	 */
	private double getResidualVariance() {
		double denom = new Integer(use_row.length).doubleValue() -
		    new Integer(parameters.length).doubleValue();
		return getResidualSumSqared()/denom;
	}
	
	/** Returns the current String value of the failure message, or null
	 * if none has occurred.
	 * 
	 * @return the String indicating why failure occured
	 */
	public String getFailureMessage() {
		return failure_message;
	}
	
	/** Returns the current iteration number
	 * 
	 * @return the iterative step
	 */
	public int getCurrentIteration() {
		return iteration;
	}
	
	/** 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 change;
	}
	
	/** Returns a new LSIterative object representing the next iteration
	 * 
	 * @return the next iteration
	 */
	public LSIterative getNextIteration() {
		return (LSIterative)(new RunNLLS(this));
	}
	
	/** 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.model = model;
		
		this.estimates = data.getCoef();
		this.parameters = data.getCoefLabels();
		
		this.covariance_matrix = data.getCovarianceMatrix();
		
		this.dependent_var = data.getDependentVariable();
		
		this.results = data.getResults();
		
		this.equation = data.getEquation();
		
		this.success = true;
			
	}
	
	/** Explicit call to perform the regression (1 iteration only)
	 * 
	 */
	public void perform() {
		// Run regression
		this.success = regress();
	}
	
	/** 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;
	}
}
