/*Uses the Roe solver to solve the formulation of an ultrarelativistic fluid described 
  in the second project here. Again depends on StdDraw.java.*/

public class UltraRel {

	//holds all grid functions: row 0 = rho, row 1 = v, row 2 = tau, row 3 = S
	double[][] vars;
	//toggles Gaussian or shock tube initial data
	boolean smooth;
	//number of time steps
	int numsteps;
	//length of time steps
	double delta_t;
	
	//number of cells
	int Nx;
	//number of ghost cells at each end
	int Ng = 2;
	
	//domain
	double xmin = 0;
	double xmax = 5;
	//distance between cells
	double delta_x;
	//holds positions of cell centers
	double[] x;
	
	//adiabatic index
	double k = 1.5;
	//conservative variables must always be above floor
	double floor = 1e-11;

	public UltraRel(int Nx, double dt, boolean smooth) {
		this.smooth = smooth;
		delta_t = dt;
		this.Nx = Nx;
		delta_x = (xmax - xmin)/Nx;
		x = new double[Nx];
		vars = new double[4][Nx];
		for (int i = 0; i < Nx; i++) {
			x[i] = xmin + (i + .5)*delta_x;
		}
		
	}
	
	public void init_fluid() {

		//if smooth = true, Gaussian initial data, otherwise shock tube
		if (smooth == true) {
			for (int i = 0; i < Nx; i++) {
				vars[0][i] = Math.exp(-(x[i] - 2.5)*(x[i] - 2.5));
			}
		}
		
		else {
			//shock tube
			double rho_L = 1.0;
			double rho_R = 0.1;
			for (int i = 0; i < Nx; i++) {
				if (i < Nx/2) vars[0][i] = rho_L;
				else vars[0][i] = rho_R;
			}
		}
		
		double v_L = 0;
		double v_R = 0;
		for (int i = 0; i < Nx; i++) {
			if (i < Nx/2) vars[1][i] = v_L;
			else vars[1][i] = v_R;
		}
		double[] P = new double[Nx];
		for (int i = 0; i < Nx; i++) {
			P[i] = (k - 1) * vars[0][i];
		}
		
		for (int i = 0; i < Nx; i++) {
			vars[2][i] = (vars[0][i] + P[i])/(1 - vars[1][i]*vars[1][i]) - P[i];
		}
		
		for (int i = 0; i < Nx; i++) {
			vars[3][i] = vars[1][i]*(vars[2][i] + P[i])/(1 - vars[1][i]*vars[1][i]);
		}
		
	}
	
	//minmod function
	private double minmod(double a, double b) {
		if (a*b < 0) return 0;
		else if(Math.abs(a) < Math.abs(b)) return a;
		else return b;
	}
	
	//calculating Jacobian matrix at x
	private double[][] calc_A(double tau, double S, double rho, double v) {
		double beta = (2 - k)/4;
		double P_tau = -2*beta + (4*beta*beta + k - 1)*tau/Math.sqrt(4*beta*beta*tau*tau + (k - 1)*(tau*tau - S*S));
		double P_s = -(k - 1)*S/Math.sqrt(4*beta*beta*tau*tau + (k - 1)*(tau*tau - S*S));
		double[][] A = new double[2][2];
		A[0][0] = 0;
		A[0][1] = 1;
		A[1][0] = -v*v + (1 - v*v)*P_tau;
		A[1][1] = 2*v + (1 - v*v)*P_s;
		return A;
	}
	
	//calculating the first eigenvector of Jacobian
	private double[] calc_eta1(double[][] A) {
		double[] eta1 = new double[2];
		double lambda1 = (A[0][0] + A[1][1] + Math.sqrt((A[0][0] - A[1][1])*(A[0][0] - A[1][1]) + 4*A[0][1]*A[1][0]))/2;
		eta1[0] = 1;
		eta1[1] = (lambda1 - A[0][0])/A[0][1];
		return eta1;
	}
	
	//calculating the other eigenvector
	private double[] calc_eta2(double[][] A) {
		double[] eta2 = new double[2];
		double lambda2 = (A[0][0] + A[1][1] - Math.sqrt((A[0][0] - A[1][1])*(A[0][0] - A[1][1]) + 4*A[0][1]*A[1][0]))/2;
		eta2[0] = 1;
		eta2[1] = (lambda2 - A[0][0])/A[0][1];
		return eta2;
	}
	
