package hep.aida.ref.fitter;

import hep.aida.IAnalysisFactory;
import hep.aida.IAnnotation;
import hep.aida.IBaseHistogram;
import hep.aida.ICloud;
import hep.aida.IDataPointSet;
import hep.aida.IFitData;
import hep.aida.IFitParameterSettings;
import hep.aida.IFitResult;
import hep.aida.IFunction;
import hep.aida.IFunctionFactory;
import hep.aida.IHistogram;
import hep.aida.IHistogram1D;
import hep.aida.IHistogramFactory;
import hep.aida.IManagedObject;
import hep.aida.IModelFunction;
import hep.aida.IProfile;
import hep.aida.IRangeSet;
import hep.aida.ITree;
import hep.aida.ITupleFactory;
import hep.aida.dev.IDevFitData;
import hep.aida.dev.IDevFitDataIterator;
import hep.aida.dev.IDevFitResult;
import hep.aida.ext.IExtFitter;
import hep.aida.ext.IFitMethod;
import hep.aida.ext.IOptimizer;
import hep.aida.ext.IOptimizerConfiguration;
import hep.aida.ext.IOptimizerFactory;
import hep.aida.ext.IVariableSettings;
import hep.aida.ref.AidaUtils;
import hep.aida.ref.fitter.fitdata.FitDataCreator;
import hep.aida.ref.function.BaseModelFunction;
import hep.aida.ref.function.FunctionCatalog;
import hep.aida.ref.function.RangeSet;
import hep.aida.ref.histogram.DataPointSet;
import hep.aida.ref.pdf.Dependent;
import hep.aida.ref.pdf.Function;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Enumeration;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.Map;
import java.util.StringTokenizer;
import java.util.regex.Pattern;

import org.freehep.util.FreeHEPLookup;
import org.openide.util.Lookup;

/**
 * @author The AIDA team @ SLAC.
 *
 */

public class Fitter implements IExtFitter {
    
    // The IOptimizer.
    private String engineType;
    private IOptimizer optimizer = null;
    
    // The IFitMethod.
    private String fitMethodType;
    private IFitMethod fitMethod = null;
    
    // The IFunction to be minimized.
    private IFunction fitFunction;
    
    // IFitParameterSettings
    private Hashtable fitParHash = new Hashtable();
    
    private boolean useGradient = true;
    
    private ArrayList constraintList = new ArrayList();
    private Hashtable simpleConstraintHash = new Hashtable();
    private boolean createClone = true;
    
    /**
     * Create a new Fitter specifying the underlying optimizing engine.
     * @param fitMethodType The type of fitter.
     * @param engineType The type of optimizer to use.
     * @throws IllegalArgumentException if the engineType does not exist.
     *
     */
    public Fitter(String fitMethodType, String engineType, String options) throws IllegalArgumentException {
        setFitMethod(fitMethodType);
        setEngine(engineType);
        
        Map opt = hep.aida.ref.AidaUtils.parseOptions(options);
        String val = (String) opt.get("noClone");
        if ( val != null && val.trim().equalsIgnoreCase("true") ) createClone = false;
    }
    
    public void setEngine(String engineType) throws IllegalArgumentException {
        if (engineType == null || engineType.length() == 0) engineType = "jminuit";
        String enType = engineType.toLowerCase();
        
        IOptimizerFactory tmpOptimizerFactory = null;
        Lookup.Template template = new Lookup.Template(IOptimizerFactory.class);
        Lookup.Result result = FreeHEPLookup.instance().lookup(template);
        Collection c = result.allInstances();
        for (Iterator i = c.iterator(); i.hasNext(); ) {
            IOptimizerFactory of = (IOptimizerFactory)i.next();
            String[] names = of.optimizerFactoryNames();
            if ( names == null || names.length == 0 )
                throw new IllegalArgumentException("IOptimizerFactory with illegal names!");
            for ( int j = 0; j < names.length; j++ ) {
                if ( enType.equals( names[j].toLowerCase() ) ) {
                    tmpOptimizerFactory = of;
                    break;
                }
            }
        }
        if (tmpOptimizerFactory == null) throw new IllegalArgumentException("Cannot create IOptimizer of type: "+engineType);
        this.engineType = engineType;
        this.optimizer = tmpOptimizerFactory.create(engineType);
    }
    public String engineName() {
        return engineType;
    }
    
