package ipp.w7x.minerva.neutralBeams;


import java.text.DecimalFormat;

import algorithmrepository.Algorithms;
import jafama.FastMath;
import oneLiners.OneLiners;
import seed.minerva.GraphicalModel;
import seed.minerva.NodeImpl;
import seed.minerva.RandomManager;
import seed.minerva.StateFullNodeImpl;
import seed.minerva.diagnostics.LineOfSightSystem;
import seed.minerva.neutralBeams.BeamDefsNode;
import seed.minerva.neutralBeams.GaussianBeamDefinition;
import seed.minerva.neutralBeams.GaussianBeamDefinitions;

/** Geometry and fuel information for the W7X Neutral Beams
 * 
 * This is largely copied from the AUG definitions, which is nice because they
 * largely copied the beams, fom the AUG beams.
 * 
 * 
 * The coordinates here were very painstakingly worked out from matching what 
 * little information was lying around to some  low red drawings in a presentation 
 * [ 1-AAO-T0018 J.Baldzuhn ] 
 * and the CAD diagrams of the port [ \\Sv1fshgw\plmcatia\..\1-aek20b--.. ]
 * and vessel structure [ \\Sv1fshgw\plmcatia\..\1-abb2--... ]
 * and some bits from the CXRS and BES report [ 1-QSC-T0002 ]
 * 
 * 
 * For the sake of information:
 * 
 * We're using what is hopefully the common coord system where the z axis
 * is vertically upwards, x axis is the middle of Module 1 (1 based, so the first module)
 * and y is half way through the second part of module 2 (i.e. the middle of '21', not '20')
 * 
 * The north wall of hall is at phi = -36deg and the actual north is phi ~ -72 deg.
 * 
 * The heating beams are in the south-west of the hall.
 * 
 * There is an XML file for the beams around, we should probably put that in XBase and use it.
 */
public class W7XHeatingBeams extends StateFullNodeImpl implements GaussianBeamDefinitions, BeamDefsNode {
	public final static double electronCharge = 1.60217646e-19;
	public static final double protonMass = 1.67262158e-27; //kg
	private final static double[] deg2Rad(double deg[]){ for(int i=0; i < deg.length; i++){ deg[i] *= Math.PI / 180; } return deg; }
	private final static double[][] deg2Rad(double deg[][]){ for(int i=0; i < deg.length; i++){ deg[i] = deg2Rad(deg[i]); } return deg; }
	
	/** The 4 beams entering the AEK20 port */
	public static final int BEAMBOX_K20 = 0;
	public static final int BEAMBOX_K21 = 1;
	
	public static final int BEAMFUEL_HYDROGEN = 0;
	public static final int BEAMFUEL_DEUTERIUM = 1;
	
	// Geometry is defined for each  BEAMBOX_* and for each of te 4 beams within a beam box 
	
	/** Beam inclination angle */
	private static final double beta[][] = deg2Rad(new double[][]{ 
				{ 4.9, 4.9, -4.9, -4.9 }, 
				{ -4.9, -4.9, 4.9, 4.9 }, 
			});
	
	private static final double ang = (90 - 18 - 15.15) * Math.PI / 180;	
	/** Angle of R to x (x axis is border between modules 10 and 11) */	
	private static final double theta[][] = deg2Rad(new double[][]{ 
				{ 72-15.15, 72-15.15, 72-15.15, 72-15.15 },	//AEK20
				{ 72+15.15, 72+15.15, 72+15.15, 72+15.15 },	//AEK21
			}); 
	
	/** Angle FROM the reverse box axis (opposite to injection direction)
	 * 	TO R (radial from (0,0) to P0). +ve for anticlockwise */
	private static final double psi[][] = deg2Rad(new double[][]{ 
				{ -7.5, -7.5, -7.5, -7.5 },
				{ +7.5, +7.5, +7.5, +7.5 },
			});
	
	 /** Horizontal angle of beam axis to box axis */
	private static final double alpha[][] = deg2Rad(new double [][]{ 
				{ -4.1, 4.1, 4.1, -4.1 },	
				{ -4.1, 4.1, 4.1, -4.1 },	
			});
	