	//returns both eigenvalues
	private double[] calc_lambda(double[][] A) {
		double[] lambda = new double[2];
		lambda[0] = (A[0][0] + A[1][1] + Math.sqrt((A[0][0] - A[1][1])*(A[0][0] - A[1][1]) + 4*A[0][1]*A[1][0]))/2;
		lambda[1] = (A[0][0] + A[1][1] - Math.sqrt((A[0][0] - A[1][1])*(A[0][0] - A[1][1]) + 4*A[0][1]*A[1][0]))/2;
		return lambda;
	}
	
	//calculates jumps in both conservative variables
	private double[] calc_omega(double tauL, double sL, double tauR, double sR, double[] eta1, double[] eta2) {
		double[] omega = new double[2];
		omega[0] = (eta2[1]*(tauR - tauL) + eta2[0]*(sL - sR))/(eta1[0]*eta2[1] - eta1[1]*eta2[0]);
		omega[1] = (eta1[0]*(sR - sL) + eta1[1]*(tauL - tauR))/(eta1[0]*eta2[1] - eta1[1]*eta2[0]);
		return omega;
	}

	//calculate Roe flux
	private double[] calc_FF(double[] fL, double[] fR, double[] lambda, double[] eta1, double[] eta2, double[] omega) {
		double[] FF = new double[2];
		FF[0] = 0.5*(fR[0] + fL[0] - Math.abs(lambda[0])*eta1[0]*omega[0] - Math.abs(lambda[1])*eta2[0]*omega[1]);
		FF[1] = 0.5*(fR[1] + fL[1] - Math.abs(lambda[0])*eta1[1]*omega[0] - Math.abs(lambda[1])*eta2[1]*omega[1]);
		return FF;
	}

	//calculates flux in terms of conservative variables
	private double[] calc_f(double tau, double S, double rho, double v) {
		double[] f = new double[2];
		double P = (k - 1)*rho;
		f[0] = S;
		f[1] = S*v + P;
		return f;
	}

	//calculates conservative variables from primitive variables
	private double[] calc_cons(double rho, double v) {
		double[] cons = new double[2];
		double P = (k - 1)*rho;
		cons[0] = (rho + P)/(1 - v*v) - P;
		cons[1] = v*(cons[0] + P)/(1 - v*v);
		cons[0] = Math.max(cons[0], floor + Math.abs(cons[1]));
		return cons;

	}

	//calculates primitive variables from conservative variables
	private double[] calc_prim(double tau, double S) {
		double[] prim = new double[2];
		double beta = (2 - k)/4;
		double P = -2*beta*tau + Math.sqrt(4*beta*beta*tau*tau + (k - 1)*(tau*tau - S*S));
		prim[0] = P/(k - 1);
		prim[1] = S/(tau + P);
		return prim;
	}



	private double dqL_1(double dq_imh, double dq_iph) {
		return minmod(dq_iph, dq_imh);
	}
	
	private double dqR_1(double dq_iph, double dq_ip3h) {
		return minmod(dq_iph, dq_ip3h);
	}
	
	private double dqL_2(double dq_imh, double dq_iph) {
		if (Math.abs(dq_iph) <= 1e-6) return 0;
		else {
			double theta = dq_imh/dq_iph;
			return Math.max(Math.max(0, Math.min(1, 2*theta)), Math.min(2, theta))*dq_iph;
		}
	}
	
	private double dqR_2(double dq_iph, double dq_ip3h) {
		if (Math.abs(dq_iph) <= 1e-6) return 0;
		else {
			double theta = dq_ip3h/dq_iph;
			return Math.max(Math.max(0, Math.min(1, 2*theta)), Math.min(2, theta))*dq_iph;
		}
	}
	
	private double dqL_3(double dq_imh, double dq_iph) {
		if (Math.abs(dq_iph) <= 1e-6) return 0;
		else {
			double theta = dq_imh/dq_iph;
			return Math.max(0, Math.min(Math.min((1 + theta)/2, 2), 2*theta))*dq_iph;
		}
	}
	