    public void setFitMethod(String fitMethodType) throws IllegalArgumentException {
        if (fitMethodType == null || fitMethodType.length() == 0) fitMethodType = "chi2";
        // Check the lookup table to look for the fitMethod of the given type.
        String fitMet = fitMethodType.toLowerCase();
        
        IFitMethod tmpFitMethod = null;
        Lookup.Template template = new Lookup.Template(IFitMethod.class);
        Lookup.Result result = FreeHEPLookup.instance().lookup(template);
        Collection c = result.allInstances();
        for (Iterator i = c.iterator(); i.hasNext(); ) {
            IFitMethod fm = (IFitMethod)i.next();
            String[] names = fm.fitMethodNames();
            if ( names == null || names.length == 0 )
                throw new IllegalArgumentException("IFitMethod with illegal names!");
            for ( int j = 0; j < names.length; j++ ) {
                if ( fitMet.equals( names[j].toLowerCase() ) ) {
                    tmpFitMethod = fm;
                    break;
                }
            }
        }
        if (tmpFitMethod == null) throw new IllegalArgumentException("Unknown IFitMethod type: "+fitMethodType);
        this.fitMethodType = fitMethodType;
        this.fitMethod = tmpFitMethod;
    }
    
    public IOptimizer optimizer() {
        return optimizer;
    }
    
    public String fitMethodName() {
        return fitMethodType;
    }
    
    public IFitParameterSettings fitParameterSettings(String name) {
        if ( fitParHash.containsKey(name) ) return (IFitParameterSettings) fitParHash.get(name);
        IFitParameterSettings fitPar = new FitParameterSettings(name);
        fitParHash.put(name,fitPar);
        return fitPar;
    }
    
    public String[] listParameterSettings() {
        int size = fitParHash.size();
        String[] parNames = new String[size];
        Enumeration e = fitParHash.keys();
        for (int i = 0; e.hasMoreElements(); i++)
            parNames[i] = (String) e.nextElement();
        return parNames;
    }
    
    public void resetParameterSettings() {
        fitParHash.clear();
    }
    
    public IFitResult fit( Function f ) {
        if ( f.numberOfDependents() != 1 )
            throw new IllegalArgumentException("Currently we only fit 1 dimensional functions.");
        Dependent x = f.getDependent(0);
        if ( ! x.isConnected() )
            throw new IllegalArgumentException("Dependent "+x.name()+" is not connected to a data set");
        
        IFitData data = x.data();
        return fit( data, f );
    }
    
    public IFitResult fit(IFitData d, IFunction originalFunction) {
        return fit(d, originalFunction, null);
    }
    
    public IFitResult fit(IFitData d, IFunction originalFunction, String range) {
        return fit(d, originalFunction, range, null);
    }
    
    private void addBounds(IRangeSet rs, int index, double[] lb, double[] ub) {
        if ( lb.length == 0 ) throw new RuntimeException("No ranges are set in the FitData. Please report the problem");
        for ( int j = 0; j < ub.length; j++ )
                rs.include(lb[j],ub[j]);
    }
    