	 /** Radial position of horzontal beam crossing */
	private static final double R_P0[][] = new double[][]{ 
				{ 6.700, 6.700, 6.700, 6.700 },	
				{ 6.700, 6.700, 6.700, 6.700 },	
			};
	
	/** Height (Z coord) of vertical beam crossing */
	private static final double Z0[][] = new double[][]{
				{ -0.305, -0.305, -0.305, -0.305 },		
				{ 0.305, 0.305, 0.305, 0.305 },		
			};
	
	/** Distance of Z=-0.305 crossing along box axis from P0 
	 * +ve is toward the beam source / away from R=0 */ 
	private static final double Dx[][] = new double[][]{ 
				{ -0.5, -0.5, -0.5, -0.5 },
				{ -0.5, -0.5, -0.5, -0.5 },
			}; 
	
	// Some geometry is just all the same (mostly stuff we made up for now)
	/** FWHM of beam at focus. From the presentation [ 1-AAO-T0018 ], the FWHM looks
	 * roughly 16cm in width and 24cm in height just after the box exit.
	 * The diagrams near the duct exit have all 4 beams merged, so I can't tell.
	 * 
	 * The physics scenarios use "beam radius = 0.1cm"
	 * [ https://validatorwww.ipp-hgw.mpg.de/Wiki/index.php/Physics_Scenarios ] 
	 */
	private static final double fwhmAtFocus = 0.19; 
	
	/** Taken from that used in the physics scenarios page:
	 * "divergence = 0.8degree"  
	 * [ https://validatorwww.ipp-hgw.mpg.de/Wiki/index.php/Physics_Scenarios ] */ 
	private static final double fwhmIncrease = 0.014; 
	
	/** FWHM of beam velocity angular spread (at any point). 
	 * FIXME: made up for AUG to match observations, copied here for W7X */ 
	private double velocityAngularVariance = 1.5*1.5 *Math.PI*Math.PI /180/180;

	/** Spread of each component as fraction of calculated energy
			(i.e the 1/2 component is expected to have 1/2 the deviation of the full) 
	 	* FIXME: made up for AUG to match observations, copied here for W7X */
	public double energyFractionalStdDeviation = 0.015;  
	
	// Fuel and acceleration info is stored for D and H, and also for each of [BEAMBOX_*][BEAMFUEL_*]
	// See http://www.ipp.mpg.de/technologie/nbi/AUG/NIS.html
	
	/** Atomic No and mass per fuel type [BEAMFUEL_*] */
	private static final int atomicNo[] = { 1, 1 };
	private static final int atomicMass[] = { 1, 2 };
	
	/** Beam accelerating voltage in V [BEAMBOX_*][BEAMFUEL_*]. Currently from AUG*/
	private static final double accelVoltage[][] = new double[][]{ { 55e3, 60e3 }, { 55e3, 60e3 }, };

	/** Perveance at optimum in AV^(-3/2) [BEAMBOX_*][BEAMFUEL_*]   Currently from AUG */
	private static final double optimumPerveance[][] = new double[][]{ { 7.0e-6, 5.3e-6 }, { 7.0e-6, 5.3e-6 }, {}, {} };
	
	/** Transmission (injected / extracted)  current/power [BEAMBOX_*][BEAMFUEL_*]  Currently from AUG */
	private double transmission[][] = new double[][]{ { 0.363, 0.525 }, { 0.363, 0.525 },  {}, {} };

	/** Fraction of power in each beam component [BEAMBOX_*][BEAMFUEL_*][component]  
	 * Taken from that used in the physics scenarios page.
	 * [ https://validatorwww.ipp-hgw.mpg.de/Wiki/index.php/Physics_Scenarios ] 
	 * This is the same as AUG SE beams */ 
	private double componentPowerFractions[][][] = new double[][][]{
			{ { 0.50, 0.30, 0.20 },  { 0.65, 0.25, 0.10 } },  // SE: H, D from AUG
			{ { 0.50, 0.30, 0.20 },  { 0.65, 0.25, 0.10 } },  // SE: H, D from AUG
		};
	
