package org.kne.codec.kif;

import jdk.internal.foreign.ArenaImpl;
import java.lang.foreign.Arena;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.ValueLayout;
import java.lang.foreign.ValueLayout.OfInt;

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

public class KIFEncodePipeline {
	
	private volatile IntImageBlock input;
	private volatile KIFCompressMode mode; 
	 private volatile TimeDebugger tdb=new TimeDebugger("编码时间",false);
	 
	 

	 public KIFEncodePipeline(IntImageBlock input, KIFCompressMode mode) {
		super();
		this.input = input;
		this.mode = mode;
		 pixels = input.getImage();
		 width = input.getWidth();
		 height = input.getHeight();
	}

	private byte[] serialData2;
	private MemorySegment pixels;
	private int width;
	private int height;
	
	public void encodeStage1() {
		Thread t =Thread. currentThread();
		if (t instanceof OpenCLContextThread) {
			OpenCLContextThread oct = (OpenCLContextThread) t;
			encodeStage1GPU(oct);
		}else {
		    encodeStage1CPU();
		}
	}

	private void encodeStage1CPU() {
		int[]pixels2=pixels.toArray(OfInt.JAVA_INT);
		if(mode.isCD()) {
 pixels2=ColorTransform.colorTransformRGBA(pixels2, width, height);
		}else {
			pixels2=pixels2.clone();
		}
      if(mode.isH())
      Predictor.predictInPlace(pixels2, width, height,-1,0);
      if(mode.isV())
       Predictor.predictInPlace(pixels2, width, height,0,-1);
      if(mode.isTL())
      Predictor.predictInPlace(pixels2, width, height,-1,-1);
      if(mode.isTR())
       Predictor.predictInPlace(pixels2, width, height,1,-1);
       tdb.mark("预测编码");
      
      byte[] serialData ;
      if(mode.isPL()) {
      serialData= PixelSerializer.serializePlannarRGBAWithZigzag(pixels2, width, height);
      }else {
      serialData= PixelSerializer.serializeRGBAWithZigzag(pixels2, width, height);
      }
      tdb.mark("序列化");
      
      if(mode.isBP()) {
		serialData2= BitPlane.encodePacked(serialData);
      }else {
		serialData2=serialData;
      }
      tdb.mark("位平面编码");
	}


	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 tlmpred1=new ThreadLocalOpenCLMemory(1024*1024*4);
	private static ThreadLocalOpenCLMemory tlmpred2=new ThreadLocalOpenCLMemory(1024*1024*4);
	private static ThreadLocalOpenCLMemory tlm1=new ThreadLocalOpenCLMemory(1024*1024*4);
	private static ThreadLocalOpenCLMemory tlm2=new ThreadLocalOpenCLMemory(1024*1024*4);
	private void encodeStage1GPU(OpenCLContextThread oct) {
		 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; 
			MemorySegment dataserpseg= m1.get();
		      tdb.mark("堆外内存分配");
	        //MemorySegment.copy(pixels,0,dataserpseg,0, pixels.byteSize());
		      tdb.mark("堆内存->堆外内存");
			ocm1=tlm0.get(oct.getOpenCLContext());
		      tdb.mark("分配显存");
			oct.getOpenCLCommandQueue().executeWriteBuffer(ocm1, pixels, false,0,0,pixels.byteSize());
		      tdb.mark("内存->显存");
			
	        if(mode.isH()) {
				OpenCLKernel krl3 = prog.createKernel("horizontalPredictParallel");
				OpenCLMemory opred1=tlmpred1.get(oct.getOpenCLContext());
				krl3.putArg(ocm1);
				krl3.putArg(opred1);
				krl3.putArg(width);
				krl3.putArg(height);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl3, 0, width*height);
				ocm1=opred1;
	        }
	        
	        if(mode.isV()) {
				OpenCLKernel krl3 = prog.createKernel("verticalPredictParallel");
				OpenCLMemory opred2=tlmpred2.get(oct.getOpenCLContext());
				krl3.putArg(ocm1);
				krl3.putArg(opred2);
				krl3.putArg(width);
				krl3.putArg(height);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl3, 0, width*height);
				ocm1=opred2;
	        }
	        if(mode.isCD()) {
				OpenCLKernel krl4 = prog.createKernel("colorTransformRGBAInPlace");
				krl4.putArg(ocm1);
				krl4.putArg(width*height);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl4, 0, width*height);
	        }
	      //  oct.getOpenCLCommandQueue().finish();
		    //  tdb.mark("预测解码");
			ocm2=tlm1.get(oct.getOpenCLContext());
	        if(mode.isPL()) {
				OpenCLKernel krl1 = prog.createKernel("serializePlannarRGBAWithZigzag");
				krl1.putArg(ocm1);
				krl1.putArg(ocm2);
				krl1.putArg(width*height);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl1, 0,width*height);
	        }else {
				OpenCLKernel krl2 = prog.createKernel("serializeRGBAWithZigzag");
				krl2.putArg(ocm1);
				krl2.putArg(ocm2);
				krl2.putArg(width*height);
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl2, 0,width*height);
	        }
	      //  oct.getOpenCLCommandQueue().finish();
	     //  tdb.mark("反序列化");
			OpenCLMemory dataser;
	        if(mode.isBP()) {
		        
		        ocm3=tlm2.get(oct.getOpenCLContext());
				OpenCLKernel krl0 = prog.createKernel("encodePacked");
				krl0.putArg(ocm2);
				krl0.putArg(ocm3);
				krl0.putArg((int)pixels.byteSize());
				oct.getOpenCLCommandQueue().execute1DRangeKernel(krl0, 0,pixels.byteSize());
	        	dataser=ocm3;//
	        }else {
	        	dataser=ocm2;
	        }
	      //  oct.getOpenCLCommandQueue().finish();
	       // tdb.mark("位平面解码");
	       
		      tdb.mark("编码");
	        serialData2=MembandBoost.allocateUninitializedByteArray(width*height*4);
		      tdb.mark("分配堆内存");

	        oct.getOpenCLCommandQueue().executeReadBuffer(dataser, dataserpseg, true,0,0,width*height*4);
		      tdb.mark("显存->内存");
	        MemorySegment.copy(dataserpseg, ValueLayout.JAVA_BYTE, 0, serialData2, 0,width*height*4);
		      tdb.mark("堆外内存->堆内存");
			
	}
	public KIFCompressedImageBlock encodeStage2() {
		 int sizeBefore=serialData2.length;
	        byte[]writeData=Compressors.compressZstd(serialData2);
	        tdb.mark("熵编码");
	        tdb.print();
			return new KIFCompressedImageBlock(writeData,sizeBefore,width,height,input.getPosx(),input.getPosy(),mode);
	}
}
