package svm;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map.Entry;
import java.util.stream.Collectors;

import solver.OJAlgoSolver;
import base.AtomicKernel;
import base.Example;
import base.Kernel;
import base.OutputKernel;
import base.SparseVector;
import base.Structure;

public class SVMStructKernel<X extends Structure, Y extends Structure> implements Serializable
{
	/**
	 * 
	 */
	private static final long serialVersionUID = -3805778215058778725L;

	public Mode m;
	public double C;
	public double eps;
	public int oracleCacheSize = 0;
	public SVMStructKernelInstance<X, Y> instance;
	public Kernel<X> inputKernel;
	public OutputKernel<X, Y> outputKernel;
	public double slack;
	public ArrayList<DualConstraint<X, Y>> constraints = new ArrayList<DualConstraint<X, Y>>();
	public ArrayList<Double> alphas = new ArrayList<Double>();
	public List<Example<X, Y>> data;
	public LinkedList<Integer>[] oracleCaches;
	public AtomicConstraint<X, Y> w;

	public SVMStructKernel(SVMStructKernelInstance<X, Y> instance, double C, double eps, Mode m)
	{
		this.instance = instance;
		this.inputKernel = instance.inputKernel();
		this.outputKernel = instance.outputKernel();
		this.m = m;
		this.C = C;
		this.eps = eps;

		// super(instance, C, eps, m);
		// TODO Auto-generated constructor stub
	}

	public double wTimesPhi(X x, Y y)
	{
		if (this.outputKernel instanceof AtomicKernel)
		{
			double sum = 0;

			return sum;
		}
		else
			return data.parallelStream().mapToDouble((xy) -> {
				double sum = 0;
				double kx = this.inputKernel.k(x, xy.x);
				for (int i = 0; i < alphas.size(); i++)
				{
					if (Math.abs(alphas.get(i)) < 0.00001)
						continue;
					if (Math.abs(alphas.get(i) * kx) < 0.000001)
						continue;
					DualConstraint<X, Y> c = constraints.get(i);
					sum += alphas.get(i) * kx * this.outputKernel.k(y, xy.y);
					sum -= alphas.get(i) * kx * this.outputKernel.k(y, c.ybar.get(xy.i));
				}
				return sum;
			}).average().getAsDouble();
	}

	public double wTimesPhiMinusPhi(X x, Y y, Y ybar)
	{
		if (this.outputKernel instanceof AtomicKernel)
		{
			double sum = 0;
			SparseVector phi = ((AtomicKernel<X, Y>) outputKernel).featureMap(x, y);
			SparseVector phiBar = ((AtomicKernel<X, Y>) outputKernel).featureMap(x, ybar);

			for (int i = 0; i < alphas.size(); i++)
			{
				AtomicConstraint<X, Y> atom = (AtomicConstraint<X, Y>) constraints.get(i);
				for (Entry<String, Double> e : phi.entrySet())
				{
					if (atom.phi.containsKey(e.getKey()))
						for (Entry<X, Double> f : atom.phi.get(e.getKey()).entrySet())
						{
							sum += alphas.get(i) * e.getValue() * f.getValue() * this.inputKernel.k(f.getKey(), x);
						}
				}
				for (Entry<String, Double> e : phiBar.entrySet())
				{
					if (atom.phi.containsKey(e.getKey()))
						for (Entry<X, Double> f : atom.phi.get(e.getKey()).entrySet())
						{
							sum -= alphas.get(i) * e.getValue() * f.getValue() * this.inputKernel.k(f.getKey(), x);
						}
				}
			}
			return sum;
		}
		else
			return data.parallelStream().mapToDouble((xy) -> {
				double sum = 0;
				double kx = this.inputKernel.k(x, xy.x);
				for (int i = 0; i < alphas.size(); i++)
				{
					if (Math.abs(alphas.get(i)) < 0.00001)
						continue;
					if (Math.abs(alphas.get(i) * kx) < 0.000001)
						continue;
					DualConstraint<X, Y> c = constraints.get(i);
					double a = alphas.get(i);
					sum += a * kx * this.outputKernel.k(y, xy.y);
					sum -= a * kx * this.outputKernel.k(ybar, xy.y);
					sum -= a * kx * this.outputKernel.k(y, c.ybar.get(xy.i));
					sum += a * kx * this.outputKernel.k(ybar, c.ybar.get(xy.i));
				}
				return sum;
			}).average().getAsDouble();
	}