	/** Molecular mass of components when accelerated (i.e. before disassociation) */
	private int componentMolecularMass[] = new int[]{ 1, 2, 3 };
	
	
	/** Beam setup for today, [beamIdx][type/geomIdx/fuel atomic no.] */	
	private int selectedBeams[][] = new int[][]{ 
			{ BEAMBOX_K20, 0, BEAMFUEL_DEUTERIUM }, 
			{ BEAMBOX_K20, 1, BEAMFUEL_DEUTERIUM }, 
			{ BEAMBOX_K20, 2, BEAMFUEL_DEUTERIUM }, 
			{ BEAMBOX_K20, 3, BEAMFUEL_DEUTERIUM }, 
			
			{ BEAMBOX_K21, 0, BEAMFUEL_DEUTERIUM }, 
			{ BEAMBOX_K21, 1, BEAMFUEL_DEUTERIUM }, 
			{ BEAMBOX_K21, 2, BEAMFUEL_DEUTERIUM }, 
			{ BEAMBOX_K21, 3, BEAMFUEL_DEUTERIUM }, 
	};
	
	
	private GaussianBeamDefinition beamDefs[];
	
	public W7XHeatingBeams() {
		setChanged(); //we have no parents, and need an init
	}
	
	public void calcBeams() {
		int nBeams = selectedBeams.length;
		
		beamDefs = new GaussianBeamDefinition[nBeams];
		
		for(int i=0; i < nBeams; i++){
			beamDefs[i] = new GaussianBeamDefinition();
			
			int type = selectedBeams[i][0];
			int geom = selectedBeams[i][1];
			int fuel = selectedBeams[i][2];
			
			//first sort out the geometry
			double A[] = new double[3];
			double Px[] = new double[3];
			double u[] = new double[3];
			
			A[0] = R_P0[type][geom] * Math.cos(theta[type][geom]);
			A[1] = R_P0[type][geom] * Math.sin(theta[type][geom]);
			A[2] = Z0[type][geom] + Dx[type][geom] * Math.tan(beta[type][geom]);
			
			Px[0] = A[0] + Dx[type][geom] * Math.cos(theta[type][geom] - psi[type][geom] - alpha[type][geom]);
			Px[1] = A[1] + Dx[type][geom] * Math.sin(theta[type][geom] - psi[type][geom] - alpha[type][geom]);
			Px[2] = Z0[type][geom] + 0;
			
			double l = Dx[type][geom] / Math.cos(beta[type][geom]); //length between A and Px, by def.
			
			u[0] = (A[0] - Px[0]) / l;
			u[1] = (A[1] - Px[1]) / l;
			u[2] = (A[2] - Px[2]) / l;
			
			beamDefs[i].uVec = u;
			
			//FIXME: Where actually is the focus? Here we assume it conincides with the beam crossing focus
			beamDefs[i].focus = Px; 
			
			beamDefs[i].fwhmAtFocus = fwhmAtFocus;
			beamDefs[i].fwhmIncrease = fwhmIncrease;
			beamDefs[i].velocityAngularVariance = velocityAngularVariance;
			
			
			//and now the fuel/accel info
			
			beamDefs[i].atomicNo = atomicNo[fuel];
			beamDefs[i].atomicMass = atomicMass[fuel];

			double extractedCurrent = optimumPerveance[type][fuel] * Math.pow(accelVoltage[type][fuel], 3.0/2.0);

			//molecules per second leaving the extraction grid
			double moleculeExtractionRate = extractedCurrent / (atomicNo[fuel] * electronCharge); //number of full energy particles (molecules or atoms) per beam

			int nComponents = componentMolecularMass.length;
			
			beamDefs[i].cmpPowerFrac = componentPowerFractions[type][fuel];
			beamDefs[i].cmpNumberFrac = new double[nComponents];
			beamDefs[i].cmpEnergyMean = new double[nComponents];
			beamDefs[i].cmpEnergyVar = new double[nComponents];
			beamDefs[i].cmpInjectionRate = new double[nComponents];
			
			// Calc the fraction of the injected atoms in each component
			double sumNumFrac = 0;			
			for(int iC=0; iC < nComponents; iC++) {
				double energyFrac = 1.0 / componentMolecularMass[iC];
				beamDefs[i].cmpNumberFrac[iC] = beamDefs[i].cmpPowerFrac[iC] / energyFrac;
				sumNumFrac += beamDefs[i].cmpNumberFrac[iC];
			}
				
			for(int iC=0; iC < nComponents; iC++) {
				beamDefs[i].cmpNumberFrac[iC] /= sumNumFrac;
				
				double energyFrac = 1.0 / componentMolecularMass[iC];
				
				beamDefs[i].cmpEnergyMean[iC] = accelVoltage[type][fuel] * energyFrac;
				beamDefs[i].cmpEnergyVar[iC] = (energyFractionalStdDeviation * beamDefs[i].cmpEnergyMean[iC]) 
												* (energyFractionalStdDeviation * beamDefs[i].cmpEnergyMean[iC]);
				
				// number of individual atoms injected
				double atomExtractionRate = moleculeExtractionRate * beamDefs[i].cmpPowerFrac[iC] * componentMolecularMass[iC]; 

				beamDefs[i].cmpInjectionRate[iC] = atomExtractionRate * transmission[type][fuel];

				/*System.out.println("beam=" + i + ", cmp=" + iC + 
						", E=" + beamDefs[i].cmpEnergyMean[iC] + 
						", particles/s=" + beamDefs[i].cmpInjectionRate[iC] +
						", power = " + (beamDefs[i].cmpEnergyMean[iC] * beamDefs[i].cmpInjectionRate[iC] * electronCharge));
						 //Power per source should be about 2.5MW
				*/ 
			}
		}
		
	}
	
