# -*- coding: UTF-8 -*-
"""
@Project:__init__.py 
@File:pspTouziDecomposition.py
@Function:
@Contact:
@Author:LVY
@Date:2021/10/14 10:11
@Version:1.0.0
"""
import os
import logging
from tool.algorithm.polsarpro.polarizationDecomposition import ModTouzi as TouziDecomp
from osgeo import gdal
import multiprocessing
from tool.algorithm.block.blockprocess import BlockProcess
import shutil
from tool.algorithm.image.ImageHandle import ImageHandler
from tool.file.fileHandle import fileHandle
logger = logging.getLogger("mylog")
file =fileHandle(False)

class PspTouziDecomposition:
    """
    Touzi分解
    """
    def __init__(self, inDic, outDir):
        """
        :param inDic:T3矩阵目录
        :param outDir:输出目录
        """
        self.__inDic = inDic
        self.__outDir = outDir
        self.__DecompostFlag = False
        if self._checkTifFileDic(self.__inDic) is False:
            return False
        if not os.path.exists(self.__outDir):
            os.makedirs(self.__outDir)

    def api_Touzi_decomposition_TIF(self, Nwin = 5):
        """
        :param Nwin:滤波窗口大小 3 5 7 9 11
        """
        bandHH = gdal.Open(self.__inDic["HH"])
        bandHV = gdal.Open(self.__inDic["HV"])
        bandVH = gdal.Open(self.__inDic["VH"])
        bandVV = gdal.Open(self.__inDic["VV"])
        bandAll = [bandHH, bandHV, bandVH, bandVV]
        decomposition = TouziDecomp(bandAll, Nwin)
        decomposition.get_result(self.__outDir)
        return True

    def Touzi_decomposition_TIF(self,hh_path,hv_path,vh_path,vv_path,out_dir,suffix,Nwin = 5):
        """
        :param Nwin:滤波窗口大小 3 5 7 9 11
        """
        bandHH = gdal.Open(hh_path)
        bandHV = gdal.Open(hv_path)
        bandVH = gdal.Open(vh_path)
        bandVV = gdal.Open(vv_path)
        bandAll = [bandHH, bandHV, bandVH, bandVV]
        decomposition = TouziDecomp(bandAll, Nwin)
        decomposition.get_result_block(out_dir, suffix)
        return True

    @staticmethod
    def _checkTifFileDic(inDic):
        file_name_in_out = ['HH', 'VV', 'HV', 'VH']
        for item in file_name_in_out:
            if item in inDic:
                print(inDic[item])
                if not os.path.exists(os.path.join(inDic[item])):
                    return False
            else:
                return False
        return True

    def Touzi_decomposition_multiprocessing(self):

        #创建工作文件夹
        src_path = os.path.join(self.__outDir, "src_img")
        block_path = os.path.join(self.__outDir, "block")
        decomposition_path = os.path.join(self.__outDir, "feature")
        file.creat_dirs([src_path,block_path,decomposition_path])

        shutil.copyfile(self.__inDic["HH"], os.path.join(src_path, "HH.tif"))
        shutil.copyfile(self.__inDic["HV"], os.path.join(src_path, "HV.tif"))
        shutil.copyfile(self.__inDic["VH"], os.path.join(src_path, "VH.tif"))
        shutil.copyfile(self.__inDic["VV"], os.path.join(src_path, "VV.tif"))
        self.__cols = ImageHandler.get_img_width(self.__inDic["HH"])
        self.__rows = ImageHandler.get_img_height(self.__inDic["HH"])
        # 分块
        bp = BlockProcess()
        block_size = bp.get_block_size(self.__rows, self.__cols)
        bp.cut(src_path, block_path, ['tif', 'tiff'], 'tif', block_size)
        logger.info('blocking tifs success!')

        img_dir, img_name = bp.get_file_names(block_path, ['tif'])
        dir_dict = bp.get_same_img(img_dir, img_name)

        hh_list, vv_list, hv_list, vh_list = None, None, None, None
        for key in dir_dict.keys():
            tmp = key.split('_', 2)[0]
            if tmp == 'HH':
                hh_list = dir_dict[key]
            elif tmp == 'VV':
                vv_list = dir_dict[key]
            elif tmp == 'HV':
                hv_list = dir_dict[key]
            elif tmp == 'VH':
                vh_list = dir_dict[key]

        processes_num = min([len(hh_list), multiprocessing.cpu_count() - 1])

        # 开启多进程处理
        pool = multiprocessing.Pool(processes=processes_num)

        for i in range(len(hh_list)):
            suffix = bp.get_suffix(os.path.basename(hh_list[i]))
            # self.Touzi_decomposition_TIF(hh_list[i], hv_list[i], vh_list[i], vv_list[i], block_path, suffix,5)
            pool.apply_async(self.Touzi_decomposition_TIF, (hh_list[i], hv_list[i], vh_list[i], vv_list[i], decomposition_path, suffix,5))
            logger.info('total:%s, block:%s touzi!', len(hh_list), i)

        pool.close()
        pool.join()
        # 合并处理后的影像
        bp.combine(decomposition_path, self.__cols, self.__rows, self.__outDir, file_type=['tif'], datetype='float16')

        file.del_folder(src_path)
        file.del_folder(block_path)
        file.del_folder(decomposition_path)
        pass


# if __name__ == '__main__':
#     dir = {}
#     dir.update({"HH":"I:\preprocessed\HH_preprocessed.tif"})
#     dir.update({"HV":"I:\preprocessed\HV_preprocessed.tif"})
#     dir.update({"VH":"I:\preprocessed\VH_preprocessed.tif"})
#     dir.update({"VV":"I:\preprocessed\VV_preprocessed.tif"})
#
#
#     p = PspTouziDecomposition(dir, "I:/preprocessed/")
#     p.Touzi_decomposition_multiprocessing()
#     pass