/*
 *  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 com.approximatrix.draco.datamodel.StatTableModel;
import com.approximatrix.draco.fit.statistics.ExtendedStatistic;
import com.approximatrix.draco.io.XMLio.AnalysisLoadStore;
import com.approximatrix.draco.utility.ResultsUtilities;

public class RunTSLS implements LSFit {
	
	private static final String RESULTS_TITLE = "Results of the Two-Stage Least Squares Regression\n";
	
	private boolean INCLUDE_CONSTANT_ON_STAGE_1 = true;
	
	private String[] indep_names = null;
	private String[] instrument_names = null;
	private String dep_name = null;
	
	// Results information
	private boolean success = false;
	
	private int variance_type = 1;
	private boolean include_constant = true;
	
	private String results_message = null;
	private String failure_message = null;
	
	private boolean[] extended_statistics = null;
	
	private StatTableModel base_model = null;
	
	private RunLS results = null;
	
	private boolean ignore_conditioning = false;
	
	/** The list of extended statistics to compute */
	private List<ExtendedStatistic> estats = null;
	
	public RunTSLS() {
		failure_message = "No data provided...";
	}
	
	public RunTSLS(String[] indep, String dep, 
			String instruments[], 
			StatTableModel passed_model,
			int in_variance_type,
			boolean calc_constant){

		base_model = passed_model;
		indep_names = indep;
		dep_name = dep;
		instrument_names = instruments;
		variance_type = in_variance_type;
		include_constant = calc_constant;

	}
	
	public RunTSLS(String[] indep, String dep, 
			String instruments[], 
			StatTableModel passed_model,
			int in_variance_type,
			boolean calc_constant,
			boolean ignore_condition){

		base_model = passed_model;
		indep_names = indep;
		dep_name = dep;
		instrument_names = instruments;
		variance_type = in_variance_type;
		include_constant = calc_constant;
		ignore_conditioning = ignore_condition;

	}
	
	private boolean regress() {

		// A separate list of independent variables to use for solving that
		// will be modified with estimated names
		String[] solving_independents = new String[indep_names.length];
		for(int i=0;i<solving_independents.length;i++)
			solving_independents[i] = new String(indep_names[i]);
		
		//StatTableModel 
		StatTableModel tsmodel = new StatTableModel(base_model);
				
		// Loop through all independents and solve for hat values if not instruments
		for(int i=0;i<indep_names.length;i++) {
			if(! is_one(indep_names[i],instrument_names)) {
				
				// This is not an instrument, you must
				// solve for the new value
				RunLS one_var = new RunLS(instrument_names,
										  indep_names[i],
										  base_model,
										  RunLS.VAR_OLS,
										  INCLUDE_CONSTANT_ON_STAGE_1,
										  ignore_conditioning);
				
				one_var.perform();
				
				if(!one_var.isSuccessful()) {
					failure_message = new String("Regression failed on stage 1: "+
							                     indep_names[i]+"\n\nRunLS Output:\n"+
							                     one_var.getFailureMessage());
					return false;
				}
				
				// Change the solving independent
				solving_independents[i] = new String(indep_names[i]+"_HAT");
				
				// Add the hat values to the ts array
				tsmodel.addColumn(one_var.generate_data_from_fit(),solving_independents[i]);
			}			
		}
		
		// Perform a new TSLS fit using HAT values
		results = new RunLS(solving_independents,
							dep_name,
							tsmodel,
							variance_type,
							include_constant,
							ignore_conditioning);
		
		results.perform();
		
		return results.isSuccessful();
	}
	
	private boolean is_one(String look_for, String[] look_in) {
		boolean return_val = false;
		
		for(int i=0;i<look_in.length;i++) {
			if(look_in[i].compareTo(look_for)==0) {
				return_val = true;
				break;
			}
		}
		
		return return_val;
	}
	
	public String[] getInstruments() {
		return instrument_names;
	}
	
	public double[] getCoef() {
		return results.getCoef();
	}

	public String[] getCoefLabels() {
		return results.getCoefLabels();
	}

	public double[] getCoefVariances() {
		return results.getCoefVariances();
	}

	public double[][] getCovarianceMatrix() {
		return results.getCovarianceMatrix();
	}

	public String getDependentVariable() {
		return this.dep_name;
	}

	public Double[] getFittedData() {
		return results.getFittedData();
	}

	public String[] getIndependentVariables() {
		//return indep_names;
		return results.getIndependentVariables();
	}

	public StatTableModel getModel() {
		return base_model;
	}

	public Double[] getResidualData() {
		return results.getResidualData();
	}

	public double getResidualSumSqared() {
		return results.getResidualSumSqared();
	}

	public double[] getResiduals() {
		return results.getResiduals();
	}

	public String getResults() {
		if(results_message == null) {
			if(!success) {
				
				if(failure_message == null)
					failure_message = new String("Regression failed on stage 2:\n\nRunLS output:\n"+ failure_message);
				
				failure_message = ResultsUtilities.PrepStringForOutput(failure_message);
				
				results_message = LSFitResultsGen.GenerateFailureResults(this, failure_message, RESULTS_TITLE).toString(); 
				
			} else {
				results_message = LSFitResultsGen.GenerateResults(this, RESULTS_TITLE).toString();
			}
		}
		return results_message;
	}

	public double getTotalSumSqared() {
		return results.getTotalSumSqared();
	}

	public boolean isSuccessful() {
		return success;
	}

	/** 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.base_model = model;
		
		this.results = new RunLS();
		this.results.setLoadedAnalysis(data, model);
		
		this.indep_names = data.getIndependentVariables();
		this.dep_name = data.getDependentVariable();
		this.instrument_names = data.getInstruments();
		
		this.results_message = data.getResults();
		this.success = true;
			
	}

	public void perform() {
		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;
	}
}
