# -*- coding: utf-8 -*-
"""
HyTools:  Hyperspectral image processing library
Copyright (C) 2021 University of Wisconsin

Authors: Adam Chlus, Zhiwei Ye, Philip Townsend.

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, version 3 of the License.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.

本模块包含应用经验性 BRDF 校正的函数，如下论文所述：

方程和常数可在以下论文中找到：

"""

import numpy as np
import ray
from scipy.interpolate import interp1d
from .kernels import calc_volume_kernel,calc_geom_kernel
from ..masks import mask_create
from ..misc import progbar, pairwise
from ..misc import update_brdf
from ..plotting import flex_diagno_plot

def flex_brdf(actors,config_dict):
    brdf_dict= config_dict['brdf']
    if brdf_dict['grouped']:
        calc_flex_group(actors,brdf_dict)
    else:
        _ = ray.get([a.do.remote(calc_flex_single,brdf_dict) for a in actors])

    if "diagnostic_plots" in brdf_dict:
        if brdf_dict['diagnostic_plots']:
            print('Exporting diagnostic plots.')
            _ = ray.get([a.do.remote(flex_diagno_plot,config_dict) for a in actors])


def ndvi_stratify(hy_obj):
    '''创建 NDVI 分箱分层掩膜
    '''

    ndvi = hy_obj.ndi()
    class_mask = np.zeros((hy_obj.lines, hy_obj.columns))

    for bin_num in hy_obj.brdf['bins']:
        start,end =  hy_obj.brdf['bins'][bin_num]
        class_mask[(ndvi > start) & (ndvi <= end)] = bin_num

    class_mask[~hy_obj.mask['calc_brdf']] = 0

    #Subsample data
    idx = np.array(np.where(class_mask!=0)).T
    idxRand= idx[np.random.choice(range(len(idx)),int(len(idx)*(1-hy_obj.brdf['sample_perc'])), replace = False)].T
    class_mask[idxRand[0],idxRand[1]] = 0
    class_mask = class_mask.astype(np.int8)
    hy_obj.ancillary['ndvi_classes'] = class_mask


def ndvi_2nd_split(ndvi_bins_dynamic, all_ndvi_array, ndvi_bin_range_thres=0.15):
    ''' 执行第二次 NDVI 分割
    '''

    ndvi_bin_range_thres = -0.015625 * (len(ndvi_bins_dynamic)-1) + 0.43125
    ndvi_bin_range = np.array(ndvi_bins_dynamic[1:]) - np.array(ndvi_bins_dynamic[:-1])

    bin_for_split = np.argwhere(ndvi_bin_range>=ndvi_bin_range_thres).ravel()

    new_break = []
    if bin_for_split.shape[0]>0:
        for bin_id in bin_for_split:
            # Use median of the bin as the new break point
            new_break += [np.median(all_ndvi_array[(all_ndvi_array >  ndvi_bins_dynamic[bin_id]) & (all_ndvi_array < ndvi_bins_dynamic[bin_id+1])]).astype(np.float64)]

    # New list of bin break points
    ndvi_bins_dynamic = sorted(ndvi_bins_dynamic + new_break)

    return ndvi_bins_dynamic

def ndvi_bins(ndvi,brdf_dict):
    '''计算 NDVI 分箱范围
    '''
    perc_range = brdf_dict['ndvi_perc_max'] - brdf_dict['ndvi_perc_min'] + 1

    ndvi_break_dyn_bin = np.percentile(ndvi[ndvi > 0],
                                       np.arange(brdf_dict['ndvi_perc_min'],
                                                brdf_dict['ndvi_perc_max'] + 1,
                                                perc_range / (brdf_dict['num_bins'] - 1)))
    ndvi_thres = [brdf_dict['ndvi_bin_min']]
    ndvi_thres += ndvi_break_dyn_bin.tolist()
    ndvi_thres += [brdf_dict['ndvi_bin_max']]
    ndvi_thres = sorted(list(set(ndvi_thres)))

    # 对 NDVI 分箱进行第二次分割
    ndvi_thres = ndvi_2nd_split(ndvi_thres, ndvi)

    bins = [[x,y] for x,y in pairwise(ndvi_thres)]
    return bins