	public GaussianBeamDefinition[] getBeamDefinitions(){ update(); return beamDefs; }
	
	public void setBeamSelection(int[][] selectedBeams) { this.selectedBeams = selectedBeams; setChanged("selectedBeams"); }

	public int[][] getBeamSelection() { return this.selectedBeams; }

	public void setComponentPowerFractions(double powerFractions[][][]){ this.componentPowerFractions = powerFractions; setChanged("powerFractions"); }
	
	public void setComponentPowerFractions(int beamBox, int fuelType, double powerFractions[]){ this.componentPowerFractions[beamBox][fuelType] = powerFractions; setChanged("powerFractions"); }
	
	public double[] getComponentPowerFractions(int beamBox, int fuelType) { return this.componentPowerFractions[beamBox][fuelType]; }
	
	public void setBeamVelocityAngularSigmaAll(double velocityAngularSigma) {
		this.velocityAngularVariance = velocityAngularSigma * velocityAngularSigma;
		setChanged("velocityAngularVariance"); 
	}
	
	public void setEnergyFractionalStdDeviation(double energyFractionalStdDeviation){
		this.energyFractionalStdDeviation = energyFractionalStdDeviation;
		setChanged("energyFractionalStdDeviation");
	}

	@Override
	public void updateState() {
		calcBeams();
	}
	
	
	@Override
	public double[][] generatePoints(double l0, double l1, double rMax, int nReq) { //Regular
		update();
		
		int j = 0; //beamIdx, err...
		
		double right[] = Algorithms.reNorm(Algorithms.cross(beamDefs[j].uVec, new double[]{0,0,1}));
		double up[] = Algorithms.cross(right, beamDefs[j].uVec);
		
		//calc number of point along axis
		int nL = (int)Math.pow(nReq * (l1-l0)*(l1-l0) / (Math.PI * rMax*rMax), 1.0/3.0);
		int nxy = (int)Math.sqrt(4 * nReq / nL / Math.PI);
		double dl = (l1-l0) / nL;
		double dxy = 2*rMax / nxy;
		double d = (dl + dxy) / 2;
		nL = (int)((l1-l0) / d);
		nxy = (int)(2.0 * rMax / d);
		
		//count actual number in CSA
		int nCS = 0;
		for(int iX=0; iX < nxy; iX++){
			double x = -rMax + d * iX;
			for(int iY=0; iY < nxy; iY++){
				double y = -rMax + d * iY;
				if( x*x + y*y <= rMax*rMax )
					nCS++;
			}
		}
		
		int n = nL * nCS;
		double ret[][] = new double[3][n]; 
		
		int i=0;
		for(int iL=0; iL < nL; iL++) {
			double l = l0 + iL * d;
			for(int iX=0; iX < nxy; iX++){
				double x = -rMax + d * iX;
				for(int iY=0; iY < nxy; iY++){
					double y = -rMax + d * iY;
					if( x*x + y*y <= rMax*rMax ){
						ret[0][i] = beamDefs[j].focus[0] + l * beamDefs[j].uVec[0] + x * right[0] + y * up[0];
						ret[1][i] = beamDefs[j].focus[1] + l * beamDefs[j].uVec[1] + x * right[1] + y * up[1];
						ret[2][i] = beamDefs[j].focus[2] + l * beamDefs[j].uVec[2] + x * right[2] + y * up[2];

						i++;
					}
				}
			}
		}
		
		return ret;
	}
	
	
	public double[][] generatePointsRandom(double l0, double l1, double rMax, int n) {
		update();
		
		double ret[][] = new double[3][n]; 
		for(int i=0; i < n; i++){
			int j = (int)(RandomManager.instance().nextUniform(0, 1) * beamDefs.length);
			double l = l0 + RandomManager.instance().nextUniform(0, 1) * (l1 - l0);
			ret[0][i] = beamDefs[j].focus[0] + l * beamDefs[j].uVec[0];
			ret[1][i] = beamDefs[j].focus[1] + l * beamDefs[j].uVec[1];
			ret[2][i] = beamDefs[j].focus[2] + l * beamDefs[j].uVec[2];
			
			double r = RandomManager.instance().nextUniform(0, 1);
			r = FastMath.sqrt(r);
			double phi = RandomManager.instance().nextUniform(0, 1) * 2 * Math.PI;
			
			double right[] = Algorithms.reNorm(Algorithms.cross(beamDefs[j].uVec, new double[]{0,0,1}));
			double up[] = Algorithms.cross(right, beamDefs[j].uVec);

			ret[0][i] += r * Math.cos(phi) * up[0] + r * Math.sin(phi) * right[0];
			ret[1][i] += r * Math.cos(phi) * up[1] + r * Math.sin(phi) * right[1];
			ret[2][i] += r * Math.cos(phi) * up[2] + r * Math.sin(phi) * right[2];
			
		}
		return ret;
	}
	
	/** Output the beam focal point and unit vectors, for copying into the optics package */
	public static void main(String[] args) {

		GraphicalModel g = new GraphicalModel("scrap");
		W7XHeatingBeams beamDefs = new W7XHeatingBeams();
		g.add(beamDefs);
		GaussianBeamDefinition beamDef[] = beamDefs.getBeamDefinitions();
	
		DecimalFormat fmt = new DecimalFormat("0.0000000");
		
		System.out.println("public static final double nbiStartAll[][] = { ");
		for(int i=0; i < beamDef.length; i++){
			System.out.println("\t{ " + fmt.format(beamDef[i].focus[0]) + 
								", " + fmt.format(beamDef[i].focus[1]) + 
								", " + fmt.format(beamDef[i].focus[2])
								+ " },");
			
		}
		
		System.out.println("};\n\npublic static final double nbiUnitAll[][] = {");
		for(int i=0; i < beamDef.length; i++){
			System.out.println("\t{ " + fmt.format(beamDef[i].uVec[0]) + 
								", " + fmt.format(beamDef[i].uVec[1]) + 
								", " + fmt.format(beamDef[i].uVec[2])
								+ " },");
		}
		System.out.println("};\n");
	}

}