	private double dqR_3(double dq_iph, double dq_ip3h) {
		if (Math.abs(dq_iph) <= 1e-6) return 0;
		else {
			double theta = dq_ip3h/dq_iph;
			return Math.max(0, Math.min(Math.min((1 + theta)/2, 2), 2*theta))*dq_iph;
		}
	}
	
	private double dqL_4(double s_imh, double s_iph) {
		if (Math.abs(s_iph) <= 1e-6) return 0;
		else {
			double theta = s_imh/s_iph;
			return s_iph*(theta + Math.abs(theta))/(1 + Math.abs(theta));
		}
	}
	
	private double dqR_4(double s_iph, double s_ip3h) {
		if (Math.abs(s_iph) <= 1e-6) return 0;
		else {
			double theta = s_ip3h/s_iph;
			return s_iph*(theta + Math.abs(theta))/(1 + Math.abs(theta));
		}
	}
	
	//reconstructs primitive variable q from the left
	private double recos_qL(double q_im1, double q_i, double q_ip1, double r_im1, double r_i, double r_ip1) {
		double dq_iph = q_ip1 - q_i;
		double dq_imh = q_i - q_im1;
		double dq_i = dqL_1(dq_imh, dq_iph);
		double rqL_iph = q_i + dq_i*0.5;
		return rqL_iph;
	}

	//reconstructs primitive variable q from the right
	private double recos_qR(double q_i, double q_ip1, double q_ip2, double r_i, double r_ip1, double r_ip2) {
		double dq_iph = q_ip1 - q_i;
		double dq_ip3h = q_ip2 - q_ip1;
		double dq_i = dqR_1(dq_iph, dq_ip3h);
		double rqR_iph = q_ip1 - dq_i*0.5;
		return rqR_iph;
	}
	
	//calculates Roe fluxes at each cell boundary (not including ghost cells)
	private double[][] calc_flux(double[][] vars, double[] x, int Nx, int Ng) {
		double[][] FF = new double[2][Nx];

		for (int i = Ng - 1; i < Nx - Ng; i++) {
			//reconstructing primitive variables
			double rrhoL_iph = recos_qL(vars[0][i-1], vars[0][i], vars[0][i+1], x[i-1], x[i], x[i+1]);
			double rrhoR_iph = recos_qR(vars[0][i], vars[0][i+1], vars[0][i+2], x[i], x[i+1], x[i+2]);
			double rrho_iph = 0.5*(rrhoL_iph + rrhoR_iph);
			double rvL_iph = recos_qL(vars[1][i-1], vars[1][i], vars[1][i+1], x[i-1], x[i], x[i+1]);
			double rvR_iph = recos_qR(vars[1][i], vars[1][i+1], vars[1][i+2], x[i], x[i+1], x[i+2]);
			double rv_iph = 0.5*(rvL_iph + rvR_iph);

			//calculating conservative variables from reconstructed values
			double[] rqL_iph = calc_cons(rrhoL_iph, rvL_iph);
			double[] rqR_iph = calc_cons(rrhoR_iph, rvR_iph);
			double[] rq_iph = calc_cons(rrho_iph, rv_iph);

			//calculating Jacobian
			double[][] A = calc_A(rq_iph[0], rq_iph[1], rrho_iph, rv_iph);
			double[] eta1 = calc_eta1(A);
			double[] eta2 = calc_eta2(A);
			double[] lambda = calc_lambda(A);
			double[] omega = calc_omega(rqL_iph[0], rqL_iph[1], rqR_iph[0], rqR_iph[1], eta1, eta2);

			double[] fR_iph = calc_f(rqR_iph[0], rqR_iph[1], rrhoR_iph, rvR_iph);
			double[] fL_iph = calc_f(rqL_iph[0], rqL_iph[1], rrhoL_iph, rvL_iph);
			double[] FF_iph = calc_FF(fR_iph, fL_iph, lambda, eta1, eta2, omega);
			FF[0][i] = FF_iph[0];
			FF[1][i] = FF_iph[1];
		}

		return FF;
	}
	
