import org.gdal.gdal.Band;
import org.gdal.gdal.Dataset;
import org.gdal.gdal.gdal;
import org.gdal.gdalconst.gdalconstConstants;
import org.opencv.core.Core;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.imgcodecs.Imgcodecs;

import java.util.ArrayList;

/**
 * @Author: HNUST_jue_chen
 * @Date: 2022/10/28/ 11:01
 * @Attention: 转载, 引用请注明出处
 */

//假彩色合成
public class FalseColorSynthesis {
    //加载本地动态链接库
    static {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);
    }

    //获得波段4,3,2
    public ArrayList<int[][]> getBands(String path) {
        //支持所有驱动
        gdal.AllRegister();
        //以只读方式读取数据存入Dataset类型里
        Dataset dataset = gdal.Open(path, gdalconstConstants.GA_ReadOnly);
        //获得图像的大小
        int rows = dataset.getRasterYSize();
        int cols = dataset.getRasterXSize();
        int count = dataset.getRasterCount();
        System.out.println(rows);
        System.out.println(cols);
        System.out.println(count);

        //获得需要的三个波段
        Band band4 = dataset.GetRasterBand(4);
        Band band3 = dataset.GetRasterBand(3);
        Band band2 = dataset.GetRasterBand(2);

        //定义波段三个二维数组存放三个波段的像素
        int[][] band4_arr = new int[rows][cols];
        int[][] band3_arr = new int[rows][cols];
        int[][] band2_arr = new int[rows][cols];

        //将三个波段的像素的像素存放入二维数组
        for (int i = 0; i < rows; i++) {
            //一次读取一行像素值
            band4.ReadRaster(0, i, cols, 1, band4_arr[i]);
            band3.ReadRaster(0, i, cols, 1, band3_arr[i]);
            band2.ReadRaster(0, i, cols, 1, band2_arr[i]);
        }

        //将三个二维数组存放在集合中
        ArrayList<int[][]> band_arr = new ArrayList<>();
        band_arr.add(band4_arr);
        band_arr.add(band3_arr);
        band_arr.add(band2_arr);

        return band_arr;
    }

    //获得本图像像素的最大值和最小值
    public int[] getMinMaxPixel(int[][] arr) {
        //先将数组的第一个数赋予最大值和最小值
        int min = arr[0][0];
        int max = arr[0][0];
        for (int i = 0; i < arr.length; i++) {
            for (int j = 0; j < arr[0].length; j++) {
                //更新最小值
                if (arr[i][j] < min) {
                    min = arr[i][j];
                }
                //更新最大值
                if (arr[i][j] > max) {
                    max = arr[i][j];
                }
            }
        }

        //将最大值和最小值存入数组
        int[] minMaxArr = new int[2];
        minMaxArr[0] = min;
        minMaxArr[1] = max;
        return minMaxArr;
    }

    //将16位像素转换为8位像素，线性拉伸[a,b]到[c,d]
    //a=min,b=max,c=1,d=255
    //新像素值=(d-c)/(max-min)*(原像素值-min)+1
    public int[][] linearTo8Bit(int[][] arr, int min, int max) {
        //定义线性拉伸后的数组
        int[][] arr_linear = new int[arr.length][arr[0].length];
        for (int i = 0; i < arr.length; i++) {
            for (int j = 0; j < arr[0].length; j++) {
                arr_linear[i][j] = (int) ((double) (255 - 1) / (max - min) * (arr[i][j] - min) + 1);
            }
        }
        return arr_linear;
    }

    //将三个波段合成为假彩色图像
    public Mat falseColorSynthesis(ArrayList<int[][]> arrayList) {
        //获得三个波段的数组
        int[][] band4_arr = arrayList.get(0);
        int[][] band3_arr = arrayList.get(1);
        int[][] band2_arr = arrayList.get(2);

        //定义存储三个波段的三通道数组
        int[][] image_arr = new int[band4_arr.length][band4_arr[0].length * 3];

        //Mat类三通道的的顺序为B,G,R,在标准假彩色合成中应该分别赋予band2,band3,band4
        for (int i = 0; i < image_arr.length; i++) {
            //将波段2赋予蓝色通道
            for (int j = 0; j < band2_arr[0].length; j++) {
                image_arr[i][3 * j] = band2_arr[i][j];
            }
            //将波段3赋予绿色通道
            for (int j = 0; j < band3_arr[0].length; j++) {
                image_arr[i][3 * j + 1] = band3_arr[i][j];
            }
            //将波段4赋予红色通道
            for (int j = 0; j < band4_arr[0].length; j++) {
                image_arr[i][3 * j + 2] = band4_arr[i][j];
            }
        }

        //定义假彩色图像
        Mat mat = new Mat(band2_arr.length, band2_arr[0].length, CvType.CV_32SC3);
        //将像素放入图像
        for (int i = 0; i < image_arr.length; i++) {
            //一次放入一行三通道像素值
            mat.put(i, 0, image_arr[i]);
        }
        return mat;
    }

    public static void main(String[] args) {
        FalseColorSynthesis fsc = new FalseColorSynthesis();
        ArrayList<int[][]> band_arr = fsc.getBands("D:\\Project\\IDEA_Project\\RS01\\src\\rs01\\img\\13.tif");

        //得到波段4的最小最大像素值
        int[][] band4_arr = band_arr.get(0);
        int[] band4_minMax = fsc.getMinMaxPixel(band4_arr);
        int band4_min = band4_minMax[0];
        int band4_max = band4_minMax[1];

        //得到波段3的最小最大像素值
        int[][] band3_arr = band_arr.get(1);
        int[] band3_minMax = fsc.getMinMaxPixel(band3_arr);
        int band3_min = band3_minMax[0];
        int band3_max = band3_minMax[1];

        //得到波段2的最小最大像素值
        int[][] band2_arr = band_arr.get(2);
        int[] band2_minMax = fsc.getMinMaxPixel(band2_arr);
        int band2_min = band2_minMax[0];
        int band2_max = band2_minMax[1];

        //波段4线性拉伸
        int[][] band4_arrLinear = fsc.linearTo8Bit(band4_arr, band4_min, band4_max);
        //波段3线性拉伸
        int[][] band3_arrLinear = fsc.linearTo8Bit(band3_arr, band3_min, band3_max);
        //波段2线性拉伸
        int[][] band2_arrLinear = fsc.linearTo8Bit(band2_arr, band2_min, band2_max);

        //将线性拉伸后的数组存放在集合中
        ArrayList<int[][]> band_arrLinear = new ArrayList<>();
        band_arrLinear.add(band4_arrLinear);
        band_arrLinear.add(band3_arrLinear);
        band_arrLinear.add(band2_arrLinear);

        //合成假彩色图像
        Mat mat = fsc.falseColorSynthesis(band_arrLinear);
        //将假彩色图像写入文件
        Imgcodecs.imwrite("D:\\Project\\IDEA_Project\\RS01\\src\\rs01\\img\\13_fsc.png", mat);
    }
}