/*  Monte-Carlo simulation code for statistical physics
    Copyright (C) 2001-2005  Sylvain Reynal
    
    Département de Physique
    Ecole Nationale Supérieure de l'Electronique et de ses Applications (ENSEA)
    6, avenue du Ponceau,    
    F-95014 CERGY CEDEX

    et
    
    Laboratoire de Physique Théorique et Modélisation (LPTM)
    CNRS/Université de Cergy-Pontoise - Site de Neuville
    F-95031 CERGY CEDEX
    
	
    Tel : 00 +33 130 736 245
    Fax : 00 +33 130 736 667
    e-mail : reynal@ensea.fr
    web page : http://www.ensea.fr/staff/reynal/
*/


package fr.ensea.montecarlo.multicanonical;

import fr.ensea.montecarlo.data.*;
import fr.ensea.chart.*;
import fr.ensea.montecarlo.model.*;

import cern.colt.matrix.DoubleFactory2D;
import cern.colt.matrix.DoubleMatrix2D;
import cern.colt.matrix.linalg.QRDecomposition;
import cern.colt.list.*;
import hep.aida.ref.*;
import java.io.*;

/**
 * DOS-histogram as well as S(E) interpolator for the Wang-Landau iteration scheme.
 */
public class WLDOSHistogram extends DynamicHistogram1D {

	//protected int eneMaxBin, eneMinBin; // bin corresponding to eneMax, or -1 (OVERFLOW) if eneMax hasn't been set yet
	protected double[] SE; // S(E) = SE[bin] if bin inside axis edges
	protected boolean[] isNonNullBin; // true = the corresponding group of enery levels exists ; false = this bin can never be filled
	protected double weight; // current weight, so that fill(x) actually add "weight" to the corresponding SE[bin]
	protected double initWeight;
	
	protected double maxSE; // for the UI

	/**
	 * Creates a histogram
	 * @param dExAxis energy levels axis bin width
	 * @param dEyAxis delta-energy axis bin width (for transition matrix only)
	 * @param weight update weight for S(E)
	 */
	public WLDOSHistogram(double dE, double weight){
		super(dE);
		SE = new double[1]; // i.e. 1 bin
		isNonNullBin = new boolean[1]; // i.e. 1 bin ; false by default i.e. this is a null bin, until it gets filled
		this.initWeight = this.weight=weight;
		maxSE=0;
	}
	
	public void init(double dE, double weight){
		resetBinWidth(dE);
		SE = new double[1]; // i.e. 1 bin
		isNonNullBin = new boolean[1]; // i.e. 1 bin ; false by default i.e. this is a null bin, until it gets filled
		this.weight=weight;
		maxSE=0;
	}
	
	// clear S(E) and axis	
	public void reinit(){
		init(binWidth, weight);
	}

	// clear S(E) and axis	
	public void reinit(int pottsQ){
		if (pottsQ==2) init(2.0, this.initWeight);
		else init(1.0, this.initWeight);
	}

	/////////////////////////////////////////////////////////////////////////////////
	// Properties
	/////////////////////////////////////////////////////////////////////////////////

	/**
	 * Return the current weight value
	 */
	public double getWeight(){
		return weight;
	}

	/**
	 * Sets the current weight value
	 */
	public void setWeight(double w){
		this.weight=w;
	}

	public double getDeltaE(){
		return getBinWidth();
	}


	/////////////////////////////////////////////////////////////////////////////////
	// Filling and related
	/////////////////////////////////////////////////////////////////////////////////

	/**
	 * Init the x-axis from the given bin center, resetting bin entries to their default value.
	 */
	protected void initAxis(double x){
		super.initAxis(x);
		this.SE = new double[1];
		SE[0] = 0.0; // i.e. bin(0) only
		maxSE=0;
		isNonNullBin[0]=true; // it's not a non-null bin, since there's at least one entry in it
	}