def get_kernel_samples(hy_obj):
    '''计算并采样 BRDF 核函数
    '''
    geom_kernel = hy_obj.geom_kernel(hy_obj.brdf['geometric'],
                                     b_r=hy_obj.brdf["b/r"] ,
                                     h_b =hy_obj.brdf["h/b"])
    geom_kernel = geom_kernel[hy_obj.ancillary['ndvi_classes'] !=0]

    vol_kernel = hy_obj.volume_kernel(hy_obj.brdf['volume'])
    vol_kernel = vol_kernel[hy_obj.ancillary['ndvi_classes'] !=0]

    classes = hy_obj.ancillary['ndvi_classes'][hy_obj.ancillary['ndvi_classes'] !=0]
    X = np.vstack([vol_kernel,geom_kernel,
                   np.ones(vol_kernel.shape),classes]).T
    return X

def get_band_samples(hy_obj,args):
    band = hy_obj.get_band(args['band_num'],
                           corrections = hy_obj.corrections)
    return band[hy_obj.ancillary['ndvi_classes'] !=0]

def calc_flex_single(hy_obj,brdf_dict):
    ''' 计算单个图像的 BRDF 系数
    '''
    hy_obj.brdf['coeffs'] ={}

    # 确定分箱维度并创建类别掩膜
    if hy_obj.brdf['bin_type'] == 'dynamic':
        bins = ndvi_bins(hy_obj.ndi()[hy_obj.mask['no_data']],brdf_dict)
        # 更新分箱数量
        hy_obj.brdf['num_bins'] = len(bins)
    else:
        bins = brdf_dict['bins']

    hy_obj.brdf['bins'] = {k:v for (k,v) in enumerate(bins,start=1)}
    ndvi_stratify(hy_obj)
    kernel_samples= get_kernel_samples(hy_obj)

    # 计算每个波段和类别的系数
    for band_num,band in enumerate(hy_obj.bad_bands):
        if ~band:
            hy_obj.brdf['coeffs'][band_num] = {}
            band_samples = hy_obj.do(get_band_samples, {'band_num':band_num})
            coeffs= []

            for bin_num in hy_obj.brdf['bins']:
                bin_mask = (kernel_samples[:,3] == bin_num)
                X = kernel_samples[:,:3][bin_mask]
                y = band_samples[bin_mask]
                coeffs.append(np.linalg.lstsq(X, y,rcond=-1)[0].flatten().tolist())
            hy_obj.brdf['coeffs'][band_num]  = coeffs

def calc_flex_group(actors,brdf_dict):
    ''' 计算一组图像的 BRDF 系数
    '''
    # 从图像聚合 NDVI 值
    ndvi = ray.get([a.ndi.remote(mask = 'no_data') for a in actors])
    ndvi = np.concatenate([n.flatten() for n in ndvi])

    # 确定分箱维度
    if  brdf_dict['bin_type'] == 'dynamic':
        bins = ndvi_bins(ndvi,brdf_dict)
        # 更新分箱数量
        _ = ray.get([a.do.remote(update_brdf,{'key':'num_bins',
                                              'value': len(bins)}) for a in actors])
    else:
        bins = brdf_dict['bins']

    bins  = {k:v for (k,v) in enumerate(bins,start=1)}

    # 更新 BRDF 系数
    _ = ray.get([a.do.remote(update_brdf,{'key':'bins',
                                          'value': bins}) for a in actors])

    # 创建 NDVI 类别掩膜并采样核函数
    _ = ray.get([a.do.remote(ndvi_stratify) for a in actors])
    kernel_samples = ray.get([a.do.remote(get_kernel_samples) for a in actors])
    kernel_samples = np.concatenate(kernel_samples)

    bad_bands = ray.get(actors[0].do.remote(lambda x: x.bad_bands))
    coeffs = {}

    for band_num,band in enumerate(bad_bands):
        if ~band:
            coeffs[band_num] = {}
            band_samples = ray.get([a.do.remote(get_band_samples,
                                     {'band_num':band_num}) for a in actors])
            band_samples = np.concatenate(band_samples)
            band_coeffs= []
            for bin_num in bins:
                bin_mask = (kernel_samples[:,3] == bin_num)
                X = kernel_samples[:,:3][bin_mask]
                y = band_samples[bin_mask]
                band_coeffs.append(np.linalg.lstsq(X, y,rcond=-1)[0].flatten().tolist())
            coeffs[band_num]  = band_coeffs
            progbar(np.sum(~bad_bands[:band_num+1]),np.sum(~bad_bands))

    print('\n')

    # 更新 BRDF 系数
    _ = ray.get([a.do.remote(update_brdf,{'key':'coeffs',
                                          'value': coeffs}) for a in actors])

