# -*- coding: UTF-8 -*-
"""
@Project:__init__.py
@File:DualPolarToPolsarproC2.py
@Function:双极化影像转成polsarpro格式C2数据
@Contact:
@Author:SHJ
@Date:2021/11/5
@Version:1.0.0
"""
import os
import numpy as np
import glob
import struct
import gc
from tool.algorithm.image.ImageHandle import ImageHandler


class DualPolarToPolsarproC2:
    """
    双极化影像转换为bin格式C2矩阵，支持polsarpro处理
    """

    def __init__(self):
        pass

    @staticmethod
    def __dual_polar_to_c2(dual_polar_dir):
        """
        双影像转S2矩阵
        :param dual_polar_dir: 双极化影像文件夹路径
        :return: C2矩阵
        """
        in_tif_paths = list(glob.glob(os.path.join(dual_polar_dir, '*.tif')))

        if in_tif_paths == []:
            in_tif_paths = list(glob.glob(os.path.join(dual_polar_dir, '*.tiff')))
        s11, s22 = None, None
        flag_list = [0, 0, 0, 0]
        for in_tif_path in in_tif_paths:
            # 读取原始SAR影像
            proj, geotrans, data = ImageHandler.read_img(in_tif_path)
            # 获取极化类型
            if 'HH' in os.path.basename(in_tif_path):
                s11 = data[0, :, :] + 1j * data[1, :, :]
                flag_list[0] = 1
            elif 'HV' in os.path.basename(in_tif_path):
                s22 = data[0, :, :] + 1j * data[1, :, :]
                flag_list[1] = 1
            elif 'VH' in os.path.basename(in_tif_path):
                s22 = data[0, :, :] + 1j * data[1, :, :]
                flag_list[2] = 1
            elif 'VV' in os.path.basename(in_tif_path):
                s11 = data[0, :, :] + 1j * data[1, :, :]
                flag_list[3] = 1
            else:
                continue
            del data
            gc.collect()

        if flag_list != [1, 1, 0, 0] and flag_list != [0, 0, 1, 1] :
            raise Exception('Dual-Polarization SAR is not in path :%s',in_tif_path)

        c11,c12,c22 = None, None, None
        c11 = np.abs(s11)** 2
        c12 = s11 * np.conj(s22)
        del s11
        gc.collect()
        c22 = np.abs(s22)**2
        return c11, c12, c22

    def __c2_to_polsarpro_c2(self, out_dir, c11, c12, c22):
        """
        C2矩阵转bin格式，支持 polsarpro处理
        :param out_dir: 输出的文件夹路径
        :param c11:
        :param c12:
        :param c21:
        :param c22:
        :return: bin格式矩阵C3和头文件
        """
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

        rows = c11.shape[0]
        cols = c11.shape[1]
        bins_dict = {
            'C11.bin': c11,
            'C12_real.bin': c12.real,
            'C12_imag.bin': c12.imag,
            'C22.bin': c22}

        for name, data in bins_dict.items():
            bin_path = os.path.join(out_dir, name)
            self.__write_img_bin(data, bin_path)
            out_hdr_path = bin_path + '.hdr'
            self.__write_bin_hdr(out_hdr_path, bin_path, rows, cols)

        self.__write_config_file(out_dir, rows, cols)

    def rows(self):
        """获取影像行数"""
        return self._rows

    def cols(self):
        """获取影像列数"""
        return self._cols

    def __write_img_bin(self, im, file_path):
        """
        写入影像到bin文件中，保存为float32类型
        :param im : 影像矩阵数据，暂支持单通道影像数据
        :param file_path: bin文件的完整路径
        """
        with open(file_path, 'wb') as f:
            self._rows = im.shape[0]
            self._cols = im.shape[1]
            for row in range(self._rows):
                im_bin = struct.pack("f" * self._cols, *np.reshape(im[row, :], (self._cols, 1), order='F'))
                f.write(im_bin)
        f.close()

    @staticmethod
    def __write_bin_hdr(out_hdr_path, bin_path, rows, cols):
        """
        写入影像的头文件
        :param out_hdr_path : 头文件的路径
        :param bin_path: bin文件的路径
        :param rows: 影像的行数
        :param cols: 影像的列数
        """
        name = os.path.split(bin_path)[1]
        h1 = 'ENVI'
        h2 = 'description = {'
        h3 = 'File Imported into ENVI. }'
        h4 = 'samples = ' + str(cols)  # 列
        h5 = 'lines   = ' + str(rows)  # 行
        h6 = 'bands   = 1  '  # 波段数
        h7 = 'header offset = 0'
        h8 = 'file type = ENVI Standard'
        h9 = 'data type = 4'  # 数据格式 浮点型
        h10 = 'interleave = bsq'  # 存储格式
        h11 = 'sensor type = Unknown'
        h12 = 'byte order = 0'
        h13 = 'band names = {'
        h14 = name + '}'
        h = [h1, h2, h3, h4, h5, h6, h7, h8, h9, h10, h11, h12, h13, h14]
        doc = open(out_hdr_path, 'w')
        for i in range(0, 14):
            print(h[i], end='', file=doc)
            print('\n', end='', file=doc)
        doc.close()

    @staticmethod
    def __write_config_file(out_config_dir, rows, cols):
        """
        写入polsarpro配置文件
        :param out_config_dir : 配置文件路径
        :param rows: 影像的行数
        :param cols: 影像的列数
        """
        h1 = 'Nrow'
        h2 = str(rows)
        h3 = '---------'
        h4 = 'Ncol'
        h5 = str(cols)
        h6 = '---------'
        h7 = 'PolarCase'
        h8 = 'monostatic'
        h9 = '---------'
        h10 = 'PolarType'
        h11 = 'pp1'
        h = [h1, h2, h3, h4, h5, h6, h7, h8, h9, h10, h11]

        out_config_path = os.path.join(out_config_dir, 'config.txt')
        doc = open(out_config_path, 'w')
        for i in range(0, 11):
            print(h[i], end='', file=doc)
            print('\n', end='', file=doc)
        doc.close()


    def api_dual_polar__to_polsarpro_c2(self, out_file_dir, dual_polar_dir):
        c11, c12, c22 = self.__dual_polar_to_c2(dual_polar_dir)
        self.__c2_to_polsarpro_c2(out_file_dir,c11, c12, c22)


# if __name__ == '__main__':
#     tp = DualPolarToPolsarproC2()
#     out_dic = 'E:\\3-GF3_KAS_FSI_020253_E110.8_N25.5_20200614_L1A_HHHV_L10004871459\\SLC_SHJ1'
#     in_dic = 'E:\\3-GF3_KAS_FSI_020253_E110.8_N25.5_20200614_L1A_HHHV_L10004871459\\'
#     # out_file_path = 'D:\\bintest0923\\'
#     tp.api_dual_polar__to_polsarpro_c2(out_dic,in_dic)
#     # atp.ahv_to_polsarpro_t3(out_file_path, ahv_path)
#
#     print("done")