	/**
	 * Prepends new bins to the left of the current axis, and shift "SE" array accordingly.
	 */
	protected void prependBins(double x){
		int oldNbBins = xAxis.bins();
		super.prependBins(x); // modifies this.xAxis
		int newNbBins = xAxis.bins();
		if (newNbBins == oldNbBins) return;

		int destPos = newNbBins-oldNbBins; // i.e. where the old axis starts wrt the new axis

		// --- copy SE[E] ---
		double[] _SE = new double[newNbBins];
		System.arraycopy(SE,0,_SE, destPos,oldNbBins);
		SE=_SE;

		// --- copy isNonNullBin[] ---
		boolean[] _isNonNullBin = new boolean[newNbBins];
		System.arraycopy(isNonNullBin,0,_isNonNullBin, destPos,oldNbBins);
		isNonNullBin=_isNonNullBin;

		// --- init S(E) for new-axis-lower-edge <= E <= old-axis-lower-edge, using linear predictor ---
		double beta = computeBetaEPredictor();
		for (int b = 0; b < destPos; b++){
			SE[b] = SE[destPos] + beta * (xAxis.binCentre(b)-xAxis.binCentre(destPos)) - weight;
		}
	}

	/**
	 * Appends new bins to the right of the current axis, and shifts "cumG0" accordingly
	 */
	protected void appendBins(double x){
		int oldNbBins = xAxis.bins(); // where old axis ends
		super.appendBins(x); // modifies this.xAxis
		int newNbBins = xAxis.bins();
		if (newNbBins == oldNbBins) return;

		double[] _SE = new double[newNbBins];
		System.arraycopy(SE,0,_SE,0,oldNbBins);
		SE=_SE;

		boolean[] _isNonNullBin = new boolean[newNbBins];
		System.arraycopy(isNonNullBin,0,_isNonNullBin,0,oldNbBins);
		isNonNullBin=_isNonNullBin;

		// --- init S(E) for old-axis-upper-edge <= E <= new-axis-upper-edge, using linear predictor ---
		for (int b = oldNbBins; b < newNbBins; b++){
			SE[b] = 0;
		}
	}

	/**
	 * Fill histogram, and SE[] using current weight
	 */
	public synchronized void fill(double x){
		// --- iterations ---
		int bin = updateAxis(x);
		fillFromBin(bin);
		SE[bin] += this.weight;
		if (SE[bin]>maxSE) maxSE = SE[bin];
		this.isNonNullBin[bin]=true;
	}

	/////////////////////////////////////////////////////////////////////////////////
	// Statistics and related properties
	/////////////////////////////////////////////////////////////////////////////////

	/**
	 * Returns a list of bins corresponding to existing energy levels, b/w eneMin and eneMax
	 */
	protected IntArrayList getEnergyLevelsIndices(){
		IntArrayList l = new IntArrayList();
		for (int bin=0; bin < SE.length; bin++){
			if (this.isNonNullBin[bin]==true) l.add(bin);
		}
		return l;
	}



	/////////////////////////////////////////
	//// Emu(E) interpolator
	///////////////////////////////////////

	/**
	 * Compute the slope of y=SE[a:b-1] where a=indices[i1] and b=indices[i2], using a LS approximation
	 * wrt the xAxis.
	 * That is, we solve:
	 *
	 *                A     *   X     ==        B
	 *
	 * i.e.
	 *
	 *               x_0 1      a               y_0         
	 *               x_1 1  *   b     ==        y_1
	 *               x_2 1                      y_2
	 *               ...                        ...
	 *
	 *          i.e. we solve the system: a * x_0 + b = y_0
	 *                                    a * x_1 + b = y_1
	 *                                    a * x_2 + b = y_2
	 *                                    ...
	 * that is, y=ax+b is the least square solution (but we only return the slope "a")
	 */
	private synchronized double computeBetaEPredictor() {

		IntArrayList nonEmptyBins = this.getIndicesOfNonEmptyBins();
		int N = (int)(nonEmptyBins.size()*0.2);
		if (N < 10) return 0.0;

		DoubleMatrix2D matA = DoubleFactory2D.dense.make(N, 2, 1.0);
		DoubleMatrix2D matB = DoubleFactory2D.dense.make(N, 1, 1.0);
		for (int i = 0; i < N; i++) {
			int bin = nonEmptyBins.get(i);
			// fill first column of A with i*h : (second column filled with ones)
			matA.setQuick(i, 0, xAxis.binCentre(bin));
			// fill first (and sole) column of B with a_i :
			matB.setQuick(i, 0, SE[bin]);
		}
		QRDecomposition qr = new QRDecomposition(matA);
		if (!qr.hasFullRank()) return 0.0;
		else {
			DoubleMatrix2D matX = qr.solve(matB);
			double beta = matX.getQuick(0, 0);
			if (beta<0) beta=0;
			//System.out.println("beta="+beta);
			return beta;
		}
	}