def apply_flex(hy_obj,data,dimension,index):
    ''' 对数据切片应用 flex BRDF 校正

    参数:
        hy_obj : Hytools 类对象。
        data (np.ndarray): 数据切片。
        index (int,list): 数据索引。

    返回:
        data (np.ndarray): BRDF 校正后的数据切片。
    '''

    if 'k_vol' not in hy_obj.ancillary:
        hy_obj.ancillary['k_vol'] = hy_obj.volume_kernel(hy_obj.brdf['volume'])
    if 'k_geom' not in hy_obj.ancillary:
        hy_obj.ancillary['k_geom'] = hy_obj.geom_kernel(hy_obj.brdf['geometric'],
                                     b_r=hy_obj.brdf["b/r"],
                                     h_b =hy_obj.brdf["h/b"])
    if ('k_vol_nadir' not in hy_obj.ancillary) or ('k_geom_nadir' not in hy_obj.ancillary):
        solar_zn = hy_obj.brdf['solar_zn_norm_radians']  * np.ones((hy_obj.lines,hy_obj.columns))
        hy_obj.ancillary['k_vol_nadir']  = calc_volume_kernel(0,solar_zn,
                                                               0,0,hy_obj.brdf['volume'])
        hy_obj.ancillary['k_geom_nadir']  = calc_geom_kernel(0,solar_zn,
                                                             0,0,hy_obj.brdf['geometric'],
                                                             b_r=hy_obj.brdf["b/r"],
                                                             h_b =hy_obj.brdf["h/b"])
    if 'apply_brdf' not in hy_obj.mask:
        hy_obj.gen_mask(mask_create,'apply_brdf',hy_obj.brdf['apply_mask'])

    if 'ndvi' not in hy_obj.ancillary:
        hy_obj.ancillary['ndvi'] =  hy_obj.ndi()

    if 'interpolators' not in hy_obj.ancillary:
        bin_centers = np.mean(list(hy_obj.brdf['bins'].values()),axis=1)
        hy_obj.ancillary['interpolators'] ={}

        # 生成插值器
        for i in hy_obj.brdf['coeffs']:
            coeffs= np.array(hy_obj.brdf['coeffs'][i])
            interpolator = interp1d(bin_centers, coeffs, kind = hy_obj.brdf['interp_kind'],
                                    axis=0,fill_value="extrapolate")
            hy_obj.ancillary['interpolators'][int(i)] = interpolator

    # 转换为浮点数
    data = data.astype(np.float32)
    brdf_bands = [int(x) for x in hy_obj.ancillary['interpolators']]
    if dimension == 'line':
        # index= 3000
        # data = hy_obj.get_line(3000)

        interpolated_f = [hy_obj.ancillary['interpolators'][band](hy_obj.ancillary['ndvi'][index,:]) for band in brdf_bands]
        interpolated_f = np.array(interpolated_f)
        fvol, fgeo, fiso  = interpolated_f[:,:,0], interpolated_f[:,:,1], interpolated_f[:,:,2]

        brdf = fvol*hy_obj.ancillary['k_vol'][index,:]
        brdf+= fgeo*hy_obj.ancillary['k_geom'][index,:]
        brdf+= fiso

        brdf_nadir = fvol*hy_obj.ancillary['k_vol_nadir'][index,:]
        brdf_nadir+= fgeo*hy_obj.ancillary['k_geom_nadir'][index,:]
        brdf_nadir+= fiso

        correction_factor = brdf_nadir/brdf
        correction_factor[:,~hy_obj.mask['apply_brdf'][index]] = 1

        data[:,brdf_bands] = data[:,brdf_bands]*correction_factor.T

    elif dimension == 'column':
        #index= 300
        #data = hy_obj.get_column(index)

        interpolated_f = [hy_obj.ancillary['interpolators'][band](hy_obj.ancillary['ndvi'][:,index]) for band in brdf_bands]
        interpolated_f = np.array(interpolated_f)
        fvol, fgeo, fiso  = interpolated_f[:,:,0], interpolated_f[:,:,1], interpolated_f[:,:,2]

        brdf = fvol*hy_obj.ancillary['k_vol'][:,index]
        brdf+= fgeo*hy_obj.ancillary['k_geom'][:,index]
        brdf+= fiso

        brdf_nadir = fvol*hy_obj.ancillary['k_vol_nadir'][:,index]
        brdf_nadir+= fgeo*hy_obj.ancillary['k_geom_nadir'][:,index]
        brdf_nadir+= fiso

        correction_factor = brdf_nadir/brdf
        correction_factor = np.moveaxis(correction_factor,0,1)
        correction_factor[:,~hy_obj.mask['apply_brdf'][index]] = 1
        data[:,brdf_bands] = data[:,brdf_bands]*correction_factor.T

    elif (dimension == 'band') & (index in brdf_bands):
        # index= 8
        # data = hy_obj.get_band(index)

        interpolated_f = hy_obj.ancillary['interpolators'][index](hy_obj.ancillary['ndvi'])
        fvol, fgeo, fiso  = interpolated_f[:,:,0], interpolated_f[:,:,1], interpolated_f[:,:,2]

        brdf = fvol*hy_obj.ancillary['k_vol']
        brdf += fgeo*hy_obj.ancillary['k_geom']
        brdf += fiso

        brdf_nadir = fvol*hy_obj.ancillary['k_vol_nadir']
        brdf_nadir += fgeo*hy_obj.ancillary['k_geom_nadir']
        brdf_nadir += fiso

        correction_factor = brdf_nadir/brdf
        correction_factor[~hy_obj.mask['apply_brdf']] = 1
        data= data* correction_factor

    elif dimension == 'chunk':
        # index = 200,501,3000,3501
        x1,x2,y1,y2 = index
        # data = hy_obj.get_chunk(x1,x2,y1,y2)

        interpolated_f = [hy_obj.ancillary['interpolators'][band](hy_obj.ancillary['ndvi'][y1:y2,x1:x2]) for band in brdf_bands]
        interpolated_f = np.array(interpolated_f)
        interpolated_f = np.swapaxes(interpolated_f,0,-1)
        fvol, fgeo, fiso  = interpolated_f[0,:,:,:], interpolated_f[1,:,:,:], interpolated_f[2,:,:,:]

        brdf = fvol*hy_obj.ancillary['k_vol'][y1:y2,x1:x2,np.newaxis]
        brdf+= fgeo*hy_obj.ancillary['k_geom'][y1:y2,x1:x2,np.newaxis]
        brdf+= fiso

        brdf_nadir = fvol*hy_obj.ancillary['k_vol_nadir'][y1:y2,x1:x2,np.newaxis]
        brdf_nadir+= fgeo*hy_obj.ancillary['k_geom_nadir'][y1:y2,x1:x2,np.newaxis]
        brdf_nadir+= fiso

        correction_factor = brdf_nadir/brdf
        correction_factor[~hy_obj.mask['apply_brdf'][y1:y2,x1:x2]] = 1
        data[:,:,brdf_bands] = data[:,:,brdf_bands]*correction_factor

    elif dimension == 'pixels':
        # index = [[2000,2001],[200,501]]
        y,x = index
        # data = hy_obj.get_pixels(y,x)

        interpolated_f = [hy_obj.ancillary['interpolators'][band](hy_obj.ancillary['ndvi'][y,x]) for band in brdf_bands]
        interpolated_f = np.array(interpolated_f)
        interpolated_f = np.swapaxes(interpolated_f,0,1)
        fvol, fgeo, fiso  = interpolated_f[:,:,0], interpolated_f[:,:,1], interpolated_f[:,:,2]

        brdf = fvol*hy_obj.ancillary['k_vol'][y,x,np.newaxis]
        brdf+= fgeo*hy_obj.ancillary['k_geom'][y,x,np.newaxis]
        brdf+= fiso

        brdf_nadir = fvol*hy_obj.ancillary['k_vol_nadir'][y,x,np.newaxis]
        brdf_nadir+= fgeo*hy_obj.ancillary['k_geom_nadir'][y,x,np.newaxis]
        brdf_nadir+= fiso

        correction_factor = brdf_nadir/brdf
        correction_factor[~hy_obj.mask['apply_brdf'][y,x]] = 1
        data[:,brdf_bands] = data[:,brdf_bands]*correction_factor

    return data