import numpy as np
import xarray as xr

import ancillary_data as anc
import utils

_dtr = np.pi/180


# this case is written for the 11-12um Cirrus Test for scenes that follow pattern 1 (see note below)
def prepare_thresholds(data, thresholds):

    coeff_values = np.empty((data.M01.shape[0], data.M01.shape[1], 2))
    coeff_values[:, :, 0] = np.full(data.M01.shape, thresholds['11-12um_Cirrus_Test']['coeffs'][0])
    coeff_values[:, :, 1] = np.full(data.M01.shape, thresholds['11-12um_Cirrus_Test']['coeffs'][1])
    cmult_values = np.full(data.M01.shape, thresholds['11-12um_Cirrus_Test']['cmult'])
    adj_values = np.full(data.M01.shape, thresholds['11-12um_Cirrus_Test']['adj'])

    thr_dict = {'coeffs': (['number_of_lines', 'number_of_pixels', 'z'], coeff_values),
                'cmult': (['number_of_lines', 'number_of_pixels'], cmult_values),
                'adj': (['number_of_lines', 'number_of_pixels'], adj_values)
                }

    return xr.Dataset(data_vars=thr_dict)


def preproc(data, thresholds):
    cosvza = np.cos(data.sensor_zenith * _dtr)
    schi = (1/cosvza).where(cosvza > 0, 99.0)

    schi = schi.values.reshape(np.prod(schi.shape))
    m15 = data.M15.values.reshape(np.prod(data.M15.shape))
    thr = anc.py_cithr(1, schi, m15)
    thr = thr.reshape(data.M15.shape)
    schi = schi.reshape(data.M15.shape)

    # thr_xr = xr.Dataset(np.full(data.sensor_zenith.shape, thresholds['coeffs']),
    #                     dims=('number_of_lines', 'number_of_pixels'))
    thr_xr = prepare_thresholds(data, thresholds)

    midpt = thr_xr.coeffs[:, :, 0].where((thr < 0.1) | (np.abs(schi-99) < 0.0001), thr)
    locut = midpt + (thr_xr.cmult * midpt)
    hicut = midpt - thr_xr.adj
    # this below is for the method 2 of computing hicut
    # hicut = midpt - (thr_xr.adj * midpt)

    thr_out = xr.DataArray(data=np.dstack((locut, midpt, hicut, np.ones(locut.shape), np.ones(locut.shape))),
                           dims=('number_of_lines', 'number_of_pixels', 'z'))
    return thr_out
    # return locut, hicut, midpt


def preproc_sst(data, thresholds):
    m31c = data.M15 - 273.16
    m32c = data.M16 - 273.16
    m31c_m32c = m31c - m32c
    sstc = data.geos_sfct - 273.16
    cosvza = np.cos(data.sensor_zenith*_dtr)

    a = thresholds['coeffs']

    modsst = 273.16 + a[0] + a[1]*m31c + a[2]*m31c_m32c*sstc + a[3]*m31c_m32c*((1/cosvza) - 1)
    sfcdif = data.geos_sfct - modsst

    return sfcdif


def preproc_nir(data, thresholds, scene):

    sza = data.solar_zenith
    band_n = 2
    # NOTE: the visud condition in the C code is equivalent to having sza <= 85
    # For the time being the visud filtering is not implemented

    a = np.array(thresholds[scene]['NIR_Reflectance_Test']['coeffs'])
    vzcpow = thresholds['VZA_correction']['vzcpow'][0]
    refang = data.sunglint_angle.values.reshape(np.prod(data.sunglint_angle.shape))
    sunglint_thresholds = thresholds['Sun_Glint']
    sunglint_flag = utils.sunglint_scene(refang, sunglint_thresholds).reshape(refang.shape)
    nir_thresh = thresholds[scene]['NIR_Reflectance_Test']

    hicut0 = a[0] + a[1]*sza + a[2]*np.power(sza, 2) + a[3]*np.power(sza, 3)
    hicut0 = (hicut0 * 0.01) + nir_thresh['adj']
    hicut0 = (hicut0 + nir_thresh['bias']).values.reshape(refang.shape)
    midpt0 = hicut0 + (nir_thresh['midpt_coeff'] * nir_thresh['bias'])
    locut0 = midpt0 + (nir_thresh['locut_coeff'] * nir_thresh['bias'])
    thr = np.array([locut0, midpt0, hicut0, nir_thresh['thr'][3]*np.ones(refang.shape)])

    cosvza = np.cos(data.sensor_zenith*_dtr).values.reshape(refang.shape)
    corr_thr = np.zeros((4, refang.shape[0]))

    corr_thr[:3, sunglint_flag == 0] = thr[:3, sunglint_flag == 0] * (1./np.power(cosvza[sunglint_flag == 0],
                                                                                  vzcpow))
    corr_thr[3, sunglint_flag == 0] = thr[3, sunglint_flag == 0]

    for flag in range(1, 4):
        if len(refang[sunglint_flag == flag]) > 0:
            dosgref = utils.get_sunglint_thresholds(refang, sunglint_thresholds, band_n, flag, thr)
            corr_thr[:3, sunglint_flag == flag] = dosgref[:3, sunglint_flag == flag] * \
                (1./np.power(cosvza[sunglint_flag == flag], vzcpow))
            corr_thr[3, sunglint_flag == flag] = dosgref[3, sunglint_flag == flag]

    corr_thr = np.transpose(corr_thr.reshape((4, sza.shape[0], sza.shape[1])), (1, 2, 0))

    return corr_thr


# NOTE: 11-12um Cirrus Test
# hicut is computed in different ways depending on the scene
# 1. midpt - adj
# - Land_Day
# - Land_Day_Coast
# - Land_Day_Desert
# - Land_Day_Desert_Coast
# - Ocean_Day
# - Ocean_Night
# - Polar_Day_Ocean
# - Polar_Night_Ocean
#
# 2. midpt - (btd_thr * adj)
# - Polar_Day_Land
# - Polar_Day_Coast
# - Polar_Day_Desert
# - Polar_Day_Desert_Coast
# - Polar_Day_Snow
#
# 3. Others
# - Land_Night
# - Polar_Night_Land
# - Polar_Night_Snow
# - Day_Snow
# - Night_Snow

# NOTE: 1.38um High Cloud Test
# thresholds are not always computed the same way. In group 1 there's no preprocessing required,
# in group 2 some calcuations are needed
# 1.
# - Land_Day
# - Land_Day_Coast
# - Land_Day_Desert
# - Land_Day_Desert_Coast
# - Polar_Day_Land
# - Polar_Day_Coast
# - Polar_Day_Desert
# - Polar_Day_Desert_Coast
# - Polar_Day_Snow
# - Day_Snow
#
# 2.
# - Ocean_Day
# - Polar_Ocean_Day