    public IFitResult fit(IFitData d, IFunction originalFunction, String range, Object correlationObject) {
        
        String name = "";
        if ( originalFunction instanceof IManagedObject ) {
            name = ((IManagedObject) originalFunction).name();
        } else {
            name = originalFunction.title();
        }
        name += " "+engineName()+" fit";
        
        IFunction fClone = cloneFunction(name,originalFunction);
        
        boolean setRange = false;
        if ( originalFunction instanceof IFunction )
            setRange = true;
        //        IModelFunction f = (IModelFunction)originalFunction;
        IModelFunction f;
        
        if ( fClone instanceof IModelFunction ) {
            f = (IModelFunction)fClone;
        } else if (fClone instanceof IFunction) {
            f = new BaseModelFunction(name, name, fClone);
        } else {
            throw new RuntimeException("Fitter for now can only use IModelFunctions and IFunctions. Please report this problem");
        }
        
        if ( range != null ) {
            Map optionsMap = AidaUtils.parseOptions( range );
            range = (String) optionsMap.get("range");
            RangeSet usersRange = new RangeSet(range);
            IRangeSet rs = d.range(0);
            rs.excludeAll();
            addBounds(rs, 0, usersRange.lowerBounds(), usersRange.upperBounds());
        }

        if ( setRange ) {
            f.excludeNormalizationAll();
                //Set ranges on function.            
                IRangeSet d_rs = null;
                for ( int i = 0; i < d.dimension(); i++ ) {
                    d_rs = d.range(i);
                    IRangeSet f_rs = f.normalizationRange(i);
                    addBounds(f_rs, i, d_rs.lowerBounds(), d_rs.upperBounds());
                }
        }        
        
        //Clear the fit method and set the correlation Object
        fitMethod.clear();
        fitMethod.setCorrelationObject(correlationObject);
                
        loadFitDataAndFunction(d,f);
        
        long startFit = System.currentTimeMillis();
        optimizer.optimize();
        long endFit = System.currentTimeMillis();
        double fitSeconds = (double)(endFit-startFit)/1000.;
        
        // Set the minimum parameters on the function
        String[] parNames = f.parameterNames();
        for( int i = 0; i < parNames.length; i++ )
            f.setParameter(parNames[i],optimizer.result().parameters()[i]);
        
        double[][] covMatrix = optimizer.result().covarianceMatrix();
        
        int status = optimizer.result().optimizationStatus();
        int dataEntries = ((FitFunction) fitFunction).dataEntries();
        int freePars = ((FitFunction) fitFunction).nFreePars();
        int nDoF = dataEntries - freePars;
        double funcVal = fitFunction.value(f.parameters());
        
        IDevFitResult result = new FitResult(fitFunction.dimension(), fitSeconds);
        result.setConstraints( constraints() );
        result.setDataDescription( d.dataDescription() );
        result.setEngineName( engineName() );
        result.setFitMethodName( fitMethodName() );
        result.setFitStatus( status );
        result.setFittedFunction( f ); // FIX ME! Replace with clone.
        result.setIsValid( true ); ////??????
        result.setNdf( nDoF );
                
        if ( fitMethod.fitType() == IFitMethod.UNBINNED_FIT )
            result.setQuality( funcVal/nDoF/Math.sqrt(2.) );
        else
            result.setQuality( funcVal/nDoF );
        
        int countI = 0;
        for( int i = 0; i < parNames.length; i++ ) {
            result.setFitParameterSettings( parNames[i], fitParameterSettings(parNames[i]) );
            int countJ = 0;
            for( int j = 0; j < parNames.length; j++ ) {
                if ( ! optimizer.variableSettings(parNames[i]).isFixed() &&  ! optimizer.variableSettings(parNames[j]).isFixed() )
                    result.setCovMatrixElement( i, j , covMatrix[countI][countJ++] );
            }
            if ( ! optimizer.variableSettings(parNames[i]).isFixed() ) countI++;
        }
        
        return (IFitResult)result;
    }
    
    
    public IFitResult fit(IBaseHistogram h, IFunction f) {
        return fit( h, f, null);
    }

    public IFitResult fit(IBaseHistogram h, IFunction f, String range) {    
        if ( h instanceof IHistogram || h instanceof IProfile )
            if ( fitMethod.fitType() == IFitMethod.UNBINNED_FIT ) throw new IllegalArgumentException("Cannot perform unbinned fit on a IHistogram!!");
            else if ( h instanceof ICloud )
                if ( fitMethod.fitType() == IFitMethod.BINNED_FIT ) throw new IllegalArgumentException("Cannot perform binned fit on a ICloud!!");
        IFitData fitData = FitDataCreator.create(h);
        return fit( fitData, f, range );
    }    
    