	//updating
	private double[][] update_q(double[][] q_n, double[] x, double[][] FF, double delta_t, int Nx, int Ng) {
		double dx = x[1] - x[0];
		double[][] q_np1 = new double[4][Nx];

		for (int i = Ng; i < Nx - Ng; i++) {
			q_np1[2][i] = q_n[2][i] - delta_t*(FF[0][i] - FF[0][i-1])/dx;
			q_np1[3][i] = q_n[3][i] - delta_t*(FF[1][i] - FF[1][i-1])/dx;
			q_np1[2][i] = Math.max(q_np1[2][i], Math.abs(q_np1[3][i]) + floor);
		}
		
//		//constant extrapolation
//		for (int i = 0; i < Ng; i++) {
//			q_np1[2][i] = q_np1[2][Ng];
//			q_np1[2][Nx - i - 1] = q_np1[2][Nx - Ng - 1];
//			q_np1[3][i] = q_np1[3][Ng];
//			q_np1[3][Nx - i - 1] = q_np1[3][Nx - Ng - 1];
//		}
		
		//first order extrapolation
		for (int i = 0; i < Ng; i++) {
			q_np1[2][i] = q_np1[2][Ng]  - (Ng - i)*(q_np1[2][Ng + 1] - q_np1[2][Ng]);
			q_np1[2][Nx - i - 1] = q_np1[2][Nx - Ng - 1] + (Ng - i)*(q_np1[2][Nx - Ng - 1] - q_np1[2][Nx - Ng - 2]);
			q_np1[3][i] = q_np1[3][Ng]  - (Ng - i)*(q_np1[3][Ng + 1] - q_np1[3][Ng]);
			q_np1[3][Nx - i - 1] = q_np1[3][Nx - Ng - 1] + (Ng - i)*(q_np1[3][Nx - Ng - 1] - q_np1[3][Nx - Ng - 2]);
		}
		
		for (int i = 0; i < Nx; i++) {
			q_np1[0][i] = calc_prim(q_np1[2][i], q_np1[3][i])[0];
			q_np1[1][i] = calc_prim(q_np1[2][i], q_np1[3][i])[1];

		}
		
		return q_np1;
	}
	
	//moving forward one step in time
	private double[][] step(double[][] var) {
		
		//half advanced time step
		double[][] FF = calc_flux(var, x, Nx, Ng);
		double[][] var_nph = update_q(var, x, FF, delta_t/2, Nx, Ng);
		
		//full advanced time step
		FF = calc_flux(var_nph, x, Nx, Ng);
		double[][] q_np1 = update_q(var_nph, x, FF, delta_t, Nx, Ng);
		
		return q_np1;
		
	}
	
	//shows the solution evolving until time tmax
	public void animate(double tmax) {
		StdDraw.setXscale(xmin, xmax);
		init_fluid();
		numsteps = (int) Math.floor(tmax/delta_t);
		double[][] curr = vars;
		for (int i = 0; i < numsteps; i++) {
			StdDraw.clear();
			for (int j = 0; j < Nx - 1; j++) {
				StdDraw.line(x[j], curr[0][j], x[j+1], curr[0][j+1]);
			}
			
			//drawing axes
			StdDraw.line(xmin, 0, xmax, 0);
			StdDraw.line(0, 0, 0, 1);
			for (int k = (int) xmin; k <= (int) xmax; k++ ) {
				if (k == 0) {}
				else {
					StdDraw.line((double) k, 0, (double) k, 0.01);
					String I = "" + k;
					StdDraw.text((double) k, 0.025, I);
				}
				
			}
			for (int k = 1; k <= 5; k++) {

					StdDraw.line(0, 0.2*k, 0.1, 0.2*k);
					String I = String.format("%.1g%n", 0.2*k);
					StdDraw.text(0.35, 0.2*k, I);
			}
			String filename = "UltraRel" + i+".png";
			StdDraw.save(filename);
			StdDraw.show(30);
			curr = step(curr);
		}
	}
	
	//displays the solution at a time tmax
	public void time(double tmax) {
		init_fluid();
		numsteps = (int) Math.floor(tmax/delta_t);
		double[][] curr = vars;
		for (int i = 0; i < numsteps; i++) {
			curr = step(curr);
		}
		StdDraw.setXscale(xmin, xmax);
		for (int j = 0; j < Nx - 1; j++) {
			StdDraw.line(x[j], curr[0][j], x[j+1], curr[0][j+1]);
		}
		//drawing axes
		StdDraw.line(xmin, 0, xmax, 0);
		StdDraw.line(0, 0, 0, 1);
		for (int i = (int) xmin; i <= (int) xmax; i++ ) {
			if (i == 0) {}
			else {
				StdDraw.line((double) i, 0, (double) i, 0.01);
				String I = "" + i;
				StdDraw.text((double) i, 0.025, I);
			}

		}
		for (int i = 1; i <= 5; i++) {

			StdDraw.line(0, 0.2*i, 0.1, 0.2*i);
			String I = String.format("%.1g%n", 0.2*i);
			StdDraw.text(0.35, 0.2*i, I);
		}
	}
	
