package shared;

import jdk.incubator.vector.FloatVector;
import jdk.incubator.vector.VectorMask;
import jdk.incubator.vector.ByteVector;
import jdk.incubator.vector.IntVector;
import jdk.incubator.vector.VectorOperators;
import jdk.incubator.vector.VectorSpecies;

final class SIMD {
	
	//Example from https://medium.com/@Styp/java-18-vector-api-do-we-get-free-speed-up-c4510eda50d2
	@SuppressWarnings("restriction")
	private static final VectorSpecies<Float> FSPECIES=FloatVector.SPECIES_256;//FloatVector.SPECIES_PREFERRED; //This needs to be final or performance drops.
	private static final int FWIDTH=FSPECIES.length();
//	private static final int boundMask=~(FWIDTH-1);
	
	@SuppressWarnings("restriction")
	private static final VectorSpecies<Byte> BSPECIES=ByteVector.SPECIES_256;
	private static final int BWIDTH=BSPECIES.length();
	
	@SuppressWarnings("restriction")
	private static final VectorSpecies<Integer> ISPECIES=IntVector.SPECIES_256;
	private static final int IWIDTH=ISPECIES.length();
	
	@SuppressWarnings("restriction")
	/** 
	 * Vectorized version of "c+=a[i]*b[i]" where a and b are equal-length arrays.
	 * @param a A vector to multiply.
	 * @param b A vector to multiply.
	 * @return Sum of products of vector elements.
	 */
	public static final float fma(final float[] a, final float[] b){
		//final int width=SPECIES.length();
		final int limit=FSPECIES.loopBound(a.length);
//		final int limit=a.length&boundMask;

		FloatVector sum=FloatVector.zero(FSPECIES);
		int i=0;
		for(; i<limit; i+=FWIDTH) {//SIMD loop
			FloatVector va=FloatVector.fromArray(FSPECIES, a, i);
			FloatVector vb=FloatVector.fromArray(FSPECIES, b, i);
			sum=va.fma(vb, sum);
		}
		float c=sum.reduceLanes(VectorOperators.ADD);
		for (; i<a.length; i++) {//Residual scalar loop
			c+=a[i]*b[i];
		}
		return c;
	}
	
	/** 
	 * Performs "a+=incr" where a and incr are equal-length arrays.
	 * @param a A vector to increment.
	 * @param b Increment amount.
	 */
	@SuppressWarnings("restriction")
	public static final void add(final float[] a, final float[] b){
		//final int width=SPECIES.length();
		final int limit=FSPECIES.loopBound(a.length);
//		final int limit=a.length&boundMask;
		
		int i=0;
		for(; i<limit; i+=FWIDTH) {//SIMD loop
			FloatVector va=FloatVector.fromArray(FSPECIES, a, i);
			FloatVector vb=FloatVector.fromArray(FSPECIES, b, i);
			FloatVector sum=va.add(vb);
			sum.intoArray(a, i);
		}
		for (; i<a.length; i++) {//Residual scalar loop; TODO: replace with vector mask
			a[i]+=b[i];
		}
	}

	/** 
	 * Performs "a+=b*mult" where a and b are equal-length arrays.
	 * @param a A vector to increment.
	 * @param b Increment amount.
	 * @param mult Increment multiplier.
	 */
	@SuppressWarnings("restriction")
	public static final void addProduct(final float[] a, final float[] b, final float mult){
		//final int width=SPECIES.length();
		final int limit=FSPECIES.loopBound(a.length);
//		final int limit=a.length&boundMask;
		
		int i=0;
		for(; i<limit; i+=FWIDTH) {//SIMD loop
			FloatVector va=FloatVector.fromArray(FSPECIES, a, i);
			FloatVector vb=FloatVector.fromArray(FSPECIES, b, i);
			FloatVector sum=va.add(vb.mul(mult));
			sum.intoArray(a, i);
		}
		for (; i<a.length; i++) {//Residual scalar loop; TODO: replace with vector mask
			a[i]+=b[i]*mult;
		}
	}
	
	//a is dest
	@SuppressWarnings("restriction")
	public static final void copy(final float[] a, final float[] b){
		final int length=Tools.min(a.length, b.length);
		//final int width=SPECIES.length();
		final int limit=FSPECIES.loopBound(length);
//		final int limit=a.length&boundMask;
		
		int i=0;
		for(; i<limit; i+=FWIDTH) {//SIMD loop
			FloatVector vb=FloatVector.fromArray(FSPECIES, b, i);
			vb.intoArray(a, i);
		}
		for (; i<length; i++) {//Residual scalar loop; TODO: replace with vector mask
			a[i]=b[i];
		}
	}
	
	/** Returns number of matches */
	@SuppressWarnings("restriction")
	public static final int countMatches(final byte[] s1, final byte[] s2, int a1, int b1, int a2, int b2){
		final int length=b2-a2+1;
		final int limit0=BSPECIES.loopBound(length);
		final int limit=a2+limit0;
		
		int i=a1, j=a2;
		int matches=0;
		for(; j<limit; i+=BWIDTH, j+=BWIDTH) {//SIMD loop
			ByteVector v1=ByteVector.fromArray(BSPECIES, s1, i);
			ByteVector v2=ByteVector.fromArray(BSPECIES, s2, j);
			VectorMask<Byte> x=v1.eq(v2);
			matches+=x.trueCount();//This might be slow, or might not
		}
		for(; j<=b2; i++, j++) {
			final byte x=s1[i], y=s2[j];
			final int m=((x==y) ? 1 : 0);
			matches+=m;
		}
		return matches;
	}
	
	/** Returns index of symbol */
	@SuppressWarnings("restriction")
	public static final int find(final byte[] a, final byte symbol, final int from, final int to){	
		final int length=to-from;
		final int limit0=BSPECIES.loopBound(length);
		final int limit=from+limit0;
		
		int pos=from;
		for(; pos<limit; pos+=BWIDTH) {//SIMD loop
			ByteVector v=ByteVector.fromArray(BSPECIES, a, pos);
			VectorMask<Byte> x=v.eq(symbol);
			if(x.anyTrue()) {break;}
		}
		while(pos<to && a[pos]!=symbol){pos++;}
		return pos;
	}
	
}