	@SuppressWarnings("unchecked")
	public void train(List<Example<X, Y>> training)
	{
		System.out.println("C: " + this.C + "\tEpsilon: " + this.eps);
		if (outputKernel instanceof AtomicKernel)
			System.out.println("Using Atomic Output Kernel and Atomic Constraints");
		data = training;
		System.out.println("Number of Training Examples: " + training.size());
		System.out.println("OJALGOSOLVER");
		DualQPSolver<X, Y> solver = new OJAlgoSolver<X, Y>(m);// new GurobiDual<X, Y>(m);//
		constraints = new ArrayList<DualConstraint<X, Y>>();
		alphas = new ArrayList<Double>();
		this.oracleCaches = new LinkedList[training.size()];
		for (int i = 0; i < training.size(); i++)
			this.oracleCaches[i] = new LinkedList<Integer>();

		int iter = 0;
		while (true)
		{
			iter++;
			System.out.print(iter + "\n");
			// System.out.println((iter)+". Iteration");
			System.out.println("Number of Constraints: " + constraints.size());
			if (constraints.size() > 0)
			{
				System.out.println("SOLVING");
				slack = solver.solve(constraints, data, alphas, C, this, iter);
				if (this.outputKernel instanceof AtomicKernel)
				{
					System.out.println("COMPRESSING w");
					int n = 0;
					int z = 0;
					this.w = new AtomicConstraint<X, Y>((AtomicKernel<X, Y>) this.outputKernel);
					for (int i = 0; i < constraints.size(); i++)
					{
						double alphai = this.alphas.get(i);
						if (alphai > 0.000001)
						{

							AtomicConstraint<X, Y> ci = (AtomicConstraint<X, Y>) constraints.get(i);
							for (Entry<String, HashMap<X, Double>> e : ci.phi.entrySet())
							{
								if (!this.w.phi.containsKey(e.getKey()))
								{
									this.w.phi.put(e.getKey(), new HashMap<X, Double>());
								}
								HashMap<X, Double> bla = this.w.phi.get(e.getKey());
								for (Entry<X, Double> f : e.getValue().entrySet())
								{
									n++;
									z++;
									Double b = bla.put(f.getKey(), f.getValue() * alphai);
									if (b != null)
									{
										z--;
										bla.put(f.getKey(), b + f.getValue() * alphai);
									}
								}

							}
						}
					}
					System.out.println("Compression Rate: " + 1.0 * z / n + "\t Number of Entries: " + z + "\t under " + w.phi.size() + " keys");
				}
				System.out.println("SOLVING DONE");
				int nonzero = 0;
				for (double alpha : alphas)
				{
					if (Math.abs(alpha) >= 0.00001)
						nonzero++;
				}
				System.out.println("Number of Support Vectors: " + nonzero);
			}
			else
			{
				slack = 0;
			}

			// System.out.println(alphas.toString());
			System.out.println("COMPUTING CUTTINGPLANE");

			DualConstraint<X, Y> nc;

			if (outputKernel instanceof AtomicKernel)
			{
				// AtomicConstraints can do some fancy cashing due to the predefined structure of the atomic kernel.
				nc = new AtomicConstraint<X, Y>((AtomicKernel<X, Y>) outputKernel);
			}
			else
				nc = new DualConstraint<X, Y>();

			// try to find a nice cutting plane doing basically nothing
			double violation = 0;
			for (int i = 0; i < data.size(); i++)
			{
				double vmax = Double.NEGATIVE_INFINITY;
				int jmax = -1;
				for (int j = 0; j < oracleCaches[i].size(); j++)
				{
					Y ybar = this.constraints.get(oracleCaches[i].get(j)).ybar.get(i);
					double l = this.constraints.get(oracleCaches[i].get(j)).losses.get(i);
					double v = l - this.wTimesPhiMinusPhi(data.get(i).x, data.get(i).y, ybar);
					if (v > vmax)
					{
						vmax = v;
						jmax = j;
					}
				}
				if (jmax != -1)
				{
					violation += vmax;
					nc.addConstraint(data.get(i).x, data.get(i).y,
							this.constraints.get(oracleCaches[i].get(jmax)).ybar.get(i), this.m,
							this.constraints.get(oracleCaches[i].get(jmax)).losses.get(i));
					int c = oracleCaches[i].get(jmax);
					oracleCaches[i].remove(jmax);
					oracleCaches[i].addFirst(c);
				}
			}
			violation /= data.size();

			if (violation < this.slack + this.eps || iter == 1)
			{
				System.out.println("Nothing in Oracle Cache");
				if (outputKernel instanceof AtomicKernel)
				{
					// AtomicConstraints can do some fancy cashing due to the predefined structure of the atomic kernel.
					nc = new AtomicConstraint<X, Y>((AtomicKernel<X, Y>) outputKernel);
				}
				else
					nc = new DualConstraint<X, Y>();

				// Parallelize this sucker...Java 8 Magic!
				List<Y> collect = data.parallelStream().map((e) ->
				{
					Y ybar = mostViolatedConstraint(e);
					return ybar;
				}).collect(Collectors.toList());

				System.out.println("COMPUTING CUTTINGPLANE MIDWAY");

				// TODO: parallelize this? Probably not a good idea...with the concurrent modification of the constraint.
				for (int i = 0; i < collect.size(); i++)
				{
					Y ybar = collect.get(i);
					// System.out.println(i);
					Example<X, Y> e = data.get(i);
					double loss = instance.loss(e.y, ybar);
					nc.addConstraint(e.x, e.y, ybar, m, loss);
					oracleCaches[i].add(0, constraints.size());
					if (oracleCaches[i].size() > oracleCacheSize)
						oracleCaches[i].removeLast();
				}

			}
			nc.lastUsed = iter;
			nc.finalizeConstraint(data.size());
			constraints.add(nc);

			// cleanup unused constraints with a crude heuristic
			for (int i = constraints.size() - 2; i >= 0; i--)
			{
				DualConstraint<X, Y> c = constraints.get(i);
				if (iter - c.lastUsed > 50)
				{
					constraints.remove(i);
					alphas.remove(i);
					solver.pruneConstraint(i);
					// System.out.println("cleanup");
				}
			}

			System.out.println("COMPUTING CUTTINGPLANE DONE");
			System.out.println("UPDATING H MATRIX");
			solver.updateH(constraints, training, this);
			System.out.println("Updating H MATRIX DONE");

			double c = 0;
			if (this.m == Mode.MARGIN_RESCALING)
			{
				for (int i = 0; i < alphas.size(); i++)
				{
					c += alphas.get(i) * solver.H.get(i).get(alphas.size());
				}
			}
			else
			{
				for (int i = 0; i < nc.ybar.size(); i++)
				{
					c += nc.losses.get(i) * wTimesPhiMinusPhi(data.get(i).x, data.get(i).y, nc.ybar.get(i));
				}
				c /= data.size();
			}

			double tmp = nc.loss - c;
			System.out.println("Loss: " + nc.loss);
			System.out.println("Margin: " + c);
			System.out.println("L-M: " + tmp);
			System.out.println("Slack: " + slack);

			if (tmp <= slack + eps)
			{
				constraints.remove(constraints.size() - 1);
				System.out.println();
				int nonzero = 0;
				for (double alpha : alphas)
				{
					if (Math.abs(alpha) >= 0.00001)
						nonzero++;
				}
				System.out.println("Number of Support Vectors: " + nonzero);
				System.out.println("Train Error: " + empiricalError(training));

				// System.out.println("L2: " + this.w.innerProduct(this.w, data, this.inputKernel));
				double ww = 0;
				for (int i = 0; i < alphas.size(); i++)
					for (int j = 0; j < alphas.size(); j++)
						ww += alphas.get(i) * alphas.get(j) * solver.H.get(i).get(j);
				System.out.println("L2: " + ww);

				// System.out.println("L2: " + w.multiply(w));
				// System.out.println("Number of selected features: " + w.countNonZero());
				System.out.println("");
				// debugH(solver);
				break;
			}
		}
	}

