package mior.model;

import java.util.Arrays;
import java.util.BitSet;

import mcmas.core.MCMCommandQueue;
import mcmas.core.MCMCommandQueueProperty;
import mcmas.core.MCMContext;
import mcmas.core.MCMEvent;
import mcmas.core.MCMKernel;
import mcmas.core.MCMMem;
import mcmas.core.MCMProgram;
import mcmas.core.MCMUtils;
import mior.model.dist.IMiorDistribution;

import org.jocl.Pointer;
import org.perf4j.slf4j.Slf4JStopWatch;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


public class OCLParaCSRModel4 extends AbstractMiorModel {
	
	/**
	 * Java view
	 */
	
	private final int nbMM;
	private final int nbOM;
	
	private final MiorMM [] mmList;
	private final MiorOM [] omList;
	private final MiorWorld [] worlds;
	
	/**
	 * OpenCL implementation
	 */
	final MCMContext context;
	final MCMCommandQueue queue;
	final MCMProgram program;
	
	final MCMKernel topoKernel;
	final MCMKernel simulateKernel;
	final MCMKernel autoliveKernel;
	
	final MCMMem mmMem;
	final MCMMem omMem;
	final MCMMem worldsMem;
	
	final MCMMem mmOffsetsMem;
	final MCMMem omOffsetsMem;
	
	MCMMem mmCSRMem;
	MCMMem omCSRMem;
	
	final MCMMem partsMem;
	
	private final int [] mmOffsets;
	private final int [] omOffsets;
	
	private int [] mmCSR;
	private int [] omCSR;
	private final int [] parts;
	
	private final int nbSim;
	private int blockSize;
	
	private final static Logger logger = LoggerFactory.getLogger(OCLParaCSRModel4.class);
	private final Slf4JStopWatch watch = new Slf4JStopWatch(logger);
	
	private final static String MODEL_SOURCE = "kernels/mior_model_multisim4.cl";
	private final IMiorDistribution dist;
	
	public OCLParaCSRModel4(int onbMM, int onbOM, int nbSim, IMiorDistribution dist) {
		this.dist = dist;
		System.out.println(dist);
		
		//final double factor = dist.getMax() / ((double) onbOM);
		this.nbOM = (int) (dist.getMaxFactor() * dist.getMeanOM());
		this.nbMM = (int) (dist.getMaxFactor() * dist.getMeanMM());
		
		this.nbSim = nbSim;
		this.blockSize = Math.max(nbOM, nbMM);
		
		// Allocate Java MIOR structures
		this.mmList       = new MiorMM[nbSim * nbMM];
		this.omList       = new MiorOM[nbSim * nbOM];
		this.worlds       = new MiorWorld[nbSim];
		
		this.mmOffsets    = new int[nbSim * nbMM];
    this.omOffsets    = new int[nbSim * nbOM];
		//this.mmCSR        = allocateStorage(nbSim * nbOM, nbMM);
		//this.omCSR        = allocateStorage(nbSim * nbMM, nbOM);
		this.parts        = new int[nbSim * nbMM * nbOM];
		
		MiorUtils.initWorldArray(worlds, dist);
		MiorUtils.initRandomMMArray(mmList, worlds[0].width);
		MiorUtils.initRandomOMArray(omList, worlds[0].width);
		
		// OpenCL allocations
		watch.start("OpenCL setup");
		
		this.context = new MCMContext();
		this.queue = context.createCommandQueue(MCMCommandQueueProperty.ENABLE_PROFILING);
		
		watch.stop();
		
		System.out.println("NB_MM: " + nbMM + ", NB_OM: " + nbOM);
		watch.start("Program compilation");
		this.program = MCMUtils.compileFile(context, MODEL_SOURCE, " -DNB_MM=" + nbMM + " -DNB_OM=" + nbOM);// + " -g -O0");
		watch.stop();
		
		watch.start("Kernel allocation");
		
		this.topoKernel = program.createKernel("topology");
		this.simulateKernel = program.createKernel("simulate");
		this.autoliveKernel = program.createKernel("autolive");
		
		watch.stop();
		watch.start("Memory allocation");
		
		this.mmMem = context.newBuffer().Using(mmList).b();
		this.omMem = context.newBuffer().Using(omList).b();
		this.worldsMem = context.newBuffer().Using(worlds).b();
		this.mmOffsetsMem = context.newBuffer().Size(nbSim * nbMM).b();
		this.omOffsetsMem = context.newBuffer().Size(nbSim * nbOM).b();
		
		//this.mmCSRMem = context.newBuffer().Using(mmCSR).b();
		//this.omCSRMem = context.newBuffer().Using(omCSR).b();
		this.partsMem = context.newBuffer().Using(new int[parts.length]).b();
		
		watch.stop();
	}
	
	/*private int[] allocateStorage(int nbList, int listSize) {
		return new int[nbList * (listSize + 1)];
	}*/

	@Override
	public void setBlockSize(int blockSize) {
		this.blockSize = blockSize;
	}
	
	@Override
	public int getBlockSize() {
		return blockSize;
	}
	
	
	@Override
	public int getNbSimulations() {
		return nbSim;
	}
	