    private void setDefaultInitialParameters(IFunction f, String model, Object data) {
        if ( model.equals("g") ) {
            if ( data instanceof IHistogram1D ) {
                IHistogram1D h1 = (IHistogram1D) data;
                double[] pars = new double[3];
                pars[0] = h1.maxBinHeight();
                pars[1] = h1.mean();
                pars[2] = h1.rms();
                f.setParameters(pars);
            }
        }        
    }
    
    public IFitResult fit(IBaseHistogram h, String model) {
        return fit(h, model, (String)null);
    }
    
    public IFitResult fit(IBaseHistogram h, String model, String range) {
        IFunction func = FunctionCatalog.getFunctionCatalog().create(model);
        setDefaultInitialParameters(func, model, h);
        return fit(h,func, range);
    }

    public IFitResult fit(IBaseHistogram h, String model, double[] initialParameters) {
        return fit(h,model, initialParameters, null);
    }
    
    public IFitResult fit(IBaseHistogram h, String model, double[] initialParameters, String range) {
        IFunction func = FunctionCatalog.getFunctionCatalog().create(model);
        if ( initialParameters.length != func.numberOfParameters() ) throw new IllegalArgumentException("Wrong number of parameters "+initialParameters.length+
        "! This function requires "+func.numberOfParameters());
        func.setParameters(initialParameters);
        return fit(h,func, range);
    }
    
    public IFitResult fit(IDataPointSet dataPointSet, IFunction f) {
        return fit(dataPointSet, f, (String)null);
    }
    public IFitResult fit(IDataPointSet dataPointSet, IFunction f, String range) {
        return fit(dataPointSet, f, range, null);
    }
    
    public IFitResult fit(IDataPointSet dataPointSet, IFunction f, double[] initialParameters) {
        return fit(dataPointSet, f, initialParameters, null, null);
    }
    
    public IFitResult fit(IDataPointSet dataPointSet, IFunction f, String range, Object correlationObject) {
        return fit(dataPointSet, f, null, range, correlationObject);
    }
    
    public IFitResult fit(IDataPointSet dataPointSet, IFunction f, double[] initialParameters, String range, Object correlationObject) {
        if ( fitMethod.fitType() == IFitMethod.UNBINNED_FIT ) throw new IllegalArgumentException("Cannot perform unbinned fit on a IDataPointSet!!");
        if ( dataPointSet.dimension() != f.dimension() + 1 ) throw new IllegalArgumentException("Wrong dimension match. DataPointSets can only be fitted if"+
        " their dimension is one unit bigger than the one of the function");
        
        if ( initialParameters != null ) {
            if ( initialParameters.length != f.numberOfParameters() ) throw new IllegalArgumentException("Wrong number of parameters "+initialParameters.length+
            "! This function requires "+f.numberOfParameters());
            f.setParameters(initialParameters);
        }
        
        IFitData fitData = FitDataCreator.create(dataPointSet);
        
        return fit( fitData, f, range, correlationObject );
    }
    
    public IFitResult fit(IDataPointSet dataPointSet, String model) {
        return fit(dataPointSet, model, (String)null);
    }
    public IFitResult fit(IDataPointSet dataPointSet, String model, String range) {
        return fit(dataPointSet,model, range, null);
    }
    
    public IFitResult fit(IDataPointSet dataPointSet, String model, String range, Object correlationObject) {
        IFunction func = FunctionCatalog.getFunctionCatalog().create(model);
        setDefaultInitialParameters(func, model, dataPointSet);
        return fit(dataPointSet,func, range, correlationObject);
    }
    
    public IFitResult fit(IDataPointSet dataPointSet, String model, double[] initialParameters) {
        return fit(dataPointSet, model, initialParameters, null, null);        
    }
    public IFitResult fit(IDataPointSet dataPointSet, String model, double[] initialParameters, String range) {
        return fit(dataPointSet, model, initialParameters, range, null);
    }
    
    public IFitResult fit(IDataPointSet dataPointSet, String model, double[] initialParameters, String range, Object correlationObject) {
        IFunction func = FunctionCatalog.getFunctionCatalog().create(model);
        
        return fit(dataPointSet,func, initialParameters, range, correlationObject);
    }
    
    
    public IFitResult fit(IFitData d, String model, double[] initialParameters) {
        return fit(d,model,initialParameters,null);
    }
    