	/**
	 * Returns S(E)
	 */
	public synchronized double getSE(double energy){

		// now, Emin < E < Emax:
		final int bin = xAxis.coordToIndex(energy);
		if (bin >= SE.length) return 0.0;
		switch (bin){
			// if energy outside current histogram axis, use predictor based on beta(E) at edges:
		case UNDERFLOW:
			return 0;
		case OVERFLOW:
			return 0;
		default: // from now on, eneMin < energy < eneMax !!!
			return SE[bin];
		}
	}

	/**
	 * Shift S(E) vertically so that min S(E) = 0
	 */
	public synchronized void normalize(){ // modify this.SE
		// compute mean SE :
		double min=Double.MAX_VALUE;
		IntArrayList nonEmptyBins = getEnergyLevelsIndices();
		for (int i=0; i<nonEmptyBins.size(); i++) {
			int bin = nonEmptyBins.get(i);
			if (SE[bin] < min) min = SE[bin];
		}
		//Log.print("WLDOS: normalizing HmuE", "shift="+shift);
		for (int i=0; i<SE.length; i++) SE[i] -= min;
		maxSE -= min;
	}
	
	///////////////////////////////////////////////////////////////
	/// UI
	///////////////////////////////////////////////////////////////
	
	public SECurve2D createSECurve2D(ILattice lattice){
		return new SECurve2D(lattice);
	}

	public class SECurve2D implements Curve2D {
		
		ILattice lattice;
		
		public SECurve2D(ILattice lattice){
			this.lattice = lattice;
		}

		public synchronized Point2D getPoint(int i){
			if (i>=SE.length) return new Point2D();
			double x = xAxis.binCentre(i)/lattice.getSpinCount();
			double y = SE[i];
			return new Point2D(x,y);
		}

		public int getNbPoints(){
			return bins();
		}

		public double getMin(int axis){
			switch (axis){
			case X_AXIS: return xAxis.binCentre(0)/lattice.getSpinCount();
			case Y_AXIS: return 0;
			default: throw new IllegalArgumentException("Wrong axis index : "+axis);
			}
		}

		public double getMax(int axis){
			switch (axis){
			case X_AXIS: return xAxis.binCentre(bins()-1)/lattice.getSpinCount();
			case Y_AXIS: return maxSE;
			default: throw new IllegalArgumentException("Wrong axis index : "+axis);
			}
		}
	}
	
	/////////////////////////////////////////////////////////////////
	//// I/O
	/////////////////////////////////////////////////////////////////
	public void store(){
		try {
			BufferedWriter bw = new BufferedWriter(new FileWriter("wldos.txt"));
			
			// new FixedAxis(bins,min,max)
			bw.write("min=" + xAxis().binLowerEdge(0)+"\n");
			bw.write("max=" + xAxis().binUpperEdge(xAxis().bins()-1)+"\n");
			bw.write("bins=" + xAxis().bins()+"\n");
			bw.write("SE=[");
			for (int i=0; i<SE.length; i++){
				bw.write(Double.toString(SE[i]));
				bw.write(" ");
			}
			bw.write("]\n");
			bw.close();
			System.out.println("WLDOS stored !");
			
		}
		catch (Exception e){
			e.printStackTrace();
		}
	}
}
