module gmle.salsa.framework;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.File;

behavior GradientDescent {
	int jump = 2;
	int nquad = 3;
	double targ = 0.1;
	double tolgrad = 0.00000001;
	double tolsearch = 0.000000001;
	double tol = tolgrad;
	double tols = tolsearch;
	int bold = 10;
	int nhessian = 10000;
	double lam = 0.1;
	double tolf = tol * tols;
	int max_line_searches = 10000;

	long start_time;

	String search_type;

	BufferedWriter evaluation_file;
	MLEvaluator evaluator;

	void set_evaluation_output(String filename) {
		try {
			System.err.println("setting evaluation output file to: " + filename);
			this.evaluation_file = new BufferedWriter(new FileWriter(new File(filename)));
		} catch (Exception e) {
			System.err.println("Unable to open evaluation output file: " + filename);
			System.err.println(e);
			e.printStackTrace();
			System.exit(0);
		}
	}

	int evaluations_done = 0;
	double track_evaluate(Matrix matrix) {
		track_evaluate(matrix.to_array()) @
		currentContinuation;
	}
	double track_evaluate(double[] parameters) {
		token result = evaluator<-evaluate(parameters);
		write_result(result, parameters) @
		currentContinuation;
	}

	double write_result(double result, double[] parameters) {
		if (evaluation_file != null) {
			evaluations_done++;
			try {
				evaluation_file.write(evaluations_done + " : " + result + " :");
				for (int i = 0; i < parameters.length; i++) {
					evaluation_file.write(" " + parameters[i]);
				}
				evaluation_file.write("\n");
				evaluation_file.flush();
			} catch (Exception e) {
				System.err.println("Unable to write to evaluation output file");
				System.err.println(e);
				e.printStackTrace();
				System.exit(0);
			}
		} else {
			evaluations_done++;
		}
		double evals_per_sec = ((double)evaluations_done) / (((double)System.currentTimeMillis()-(double)start_time)/1000.0);
		System.out.println(evaluations_done + ": evals/sec: " + evals_per_sec + ", result: " + result);

		return result;
	}


	public void initialize(Object[] init_params) {
		evaluator<-initialize(init_params) @ currentContinuation;
	}

	public GradientDescent(String evaluatorClass, String theaters_file) {
		this.evaluator = new MLEvaluator(evaluatorClass, theaters_file);
	}

	public GradientDescent(String evaluatorClass, int number_actors) {
		this.evaluator = new MLEvaluator(evaluatorClass, number_actors);
	}

	double[] step_sizes;

	double[] gradient_descent(double[] initial_parameters, double[] step_sizes) {
		search_type = "gd";
		this.step_sizes = step_sizes;
		
		descent(initial_parameters) @
		currentContinuation;
	}

	double[] conjugate_gradient_descent(double[] initial_parameters, double[] step_sizes) {
		search_type = "cgd";
		this.step_sizes = step_sizes;

		descent(initial_parameters) @
		currentContinuation;
	}

	double[] hessian_conjugate_gradient_descent(double[] initial_parameters, double[] step_sizes) {
		search_type = "cghd";
		this.step_sizes = step_sizes;

		descent(initial_parameters) @
		currentContinuation;
	}


	Matrix get_gradient(double[] parameters) {
		incoming_gradient = new double[parameters.length*2];

		System.err.println("getting gradient of size: " + incoming_gradient.length);

		double[][] parameter_copies = new double[parameters.length*2][parameters.length];
		for (int i = 0; i < parameter_copies.length; i++) {
			for (int j = 0; j < parameter_copies[i].length; j++) {
				parameter_copies[i][j] = parameters[j];
			}
		}

		token wait_token = dummy();
		for (int i = 0; i < parameters.length; i++) {
			System.err.println("setting copies: " + (i*2) + " and " + ((i*2)+1));

			parameter_copies[i*2][i] += step_sizes[i];
			parameter_copies[(i*2)+1][i] -= step_sizes[i];

			track_evaluate(parameter_copies[i*2]) : waitfor(wait_token) @
			set_gradient(new Integer(i*2), token) @
			track_evaluate(parameter_copies[(i*2)+1]) @
			wait_token = set_gradient(new Integer((i*2)+1), token);
		}

		combine_gradient() : waitfor(wait_token) @
		currentContinuation;
	}

	void dummy() {}

	double[] incoming_gradient;
	void set_gradient(int position, double value) {
		incoming_gradient[position] = value;
	}

	Matrix combine_gradient() {
		Matrix gradient = new Matrix(incoming_gradient.length/2, 1);

		System.err.println("calculated gradient:");
		for (int i = 0; i < incoming_gradient.length/2; i++) {
			gradient.values[i][0] = (incoming_gradient[2*i]-incoming_gradient[(2*i)+1])/(2.0*step_sizes[i]);
			System.err.println("\tgradient.values[" + i + "]: " + gradient.values[i][0]);			
		}
		return gradient;
	}

	int num_line_searches = 0;

	double d_elf, f_prev;
	Matrix h, h_prev, d_pres, g_prev, d_prev, x_prev, d_elg, d_elx;

	double[] descent(double[] initial_parameters) {
		g_prev = new Matrix(initial_parameters.length, 1);
		x_prev = new Matrix(initial_parameters.length, 1);
		d_prev = new Matrix(initial_parameters.length, 1);
		h_prev = new Matrix(initial_parameters.length, initial_parameters.length);
		for (int i = 0; i < initial_parameters.length; i++) {
			h_prev.values[i][i] = 1.0;
		}

		Matrix x_pres = new Matrix(initial_parameters.length, 1);
		for (int i = 0; i < initial_parameters.length; i++) {
			x_pres.values[i][0] = initial_parameters[i];
		}
		d_elg = new Matrix(initial_parameters.length, 1);
		d_elx = new Matrix(initial_parameters.length, 1);

		start_time = System.currentTimeMillis();

		f_prev = 0.0;
		num_line_searches = 0;
		token f_pres = track_evaluate(initial_parameters);
		descent_loop(f_pres, x_pres) @
		currentContinuation;
	}

	double[] descent_loop(double f_pres, Matrix x_pres) {
		System.err.print("iteration: " + num_line_searches + ", f_pres: " + f_pres + ", x_pres: ");
		x_pres.print();

		if (num_line_searches >= max_line_searches) {
			return x_pres.to_array();
		} else {
			num_line_searches++;
		}
		token g_pres = get_gradient(x_pres.to_array());
		descent_loop2(new Double(f_pres), x_pres, g_pres) @
		currentContinuation;
	}

	double[] descent_loop2(double f_pres, Matrix x_pres, Matrix g_pres) {
		d_elg = g_pres.sub(g_prev);
		d_elx = x_pres.sub(x_prev);
		d_elf = f_pres - f_prev;

		System.err.print("gradient: " );
		g_pres.print();

		if (search_type.equals("cghd")) {
			//t1 = h_prev * d_elg;
			Matrix t1 = h_prev.mul(d_elg);

			//t2 = d_elg' * t1;
			Matrix t2 = d_elg.invert().mul(t1);

			//t3 = d_elx' * d_elg;
			Matrix t3 = d_elx.invert().mul(d_elg);

			//u = d_elx/t3 - t1/t2;
			Matrix temp1 = d_elx.rdiv(t3);
			Matrix temp2 = t1.rdiv(t2);
			Matrix u = temp1.sub(temp2);

			//h = h_prev + (d_elx*d_elx')/t3 - (t1*t1')/t2 + t2*(u*u');
			h = h_prev.add(d_elx.mul(d_elx.invert()).rdiv(t3).sub(t1.mul(t1.invert()).rdiv(t2)).add(t2.mul(u.mul(u.invert()))));

			double hm = h.mean();
			if (Double.isNaN(hm)) {
				System.err.println("HM is NAN, quitting\n");
				System.exit(0);
			}
			h_prev = h;
		}

		double ng = 0.0;
		double step = 0.0;
		if (search_type.equals("gd")) {
			d_pres = g_pres.neg();
			step = 0.5;
			ng = g_pres.norm();
		} else if (search_type.equals("cgd")) {
			Matrix bet;
			if (g_prev.norm()  == 0) {
				bet = new Matrix(1,1);
			} else {
				//bet = g_pres'*(g_pres-g_prev)/(g_prev'*g_prev);
				bet = g_pres.invert().mul(g_pres.sub(g_prev)).rdiv(g_prev.invert().mul(g_prev));
			}

			//d_pres = -g_pres + bet*d_prev;
			d_pres = g_pres.neg().add(bet.mul(d_prev));

			ng = g_pres.norm();
//			step = Math.abs(f_pres - targ)/ng;
			step = 0.5;
		} else if (search_type.equals("cghd")) {
			//bet=gpres'*delg/(gprev'*gprev);
			Matrix bet = g_pres.invert().mul(d_elg.rdiv(g_prev.invert().mul(g_prev)));
			if (Double.isInfinite(bet.values[0][0])) bet.values[0][0] = 0;

			//d1=-gpres+bet*dprev;
			//nd1=norm(d1);
			Matrix d1 = g_pres.neg().add(bet.mul(d_prev));
			double nd1 = d1.norm();

			//d2=-H*gpres;
			//nd2=norm(d2);
			Matrix d2 = h.neg().mul(g_pres);
			double nd2 = d2.norm();

			//dpres=(lam*d1/nd1+d2/nd2)/(lam+1);
			d_pres = d1.mul(lam/nd1).add(d2.rdiv(nd2)).rdiv(lam+1);

			ng = g_pres.norm();
			step = (lam*Math.abs(f_pres-targ)/ng+nd2)/(lam+1);
		}

		token x_fut = null;
		if (ng > tol) {
			System.err.print("direction: " );
			d_pres.print();
			x_fut = line_search(x_pres, d_pres, new Double(step));
		} else if (Double.isNaN(ng-ng)) {
			System.err.println("ng is NaN, quitting\n");
			return x_pres.to_array();
		} else {
			System.err.println("gradient is below tolgrad, quitting\n");
			return x_pres.to_array();
		}

		token f_fut = track_evaluate(x_fut);
		if (search_type.equals("cghd")) {
			if (d_elf >= -tolf) {
				lam = Math.pow(bold,10);
				nhessian = num_line_searches+5;
			} else {
				if (lam > Math.pow(bold,-10)) {
					lam = lam/bold;
				}
			}
		}
		d_prev = d_pres;
		g_prev = g_pres;
		x_prev = x_pres;
		f_prev = f_pres;

		descent_loop(f_fut, x_fut) @
		currentContinuation;
	}

	Matrix d;
	double step;
	double e1, e2, e3;
	Matrix x0, x1, x2, x3;
	Matrix line_search(Matrix x0, Matrix d_ser, double step) {
		this.step = step;
		this.x0 = x0;
		d = d_ser.rdiv(d_ser.norm());

		token e1_token = track_evaluate(x0.to_array());
		token e2_token = track_evaluate(x0.add(d.mul(step)).to_array());
		line_search_loop1(e1_token, e2_token) @
		currentContinuation;
	}

	double l1, l2, l3;
	Matrix line_search_loop1(double e1_p, double e2_p) {
		this.e1 = e1_p;
		this.e2 = e2_p;

		if (e2 < e1) {
			l1 = 0;
			l2 = 1;
			l3 = 2;
			token e1_token = track_evaluate( x0.add(d.mul(l1*step)) );
			token e2_token = track_evaluate( x0.add(d.mul(l2*step)) );
			token e3_token = track_evaluate( x0.add(d.mul(l3*step)) );
			line_search_loop2(e1_token, e2_token, e3_token) @ currentContinuation;
		} else {
			if (step < tol) {
				System.err.println("Line search could not progress beyond tolsearch\n");
				System.exit(0);
			}
			step /= 2.0;
		}
		token e2_token = track_evaluate(x0.add(d.mul(step)).to_array());
		line_search_loop1(new Double(e1), e2_token) @
		currentContinuation;
	}

	Matrix line_search_loop2(double e1_p, double e2_p, double e3_p) {
		this.e1 = e1_p;
		this.e2 = e2_p;
		this.e3 = e3_p;

		if (e3 >= e2) {
			loop3_iterations = 1;
			line_search_loop3() @
			currentContinuation;
		} else if (Double.isNaN(e3-e3)) {
			System.err.println("premature exit of linesearch due to NaN\n");
			exit(0);
		} else {
			l1 = l2;
			e1 = e2;

			l2 = l3;
			e2 = e3;

			l3 = jump * l3;
			token e3_token = track_evaluate( x0.add(d.mul(l3*step)) );
			line_search_loop2(new Double(e1), new Double(e2), e3_token) @
			currentContinuation;
		}
		return null;
	}

	double a, b, es, lstar;
	int loop3_iterations;
	Matrix x;
	Matrix line_search_loop3() {	
		if (loop3_iterations >= nquad) {
			return x0.add(d.mul(lstar*step));
		}

		loop3_iterations++;

		a = (l1*l1)*(e2-e3) + (l2*l2)*(e3-e1) + (l3*l3)*(e1-e2);
		b = l1*(e2-e3) + l2*(e3-e1) + l3*(e1-e2);
		lstar = 0.5*(a/b);	

		token es_token = track_evaluate( x0.add(d.mul(lstar*step)) );
		line_search_loop3_body(es_token) @
		currentContinuation;
	}

	Matrix line_search_loop3_body(double es_p) {
		this.es = es_p;

		if (Double.isNaN(es)) {
			System.err.println("premature exit of linesearch due to NaN 2\n");
			System.exit(0);
		}

		if (lstar > l2 + tol) {
			if (es > e2) {
				l3 = lstar;
				e3 = es;
			} else {
				l1 = l2;
				e1 = e2;
				l2 = lstar;
				e2 = es;
			}
		} else if (lstar < l2 - tol) {
			if (es > e2) {
				l1 = lstar;
				e1 = es;
			} else {
				l3 = l2;
				e3 = e2;
				l2 = lstar;
				e2 = es;
				}
		} else {
			lstar = l2 + tol;
			//es = track_evaluate(x0 + lstar*step*d);
			token es_token = track_evaluate( x0.add(d.mul(lstar*step)) );

			line_search_loop3a(es_token) @
			currentContinuation;
		}
		line_search_loop3() @
		currentContinuation;
	}

	Matrix line_search_loop3a(double es_p) {
		this.es = es_p;
	
		if (es < e2) {
			l1 = l2;
			e1 = e2;
			l2 = lstar;
			e2 = es;
		} else {
			l3 = lstar;
			e3 = es;
		}
		line_search_loop3() @
		currentContinuation;
	}
}