    public IFitResult fit(IFitData d, String model, double[] initialParameters, String range) {
        return fit(d,model,initialParameters, range, null);
    }

    public IFitResult fit(IFitData d, String model, double[] initialParameters, String range, Object correlationObject) {
        IFunction func = FunctionCatalog.getFunctionCatalog().create(model);
        
        if ( initialParameters != null ) {
            if ( initialParameters.length != func.numberOfParameters() ) throw new IllegalArgumentException("Wrong number of parameters "+initialParameters.length+
            "! This function requires "+func.numberOfParameters());
            func.setParameters(initialParameters);
        }
        return fit(d,func, range, correlationObject);
    }
    
    public IFitResult fit(IFitData d, String model) {
        return fit(d,model,(String)null);
    }
    
    public IFitResult fit(IFitData d, String model, String range) {
        return fit(d,model, range, null);
    }

    public IFitResult fit(IFitData d, String model, String range, Object correlationObject) {
        IFunction func = FunctionCatalog.getFunctionCatalog().create(model);
        return fit(d,func, range, correlationObject);
    }
    
    public void setConstraint(String expression) throws IllegalArgumentException {
        if ( ! optimizer.acceptsConstraints() )
            throw new UnsupportedOperationException("Optimizer "+engineName()+" does not accept constraints.");
        StringTokenizer st = new StringTokenizer(expression,"=");
        if ( st.countTokens() != 2 ) throw new IllegalArgumentException("Only constraints of the form \" parName = SomeExpression \" are supported");
        String parName = st.nextToken().trim();
        if ( ! Pattern.matches("\\w+", parName) ) throw new IllegalArgumentException("Incorrect parameter name "+parName);
        String constrExpression = st.nextToken().trim();
        
        if ( Pattern.matches("\\w+",constrExpression) ) {
            simpleConstraintHash.put(parName,constrExpression);
        } else {
            throw new IllegalArgumentException("This type of constraint is not supported yet "+parName+" = "+constrExpression);
        }
        constraintList.add(expression);
    }
    public String[] constraints() {
        int size = constraintList.size();
        String[] constraints = new String[size];
        for ( int i = 0; i pmax ) throw new IllegalArgumentException("Incorret parameter limits : "+pmin+" has to be less than "+pmax);
        if ( npts < 1 ) throw new IllegalArgumentException("The number of points has to be greater than zero : "+npts);
        
        IFunction fClone = cloneFunction(null, originalFunction);
        
        IModelFunction f;
        if ( fClone instanceof IModelFunction ) {
            f = (IModelFunction)fClone;
        } else {
            f = new BaseModelFunction("", "", fClone);
        }
        
        loadFitDataAndFunction(d,f);
        
        int index = f.indexOfParameter( parName );
        
        IFitParameterSettings fitParSet = fitParameterSettings(parName);
        if ( fitParSet.isFixed() ) throw new IllegalArgumentException("Parameter "+parName+" is fixed");
        
        IDataPointSet dps = new DataPointSet("",parName+" scan",2,npts);
        double step = (pmax - pmin)/(npts-1);
        
        double startValue = optimizer.variableSettings(parName).value();
        
        for ( int i = 0; i -1 && ind2 > -1 ) {
                varSimpleConstraint1.add( new Integer(ind1) );
                varSimpleConstraint2.add( new Integer(ind2) );
            }
        }
        
        protected boolean isValidSimpleConstraint(String varName1, String varName2 ) {
            int ind1 = indexOfVariable( varName1 );
            int ind2 = indexOfVariable( varName2 );
            if ( ind1 > -1 && ind2 > -1 ) return true;
            return false;
        }
        
        protected void applySimpleConstraint( double[] x ) {
            for ( int i = 0; i < varSimpleConstraint1.size(); i++ ) {
                int ind1 = ( (Integer) varSimpleConstraint1.get(i) ).intValue();
                int ind2 = ( (Integer) varSimpleConstraint2.get(i) ).intValue();
                x[ind1] = x[ind2];
            }
        }
        
    }
}