	@Override
	protected void resetImpl() {
		// Initialise models
		/*for (int i = 0; i < nbSim; i++) {
			this.worlds[i] = new MiorWorld(nbMM, nbOM);
		}*/
		
		MiorUtils.initWorldArray(worlds, dist);
		MiorUtils.initRandomMMArray(mmList, worlds[0].width);
		MiorUtils.initRandomOMArray(omList, worlds[0].width);
		
		if (mmCSR != null) Arrays.fill(mmCSR, 0);
		if (omCSR != null) Arrays.fill(omCSR, 0);
		Arrays.fill(parts, 0);
		
		queue.enqueueWriteBuffer(mmMem, mmList, 0, mmMem.getSize());
		queue.enqueueWriteBuffer(omMem, omList, 0, omMem.getSize());
		queue.enqueueWriteBuffer(worldsMem, worlds, 0, worldsMem.getSize());
		
		if (mmCSR != null) { 
			queue.enqueueWriteBuffer(mmCSRMem, Pointer.to(mmCSR), 0, mmCSRMem.getSize());
		}
		if (omCSR != null) {
			queue.enqueueWriteBuffer(omCSRMem, Pointer.to(omCSR), 0, omCSRMem.getSize());
		}
		queue.enqueueWriteBuffer(partsMem, Pointer.to(parts), 0, partsMem.getSize());
		queue.finish();
	}
	
	@Override
	protected void doTopologyImpl() {
		troveTopology();
		
		queue.enqueueWriteBuffer(mmOffsetsMem, Pointer.to(mmOffsets), 0, mmOffsetsMem.getSize());
		queue.enqueueWriteBuffer(mmOffsetsMem, Pointer.to(mmOffsets), 0, mmOffsetsMem.getSize());
		this.mmCSRMem = context.newBuffer().Using(mmCSR).b();
		this.omCSRMem = context.newBuffer().Using(omCSR).b();
		queue.finish();
	}
	
	protected void doTopologyImplOld() {
		// System.out.println(topoKernel.getArgumentsNumber());
		topoKernel.setArguments(mmMem, omMem, worldsMem, mmCSRMem, omCSRMem); // associationsMem
		
		MCMEvent event = queue.enqueueKernel(topoKernel, 2,
				new long[] { nbSim * nbMM , nbOM} // Global size
		);
		
		/*OCLEvent event = queue.enqueueKernel(topoKernel, 2,
				new long[] { nbMM , nbSim * nbOM}, // Global size
				new long[] { 1, nbOM }
		);*/
		
		MCMEvent.waitFor(event);
		MCMUtils.printEventStats("topology", event);
		
		if (! isBatchModeEnabled()) {
			queue.blockingReadBuffer(mmCSRMem, Pointer.to(mmCSR), 0, mmCSRMem.getSize());
			queue.blockingReadBuffer(omCSRMem, Pointer.to(omCSR), 0, omCSRMem.getSize());
			//System.out.println(Arrays.toString(mmCSR));
			//queue.blockingReadBuffer(associationsMem, Pointer.to(associations), 0, associationsMem.getSize());
		}
	}
	
	@Override
	protected void doLiveImpl() {
		//throw new UnsupportedOperationException();
		simulateKernel.setArguments(mmMem, omMem, worldsMem, mmOffsetsMem, omOffsetsMem,  mmCSRMem, omCSRMem, partsMem);
		
		if (blockSize < Math.max(nbOM, nbMM)) {
			throw new RuntimeException("blockSize (" + blockSize + ") is too small to execute the simulation");
		}
		
		MCMEvent event = queue.enqueue1DKernel(simulateKernel, nbSim * blockSize, blockSize);
		
		MCMEvent.waitFor(event);
		MCMUtils.printEventStats("simulate", event);
		
		if (! isBatchModeEnabled()) {
			queue.blockingReadBuffer(mmMem, mmList, 0, mmMem.getSize());
			queue.blockingReadBuffer(omMem, omList, 0, omMem.getSize());
			queue.blockingReadBuffer(worldsMem, worlds, 0, worldsMem.getSize());
			System.out.println("copy");
		}
	}
	
	@Override
	public void doAutoLive() {
		resetImpl();
		doTopologyImpl();
		
		if (blockSize < Math.max(nbOM, nbMM)) {
			throw new RuntimeException("blockSize (" + blockSize + ") is too small to execute the simulation");
		}
		
		autoliveKernel.setArguments(mmMem, omMem, worldsMem,  mmOffsetsMem, omOffsetsMem, mmCSRMem, omCSRMem, partsMem);
		
		if (! isBatchModeEnabled()) {
				int CO2 = -1;
				
				while (CO2 == -1 || CO2 != worlds[0].CO2) {
					CO2 = worlds[0].CO2;
					System.out.println("CO2 = " + CO2);
					doLive();
				}
				
		} else { 
			autoliveKernel.setArguments(mmMem, omMem, worldsMem, mmCSRMem, omCSRMem, partsMem);
			
			MCMEvent event = queue.enqueue1DKernel(autoliveKernel, nbSim * blockSize, blockSize);
			
			MCMEvent.waitFor(event);
			MCMUtils.printEventStats("autolive", event);
		}
		
		onSimulationFinished();
	}
	
