package org.kne.opencl64;

import java.lang.foreign.Arena;
import java.lang.foreign.MemorySegment;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicInteger;

public class HighBandwidthCopy {
	private static int defaultCores=4;
public static int getDefaultCores() {
		return defaultCores;
	}
	public static void setDefaultCores(int defaultCores) {
		HighBandwidthCopy.defaultCores = defaultCores;
	}
private static ExecutorService exc=Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
public static void multicorecopy(MemorySegment src,long srcPos,MemorySegment dst,long dstPos,long bytes,int cores) {
	long count=bytes/cores;
	long mod=bytes%cores;
	AtomicInteger aig=new AtomicInteger(0);
	for (int i = 0; i < cores; i++) {
		int j=i;
		exc.execute(()->{
			defaultcopy(src, srcPos+count*j, dst, dstPos+count*j, count);
			aig.incrementAndGet();
		});
	}
	defaultcopy(src, srcPos+count*cores, dst, dstPos+count*cores, mod);
	while(aig.get()<cores) {}
}
public static void multicorecopy(MemorySegment src,long srcpos,MemorySegment dst,long dstpos,long bytes) {
	multicorecopy(src,srcpos,dst,dstpos,bytes,defaultCores);
}
	public static void defaultcopy(MemorySegment src,long srcpos,MemorySegment dst,long dstpos,long bytes) {
		MemorySegment.copy(src, srcpos, dst, dstpos, bytes);
	}
	public static void main(String[] args) {
		for(int i=0;i<34;i++) {
			long bytes=1L<<i;
			System.out.print("Size: "+bytesUnit(bytes)+"\t");
			testMemoryBandwidth(bytes);
			System.out.println();
		}
	}
	private static void testMemoryBandwidth(long bytes) {
		MemorySegment src=Arena.ofAuto().allocate(bytes);
		MemorySegment dst=Arena.ofAuto().allocate(bytes);
		//int maxcores=Runtime.getRuntime().availableProcessors();
		int maxcores=10;
		for(int i=1;i<maxcores;i++) {
			testMemoryBandwidth(bytes,i,src,dst);
		}
	}
	private static void testMemoryBandwidth(long bytes, int cores, MemorySegment src, MemorySegment dst) {
		long speedsum=0;
		int tests=5;
		for(int i=0;i<tests;i++) {
		long start=System.nanoTime();
		multicorecopy(src, 0, dst, 0, bytes,cores);
		//MemorySegment.copy(src, 0, dst, 0, bytes);
		//defaultcopy(src, 0, dst, 0, bytes);
		long end=System.nanoTime()-start;
		long speed=bytes*1000000000L/end;
		speedsum+=speed;
		}
		System.out.print(bytesUnit(speedsum/tests)+"/s\t");
	}
	public static String bytesUnit(long v) {
		if (v >= 1024L * 1024 * 1024 * 1024 * 1024) {
			return String.format("%.2f", v / (1024.0 * 1024.0 * 1024.0 * 1024.0 * 1024.0)) + "PB";
		} else if (v >= 1024L * 1024 * 1024 * 1024) {
			return String.format("%.2f", v / (1024.0 * 1024.0 * 1024.0 * 1024.0)) + "TB";
		} else if (v >= 1024L * 1024 * 1024) {
			return String.format("%.2f", v / (1024.0 * 1024.0 * 1024.0)) + "GB";
		} else if (v >= 1024L * 1024) {
			return String.format("%.2f", v / (1024.0 * 1024.0)) + "MB";
		} else if (v >= 1024L) {
			return String.format("%.2f", v / (1024.0)) + "KB";
		} else {
			return v + "B";
		}
	}
}