	private void debugH(DualQPSolver<X, Y> solver)
	{
		for (int i = 0; i < solver.H.size() - 1; i++)
			for (int j = 0; j < solver.H.size() - 1; j++)
			{
				double loss = 0;
				for (int k = 0; k < constraints.get(i).ybar.size(); k++)
					loss += this.instance.loss(constraints.get(i).ybar.get(k), constraints.get(j).ybar.get(k));
				loss /= constraints.get(i).ybar.size();
				System.out.println(solver.H.get(i).get(j) + "\t" + loss);
			}

	}

	public double empiricalError(List<Example<X, Y>> training)
	{

		return training.parallelStream().mapToDouble((e) ->
		{
			Y ybar = predict(e.x);
			double l = instance.loss(e.y, ybar);// 0.5 * (outputKernel.k(e.y, e.y) - 2 * outputKernel.k(e.y, ybar) + outputKernel.k(ybar, ybar));
				return l;
			}).average().getAsDouble();

	}

	public Y predict(X x)
	{
		if (this.m == Mode.MARGIN_RESCALING)
		{
			if (this.outputKernel instanceof AtomicKernel)
			{
				return ((AtomicKernel<X, Y>) outputKernel).preImage(this.w, x, this.inputKernel);
			}
			else
			{
				HashMap<Y, Double> lin = new HashMap<>();
				for (int j = 0; j < data.size(); j++)
				{
					double kx = this.inputKernel.k(x, data.get(j).x) / data.size();
					for (int i = 0; i < alphas.size(); i++)
					{
						if (Math.abs(alphas.get(i)) < 0.00001)
							continue;
						if (Math.abs(alphas.get(i) * kx) < 0.000001)
							continue;
						DualConstraint<X, Y> c = constraints.get(i);
						Double prev = lin.put(data.get(j).y, alphas.get(i) * kx);
						if (prev != null)
							lin.put(data.get(j).y, alphas.get(i) * kx + prev);
						prev = lin.put(c.ybar.get(j), -alphas.get(i) * kx);
						if (prev != null)
							lin.put(c.ybar.get(j), -alphas.get(i) * kx + prev);
					}
				}
				return outputKernel.preImage(lin, x);
			}
		}
		else
		{
			throw new RuntimeException("Not yet implemented");
		}
	}

