/*
 *  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.ArrayList;
import java.util.List;


// For equation parsing
import com.approximatrix.draco.datamodel.Constant;
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.graphbuilder.math.Expression;
import com.graphbuilder.math.ExpressionTree;
import com.graphbuilder.math.ExpressionParseException;
import com.graphbuilder.math.VarMap;
import com.graphbuilder.math.FuncMap;

// LAPACK/BLAS stuff
import org.netlib.blas.DAXPY;
import org.netlib.blas.DGEMM;
import org.netlib.blas.DGEMV;
import org.netlib.lapack.DGESV;
import org.netlib.lapack.DGETRF;
import org.netlib.lapack.DGETRI;
import org.netlib.util.intW;
import org.netlib.blas.DDOT;
import org.netlib.blas.DASUM;

public class RunNLLS implements LSIterative {

	private static final double INITIAL_STEP = 1.0E-10;
	
	private static final int MAX_TRIES_FOR_NUM_DERIVATIVE = 10000;
	
	private static final double RELAXATION_MULTIPLIER = 1.0;
	
	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;
	
	/** 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(covariance_matrix != null) return covariance_matrix;
		if(!success) return null;
		
		// The result matrix
		double[][] ls_coef = MatrixSupport.Identity(parameters.length);
		
		// The independent variable matrices
		double[][] A = build_A();				
		
		// Matrix multiply to determine (X'X)
		DGEMM.DGEMM("T","N",parameters.length,parameters.length,use_row.length,
				    1.0,A,A,0.0,ls_coef);
				
		int[] ipiv = new int[parameters.length];
		double[] work = new double[parameters.length];
		intW result = new intW(0);
		
		// Together, these invert a matrix (temp)
		DGETRF.DGETRF(parameters.length,parameters.length,ls_coef,ipiv,result);
		DGETRI.DGETRI(parameters.length,ls_coef,ipiv,work,work.length,result);
		
		// Multiply by residual variance
		MatrixSupport.ScalarTimesMatrix(getResidualVariance(),ls_coef,parameters.length,parameters.length);
		
		// Do we need white estimates???
		//if(variance_type == VAR_OLS) {
		//	MatrixSupport.ScalarTimesMatrix(residual_variance(),temp,indep_size[1],indep_size[1]);
		//} else if(variance_type == VAR_WHITE) {				
		//	temp = calculateWhiteVCVMatrix(temp);
		//}
		
		covariance_matrix = ls_coef;
		
		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;
	}

	/** Performs a single NLLS regression step and updates all estimates accordingly
	 * 
	 * @return true if successful, false otherwise
	 */
	private boolean regress() {
		if(!success) return false;
		
		// Construct the A matrix
		double[][] A = build_A();
		if(!success) {
			failure_message = "Failed to properly construct derivative matrix.";
			return false;
		}
		
		// calculate db
		double[] db = getResiduals();
		if(db == null) return false;
		
		// Compute the array and vector needed for the LS fit
		double[][] ls_coef = new double[parameters.length][parameters.length];
		DGEMM.DGEMM("T","N",parameters.length,parameters.length,use_row.length,
				    1.0,A,A,0.0,ls_coef);
		
		double[] ls_sol = new double[parameters.length];
		DGEMV.DGEMV("T",use_row.length,parameters.length,1.0,A,db,1,0.0,ls_sol,1);
		
		// Determine if the matrix is acceptably conditioned
		double rcondition = MatrixSupport.GetRecipricolCondition(ls_coef);
		
		// This is a perfectly-determined system - calculate the corrections
		// to the estimates
		if(rcondition > RunLS.MINIMUM_R_CONDITION) {
			
			// Need a 2d solution vector
			double[][] solved = new double[ls_sol.length][1];
			for(int i=0;i<ls_sol.length;i++) {
				solved[i][0] = ls_sol[i];
			}
			
			org.netlib.util.intW info = new org.netlib.util.intW(0);
			int[] ipiv = new int[parameters.length];
			DGETRF.DGETRF(ipiv.length,ipiv.length,ls_coef,ipiv,info);
			if(info.val == 0) {
				DGESV.DGESV(parameters.length,1,ls_coef,ipiv,solved,info);
			}
			
			success = (info.val==0);
			
			// Update the estimates
			for(int i=0;i<ls_sol.length && success;i++) {
				ls_sol[i] = solved[i][0];
			}
			
			// Compute the new estimates
			if(success) {
				DAXPY.DAXPY(estimates.length,RELAXATION_MULTIPLIER*relaxation,ls_sol,1,estimates,1);
				change = DASUM.DASUM(ls_sol.length,ls_sol,1);
			}
			
		} else {
			failure_message = "Bad conditioning of independent data matrix\nRCOND="+(new Double(rcondition).toString()); 
			success = false;				
		}
		
		return success;
	}
	
	/** 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
			return evaluate_equation(equation);

	}
	
	/** Constructs the A matrix composed of evaluated derivatives w.r.t. each
	 * parameter evaluated at each row.  The dimensions will be 
	 * (Number of model rows)x(Number of parameters).  If the function fails,
	 * the success boolean will be set to false.
	 * 
	 * @return a 2D array of double values, or null if there was a problem
	 */
	private double[][] build_A() {
		
		if(derivative_eqns != null) {
			return build_A_symbolic();
		} else {
			return build_A_numeric();
		}
		
	}
	
	private double[] evaluate_equation(String eqn) {
		Expression x = ExpressionTree.parse(eqn);
		
		double[] b = new double[model.getRowCount()];
		VarMap variables = set_estimates(null);
		variables = set_constants(variables);
		
		FuncMap fm = new FuncMap();
		fm.loadDefaultFunctions();
		
		for(int i=0;i<use_row.length && success;i++) {
			variables = set_variables(variables,i);
			
			try {
				b[i] = x.eval(variables,fm);
			} catch(java.lang.RuntimeException re) {
				//JOptionPane.showMessageDialog(jContentPane,re.getMessage(), "Equation Error", JOptionPane.ERROR_MESSAGE);
				b[i] = Double.NaN;
			}
			
			variables = clear_variables(variables);
		}
		
		return b;
	}
	
	private double[] evaluate_derivative(String eqn, int parameter_index) {

		Expression x = ExpressionTree.parse(eqn);
		
		double[] b = new double[model.getRowCount()];

		FuncMap fm = new FuncMap();
		fm.loadDefaultFunctions();
		
		for(int i=0;i<use_row.length && success;i++) {
			VarMap variables = set_estimates(null);
			variables = set_constants(variables);
			variables = set_variables(variables,i);
			
			b[i] = 0.0;
			
			int passes = 0;
			
			double multiplier = INITIAL_STEP;
			//while((b[i] == 0 || b[i] == Double.NaN) && (passes < MAX_TRIES_FOR_NUM_DERIVATIVE)) {
				// Set the increment
				double increment = 0.0;
				if(estimates[parameter_index] == 0.0) {
					increment = multiplier;
				} else {
					increment = multiplier*estimates[parameter_index];
				}
				
				double nominal = 0.0;
				try {
					nominal = x.eval(variables,fm);
				} catch(java.lang.RuntimeException re) {
					//JOptionPane.showMessageDialog(jContentPane,re.getMessage(), "Equation Error", JOptionPane.ERROR_MESSAGE);
					b[i] = Double.NaN;
					break;
				}
				
				double incremented = 0.0;
				variables.setValue(parameters[parameter_index],estimates[parameter_index]+increment);
				try {
					incremented = x.eval(variables,fm);
				} catch(java.lang.RuntimeException re) {
					//JOptionPane.showMessageDialog(jContentPane,re.getMessage(), "Equation Error", JOptionPane.ERROR_MESSAGE);
					b[i] = Double.NaN;
					break;
				}
				
				b[i] = (incremented-nominal)/increment;
				multiplier = 10.0*multiplier;
				passes++;
			//}
			
			if(passes>=MAX_TRIES_FOR_NUM_DERIVATIVE) {
				success = false;
				b=null;
				break;
			}
			
		}
		
		return b;
	}
	
	private double[][] build_A_symbolic() {
		// Construct the A matrix
		double[][] A = new double[use_row.length][Math.min(estimates.length,parameters.length)];
		
		for(int j=0;j<Math.min(estimates.length,parameters.length);j++) {
			double[] results = evaluate_equation(derivative_eqns[j]);
			for(int i=0;i<use_row.length;i++) {
				A[i][j] = results[i];
			}
		}
		return A;
	}
	
	private double[][] build_A_numeric() {
		// Construct the A matrix
		double[][] A = new double[use_row.length][Math.min(estimates.length,parameters.length)];
		
		for(int i=0;i<Math.min(estimates.length,parameters.length);i++){
			double[] derivs = evaluate_derivative(equation,i);
			for(int j=0;j<use_row.length && success;j++)
				A[j][i] = derivs[j];
		}
		if(success)
			return A;
		else
			return null;
	}
	
	/** Updates the object's parser with all the estimates of parameters
	 * 
	 */
	private VarMap set_estimates(VarMap variables) {
		if(variables == null)
			variables = new VarMap();
		
		for(int i=0;i<Math.min(estimates.length,parameters.length);i++) {
			variables.setValue(parameters[i],estimates[i]);
		}
		return variables;
	}
	
	
	
	/** Updates the parser with all variable values for a given row
	 * 
	 * @param row the current row of the StatTableModel to use
	 */
	private VarMap set_variables(VarMap variables, int row) {
		if(variables == null)
			variables = new VarMap();
		
		for(int i=0;i<model.getColumnCount();i++) {
			Object value = model.getValueAt(row,i);
			if(value == null) continue;
			if(value.getClass().getSuperclass() != Number.class) continue;
			variables.setValue(model.getColumnName(i),((Number)value).doubleValue());
		}
		return variables;
	}
	
	private VarMap clear_variables(VarMap variables) {
		if(variables == null)
			return null;
		
		for(int i=0;i<model.getColumnCount();i++) {
			variables.remove(model.getColumnName(i));
		}
		return variables;
	}
	
	/** Updates JEP with all constant values in the StatTableModel
	 * 
	 */
	private VarMap set_constants(VarMap variables) {
		if(variables == null)
			variables = new VarMap();
		
		ArrayList constants = model.getConstants();
		for(int i=0;i<constants.size();i++) {
			Constant this_one = (Constant)constants.get(i);
			variables.setValue(this_one.getName(),this_one.getValue().doubleValue());
		}
		return variables;
	}
	
	/** Simplistic method for determining which rows are acceptable for use in 
	 * the fitting.  If the row contains ANY null values, the row is deemed unfit,
	 * regardless of whether the variable appears in the equation.
	 *
	 */ 
	private void determine_allowable_rows() {
		use_row = model.getNonNullRows();
	}
	
	/** 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;
	}
}
