package org.kne.codec.kif;

import java.awt.geom.Area;
import java.io.IOException;
import java.lang.foreign.Arena;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.ValueLayout;

import org.kne.debug.TimeDebugger;
import org.kne.membandboost.MembandBoost;
import org.kne.opencl64.OpenCLKernel;
import org.kne.opencl64.OpenCLMemory;
import org.kne.opencl64.OpenCLMemoryCopyType;
import org.kne.opencl64.OpenCLProgram;
import org.kne.opencl64.concurrent.OpenCLContextThread;
import org.kne.opencl64.concurrent.ThreadLocalOpenCLMemory;

import jdk.internal.foreign.ArenaImpl;

public class KIFDecodePipeline {
	private KIFCompressedImageBlock input;
	
	
	
	 public KIFDecodePipeline(KIFCompressedImageBlock input) {
		super();
		this.input = input;
		data = input.getData();
		width = input.getWidth();
		height = input.getHeight();
        mode = input.getPredictMode();
	}
	private TimeDebugger tdb=new TimeDebugger("解码时间",false);
	private volatile byte[] data;
	private volatile int width;
	private volatile int height;
	private volatile KIFCompressMode mode;
	private volatile byte[] dataserp;
	
	
	public void decodeStage1() throws IOException {
		  dataserp = Compressors.decompressZstd(data,input.getSizeBeforeCompress());
	        tdb.mark("熵解码");
	}
	
	public IntImageBlock decodeStage2() throws IOException {

		Thread t =Thread. currentThread();
		if (t instanceof OpenCLContextThread) {
			OpenCLContextThread oct = (OpenCLContextThread) t;
			return decodeStage2GPU(oct);
		}else {
		    return decodeStage2CPU();
		}
	}
	private static ThreadLocal<MemorySegment>m1=ThreadLocal.withInitial(()->{
		return ((ArenaImpl)Arena.ofAuto()).allocateNoInit(1024*1024*4,4);
	});

	private static ThreadLocalOpenCLMemory tlm0=new ThreadLocalOpenCLMemory(1024*1024*4);
	private static ThreadLocalOpenCLMemory tlm1=new ThreadLocalOpenCLMemory(1024*1024*4);
	private static ThreadLocalOpenCLMemory tlm2=new ThreadLocalOpenCLMemory(1024*1024*4);
	private IntImageBlock decodeStage2GPU(OpenCLContextThread oct) throws IOException {
	      tdb.mark("等待调度");
		 OpenCLProgram prog = oct.getProgram("/org/kne/codec/kif/KIFCodec.cl");
	      tdb.mark("加载程序");
			//OpenCLKernel krl4 = prog.createKernel("colorInverseRGBAInPlace");
			OpenCLMemory ocm1 = null; 
			OpenCLMemory ocm2 = null; 
			OpenCLMemory ocm3 = null; 
		       int[]des;
			MemorySegment dataserpseg= m1.get();
		      tdb.mark("堆外内存分配");
			dataserpseg.asByteBuffer().put(dataserp).flip();
		      tdb.mark("堆内存->堆外内存");
			ocm1=tlm0.get(oct.getOpenCLContext());
		      tdb.mark("分配显存");
			oct.getOpenCLCommandQueue().executeWriteBuffer(ocm1, dataserpseg, false,0,0,dataserp.length);
		      tdb.mark("内存->显存");
			
			OpenCLMemory dataser;
	        if(mode.isBP()) {
				ocm2=tlm1.get(oct.getOpenCLContext());
				OpenCLKernel krl0 = prog.createKernel("decodePacked");
				krl0.putArg(ocm1);
				krl0.putArg(ocm2);
				krl0.putArg(dataserp.length);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl0, 0, dataserp.length);
	        	dataser=ocm2;//
	        }else {
	        	dataser=ocm1;
	        }
	      //  oct.getOpenCLCommandQueue().finish();
	       // tdb.mark("位平面解码");
	        
	        ocm3=tlm2.get(oct.getOpenCLContext());
	        if(mode.isPL()) {
				OpenCLKernel krl1 = prog.createKernel("deserializePlannarRGBAWithZigzag");
				krl1.putArg(dataser);
				krl1.putArg(ocm3);
				krl1.putArg(width*height);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl1, 0,width*height);
	        }else {
				OpenCLKernel krl2 = prog.createKernel("deserializeRGBAWithZigzag");
				krl2.putArg(dataser);
				krl2.putArg(ocm3);
				krl2.putArg(width*height);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl2, 0,width*height);
	        }
	      //  oct.getOpenCLCommandQueue().finish();
	     //  tdb.mark("反序列化");
	       
	        
	        if(mode.isCD()) {
				OpenCLKernel krl4 = prog.createKernel("colorInverseRGBAInPlace");
				krl4.putArg(ocm3);
				krl4.putArg(width*height);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl4, 0, width*height);
	        }
	        if(mode.isV()) {
				OpenCLKernel krl3 = prog.createKernel("verticalInverseInPlace");
				krl3.putArg(ocm3);
				krl3.putArg(width);
				krl3.putArg(height);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl3, 0, width);
	        }
	        if(mode.isH()) {
				OpenCLKernel krl3 = prog.createKernel("horizontalInverseInPlace");
				krl3.putArg(ocm3);
				krl3.putArg(width);
				krl3.putArg(height);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl3, 0, height);
	        }
	      //  oct.getOpenCLCommandQueue().finish();
		    //  tdb.mark("预测解码");
		      tdb.mark("解码");
	        IntImageBlock result=new IntImageBlock( width, height, input.getPosx(), input.getPosy());
		      tdb.mark("分配堆外内存");

	        oct.getOpenCLCommandQueue().executeReadBuffer(ocm3, result.getImage(), true,0,0,width*height*4);
		      tdb.mark("显存->内存");
	     //   MemorySegment.copy(dataserpseg, ValueLayout.JAVA_INT, 0, des, 0, des.length);
		   //   tdb.mark("堆外内存->堆内存");
	      tdb.print();
			
		return result;
	}

	private IntImageBlock decodeStage2CPU() throws IOException {
	    byte[]dataser;
        if(mode.isBP()) {
        	dataser=BitPlane.decodePacked(dataserp);
        }else {
        	dataser=dataserp;
        }
        tdb.mark("位平面解码");
       int[]des;
        if(mode.isPL()) {
       des=PixelSerializer.deserializePlannarRGBAWithZigzag(dataser, width, height);
        }else {
       des=PixelSerializer.deserializeRGBAWithZigzag(dataser, width, height);
        }
       tdb.mark("反序列化");
        if(mode.isTR())
       Predictor.inverseInPlace(des, width, height,1,-1);
        if(mode.isTL())
       Predictor.inverseInPlace(des, width, height,-1,-1);
        if(mode.isV())
       Predictor.inverseInPlace(des, width, height,0,-1);
        if(mode.isH())
       Predictor.inverseInPlace(des, width, height,-1,0);
        if(mode.isCD())
      ColorTransform.colorInverseRGBAInPlace(des, width, height);
      tdb.mark("预测解码");
      tdb.print();
	return new IntImageBlock(MemorySegment.ofArray(des), width, height, input.getPosx(), input.getPosy());
	}
}