	private Y mostViolatedConstraint(Example<X, Y> e)
	{
		X x = e.x;
		Y y = e.y;
		if (m == Mode.MARGIN_RESCALING)
		{
			if (this.outputKernel instanceof AtomicKernel)
			{
				return ((AtomicKernel<X, Y>) outputKernel).preImageL2Penelized(this.w, y, x, this.inputKernel);
			}
			else
			{
				HashMap<Y, Double> lin = new HashMap<>();
				for (int j = 0; j < data.size(); j++)
				{
					double kx = this.inputKernel.k(x, data.get(j).x) / data.size();
					double b = 0;
					for (int i = 0; i < alphas.size(); i++)
					{
						if (Math.abs(alphas.get(i)) < 0.00001)
							continue;
						if (Math.abs(alphas.get(i) * kx) < 0.000001)
							continue;
						DualConstraint<X, Y> c = constraints.get(i);
						b += alphas.get(i) * kx;
						Double prev = lin.put(c.ybar.get(j), -alphas.get(i) * kx);
						if (prev != null)
							lin.put(c.ybar.get(j), -alphas.get(i) * kx + prev);
					}
					Double prev = lin.put(data.get(j).y, b);
					if (prev != null)
						lin.put(data.get(j).y, b + prev);
				}

				return outputKernel.preImagePenelized(lin, y, x);
			}
		}
		else
			throw new RuntimeException("Not yet implemented");
	}

}
