package ec.gp.semantic.geometry;

import ec.EvolutionState;
import ec.app.semanticGP.func.numeric.Constant;
import ec.app.semanticGP.func.numeric.Mul;
import ec.app.semanticGP.func.numeric.Sum;
import ec.app.semanticGP.func.numeric.Xn;
import ec.gp.GPNode;
import ec.gp.semantic.ISemantics;
import ec.gp.semantic.func.SimpleNodeBase;
import ec.gp.semantic.utils.InstructionJoiner;
import gnu.trove.set.hash.TCustomHashSet;
import gnu.trove.set.hash.THashSet;
import gnu.trove.strategy.HashingStrategy;
import library.semantics.TestCase;
import org.apache.commons.math3.linear.*;

import java.lang.Double;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;

/**
 * @author Tomasz Pawlak
 */
public class PolynomialInterpolator implements IInterpolator<Double> {

	/**
	 * used to detect for syntactically equivalent programs
	 */
	/*TCustomHashSet<SimpleNodeBase<?>> duplicateDetector = new TCustomHashSet<>(new HashingStrategy<SimpleNodeBase<?>>() {
		@Override
		public int computeHashCode(SimpleNodeBase<?> simpleNodeBase) {
			return simpleNodeBase.rootedTreeHashCode();
		}

		@Override
		public boolean equals(SimpleNodeBase<?> simpleNodeBase, SimpleNodeBase<?> t1) {
			return simpleNodeBase.rootedTreeEquals(t1);
		}
	});*/
	HashSet<ISemantics> duplicateDetector = new HashSet<>();

	@Override
	public SimpleNodeBase<Double> getProgram(EvolutionState state, List<TestCase<Double>> points) {

		// returns program in form
		// a1*P1(x1, x2,...) + a2*P2(x1, x2,...) + ... + an*Pn(x1, x2,...)
		// Pis are base polynomials constructed by iteratively multiplying consecutive variables
		// ais are coefficients to be determined by solving system of linear equations

		assert points.size() > 0;

		// number of examples in data
		final int examples = points.size();
		// number of variables in each example
		final int variables = points.get(0).getArguments().length;
		// prototypes of variable instructions
		final Xn[] xs = new Xn[variables];
		// base polynomials P
		final SimpleNodeBase[] Ps = new SimpleNodeBase[examples];
		// values of base polynomials in example points
		final RealMatrix A = new Array2DRowRealMatrix(examples, examples);
		// vector of output values of points
		final RealVector b = new ArrayRealVector(examples);

		// initialize prototypes
		for (int i = 0; i < variables; ++i) {
			xs[i] = Xn.get((byte)i);
		}

		// initialize base polynomials and fill matrix of system of linear equations with values of Ps in given points
		this.initStructures(Ps, xs, A, b, points);

		// solve system of linear equations
		// QRDecomposition should be right choice if number of variables is equal to number of equations
		RealVector weights = new QRDecomposition(A).getSolver().solve(b);

		// construct final program
		List<SimpleNodeBase<?>> toJoin = new ArrayList<>(examples);
		for (int i = 0; i < examples; ++i) {
			double w = weights.getEntry(i);
			if (Math.abs(w) > 1E-15) {
				Mul mul = new Mul();
				mul.children = new GPNode[]{Ps[i], new Constant(w)};
				toJoin.add(mul);
			}
		}

		SimpleNodeBase<?> program = InstructionJoiner.join(toJoin, Sum.class, 2);
		return (SimpleNodeBase<Double>) program;
	}

	private void initStructures(final SimpleNodeBase[] Ps, final Xn[] xs, final RealMatrix matrix, final RealVector b, final List<TestCase<Double>> points) {
		try {
			int fillIndex = 0;
			for (int i = 0; i < xs.length; ++i) {

				Ps[fillIndex] = xs[i];
				for (int r = 0; r < points.size(); ++r) {
					TestCase<Double> tc = points.get(r);
					double arg = tc.getArguments()[i];
					matrix.setEntry(r, fillIndex, arg);
					b.setEntry(r, tc.getValue());
				}

				if (++fillIndex >= Ps.length) {
					throw new BreakException();
				}
			}

			this.initStructures(Ps, xs, matrix, points, 0, fillIndex);
		} catch (BreakException e) {
		}
	}

	private void initStructures(final SimpleNodeBase[] Ps, final Xn[] xs, final RealMatrix matrix, final List<TestCase<Double>> points, final int from, int fillIndex) throws BreakException {

		int to = fillIndex;
		for (int curr = from; curr < to; ++curr) {
			instructionLoop: for (int i = 0; i < xs.length; ++i) {
				SimpleNodeBase<Double> copy = (SimpleNodeBase<Double>) Ps[curr].clone();
				Mul mul = new Mul();
				mul.children = new GPNode[]{copy, xs[i]};

				ISemantics mulSemantics = mul.getSemantics();
				if (this.duplicateDetector.contains(mulSemantics)) {
					continue;
				}

				for (ISemantics other : this.duplicateDetector) {
					if (mulSemantics.distanceTo(other) < 1.0) {
						continue instructionLoop;
					}
				}

				this.duplicateDetector.add(mulSemantics);
				Ps[fillIndex] = mul;

				for (int r = 0; r < points.size(); ++r) {
					double arg = matrix.getEntry(r, curr) * points.get(r).getArguments()[i];
					matrix.setEntry(r, fillIndex, arg);
				}

				if (++fillIndex >= Ps.length) {
					this.duplicateDetector.clear();
					throw new BreakException();
				}
			}
		}

		this.duplicateDetector.clear();
		this.initStructures(Ps, xs, matrix, points, to, fillIndex);
	}

	class BreakException extends Exception {
	}
}