	//returns values of rho at each mesh point after time tmax
	public double[] returnrho(double tmax) {
		init_fluid();
		numsteps = (int) Math.floor(tmax/delta_t);
		double[][] curr = vars;
		for (int i = 0; i < numsteps; i++) {
			curr = step(curr);
		}
		return curr[0];
	}
	
	//takes the L2 measure of difference between two arrays a and b of the same length, 
	//w/o contributions from j cells at each end
	public static double norm(double[] a, double[] b, int j) {
		double c = 0;
		for (int i = j; i < a.length - j; i++) {
			c += (a[i] - b[i])*(a[i] - b[i]);
		}
		double norm = Math.sqrt(c/a.length);
		return norm;
	}
	
	//given a starting number of mesh points Nx and Gaussian initial data, calculates "convergence factor" 
	//norm(rho(Nx) - rho(2*Nx))/norm(rho(2*Nx) - rho(4*Nx)) - should be 2^n for an nth order method.
	public static double test(int Nx, double time) {
		double dt = 5.0/(Nx * 20);
		UltraRel one = new UltraRel(Nx, dt, true);
		UltraRel two = new UltraRel(2*Nx, dt, true);
		UltraRel four = new UltraRel(4*Nx, dt, true);
		
		double[] One = one.returnrho(time);
		double[] Two1 = new double[Nx];
		double[] Two2 = new double[2*Nx];
		double[] Four = new double[2*Nx];
		for (int i = 0; i < Nx; i++) {
			Two1[i] = (two.returnrho(time)[2*i] + two.returnrho(time)[(2*i)+1])/2;
			Two2[i] = two.returnrho(time)[i];
			Four[i] = (four.returnrho(time)[2*i] + four.returnrho(time)[(2*i)+1])/2;
		}
		
		double Q = norm(One, Two1, 0)/norm(Two2, Four, 0);
		return Q;
		
		
	}
	
	public static void convergence(int Nx, double time) {
		double dt = 5.0/(Nx * 20);
		UltraRel one = new UltraRel(Nx, dt, true);
		UltraRel two = new UltraRel(2*Nx, dt, true);
		UltraRel four = new UltraRel(4*Nx, dt, true);
		
		double[] One = one.returnrho(time);
		double[] Two = two.returnrho(time);
		double[] Four = four.returnrho(time);
		double[] Q = new double[Nx];
		StdDraw.setXscale(one.xmin, one.xmax);
		StdDraw.setYscale(0, 10);
		//drawing axes
		StdDraw.line(one.xmin, 0, one.xmax, 0);
		StdDraw.line(0, 0, 0, 10);
		for (int i = (int) one.xmin; i <= (int) one.xmax; i++ ) {
			if (i == 0) {}
			else {
				StdDraw.line((double) i, 0, (double) i, 0.1);
				String I = "" + i;
				StdDraw.text((double) i, 0.25, I);
			}
			
		}
		for (int i = 0; i <= 10; i++) {
			if (i == 0) {}
			else {
				StdDraw.line(0, (double) i, 0.1, (double) i);
				String I = "" + i;
				StdDraw.text(0.25, (double) i, I);
			}
		}
		
		StdDraw.setPenRadius(0.005);
		for (int i = 0; i < Nx; i++) {
			double e1 = Math.abs(One[i] - ((Two[2*i] + Two[(2*i)+1])/2));
			double e2 = Math.abs((((Two[2*i] + Two[(2*i)+1])/2) - (Four[4*i] + Four[4*i + 1] + Four[4*i + 2] + Four[4*i + 3])/4));
			Q[i] = e1/e2;
			StdDraw.point(one.x[i], e1/e2);
		}
	}
	
	public static void main(String[] args) {
		UltraRel hi = new UltraRel(200, .01, true);
		hi.animate(3);
	}
}