	@Override
	protected void onSimulationFinished() {
		
		queue.blockingReadBuffer(mmMem, mmList, 0, mmMem.getSize());
		queue.blockingReadBuffer(omMem, omList, 0, omMem.getSize());
		queue.blockingReadBuffer(worldsMem, worlds, 0, worldsMem.getSize());
		//queue.blockingReadBuffer(associationsMem, Pointer.to(associations), 0, associationsMem.getSize());
		queue.blockingReadBuffer(partsMem, Pointer.to(parts), 0, partsMem.getSize());
		
		super.onSimulationFinished();
	}
	
	@Override
	protected void releaseImpl() {
		partsMem.release();
		
		omCSRMem.release();
		mmCSRMem.release();
		//associationsMem.release();
		
		worldsMem.release();
		omMem.release();
		mmMem.release();
		
		autoliveKernel.release();
		simulateKernel.release();
		topoKernel.release();
		
		program.release();
		queue.release();
		context.release();
	}

	@Override
	public MiorWorld getWorld() {
		return worlds[0];
	}

	@Override
	public MiorOM[] getOMList() {
		return Arrays.copyOfRange(omList, 0, nbOM);
	}

	@Override
	public MiorMM[] getMMList() {
		return Arrays.copyOfRange(mmList, 0, nbMM);
	}

	@Override
	public boolean isAccessible(int iMM, int iOM) {
		return false;
		//return associations[iMM * nbOM + iOM] != -1;
	}

	@Override
	public boolean isAccessible(int iMM, int iOM, int iSim) {
		return false;
		//return associations[iSim * nbMM * nbOM + iMM * nbOM + iOM] != -1;
	}
	
	private void troveTopology() {
		//int[] mmOffsets = new int[nbSim * nbMM];
		//int[] omOffsets = new int[nbSim * nbOM];
		int[] mmSizes = new int[nbSim * nbMM];
		int[] omSizes = new int[nbSim * nbOM];
		DynamicArray mmCSR = new DynamicArray();
		DynamicArray omCSR = new DynamicArray();
		BitSet bitmap = new BitSet();
		
		for (int iSim = 0; iSim < nbSim; iSim++) {
			final MiorWorld world = worlds[iSim];
			
			for (int iMM = 0; iMM < world.nbMM; iMM++) {			
				
				for (int iOM = 0; iOM < world.nbOM; iOM++) {
					final float dx = omList[iSim * nbOM + iOM].x - mmList[iSim * nbMM + iMM].x;
					final float dy = omList[iSim * nbOM + iOM].y - mmList[iSim * nbMM + iMM].y;
					final double distance = Math.sqrt(dx * dx + dy * dy);
					
					if (distance <= world.RA) {
						mmSizes[iSim * nbMM + iMM]++;
						omSizes[iSim * nbOM + iOM]++;
						bitmap.set(iSim * nbOM * nbMM + iMM * nbOM + iOM);
						//System.out.println("Store link " + iMM + " - " + iOM);
					}
				}
			}
		}
		
		for (int i = 1; i < nbSim * nbMM; i++) {
			mmOffsets[i] = mmOffsets[i - 1] + mmSizes[i - 1] + 1;
			mmCSR.set(mmOffsets[i], 0);
		}
		
		for (int i = 1; i < nbSim * nbOM; i++) {
			omOffsets[i] = omOffsets[i - 1] + omSizes[i - 1] + 1;
			omCSR.set(omOffsets[i], 0 /*omSizes[i]*/);
		}
		
		for (int iSim = 0; iSim < nbSim; iSim++) {
			final MiorWorld world = worlds[iSim];
			
			for (int iMM = 0; iMM < world.nbMM; iMM++) {
				final int mmIndex = iSim * nbMM + iMM;
				
				for (int iOM = 0; iOM < world.nbOM; iOM++) {
					final int omIndex = iSim * nbOM + iOM;
					
					if (bitmap.get(iSim * nbOM * nbMM + iMM * nbOM + iOM)) {
						list_add(omCSR, omOffsets[omIndex], mmIndex);
						list_add(mmCSR, mmOffsets[mmIndex], omIndex);
						//System.out.println("Add link " + iMM + " - " + iOM);
					}
				}
			}
		}
		
		this.mmCSR = mmCSR.toArray();
		this.omCSR = omCSR.toArray();
		
		/*
		System.out.println("nbSim: " + nbSim);
		System.out.println("mmSizes: " + Arrays.toString(mmSizes));
		System.out.println("mmOffsets: " + Arrays.toString(mmOffsets));
		System.out.println("mmCSR: " + Arrays.toString(this.mmCSR));
		System.out.println("omSizes: " + Arrays.toString(omSizes));
		System.out.println("omOffsets: " + Arrays.toString(omOffsets));
		System.out.println("omCSR: " + Arrays.toString(this.omCSR));*/
	}
	
	private void list_add(DynamicArray list, int index, int val) {
		final int offset = list.get(index);
		list.set(offset + 1 + index, val);
		list.set(index, offset + 1);
	}
	
}