"""
This code is part of the Python PolSARpro software:

"A re-implementation of selected PolSARPro functions in Python,
following the scientific recommendations of PolInSAR 2021"

developed within an ESA funded project with SATIM.

Author: Olivier D'Hondt, 2026.
Scientific advisors: Armando Marino and Eric Pottier.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

-----

# Description: module containing physical inversion functions

"""
import dask.array as da
import numpy as np
import xarray as xr

from polsarpro.auxil import validate_dataset
from polsarpro.util import C4_to_C3, S_to_C3, T3_to_C3, T4_to_C3


def dubois_surface_inversion(
    input_data: xr.Dataset,
    incidence_angle: xr.DataArray,
    freq_ghz: float,
    thresh1: float,
    thresh2: float,
    calibration_coeff: float | None = None,
) -> xr.Dataset:
    """Run the Dubois surface inversion on a PolSAR covariance dataset.

    The function accepts a single-scattering input product in any of the
    supported polarimetric representations, converts it to C3 form when
    needed, and applies the Dubois empirical model using the supplied
    incidence-angle raster.

    Args:
        input_data (xr.Dataset): Input polarimetric dataset. Supported
            products are:

            - "S": Sinclair scattering matrix

            - "C3": Lexicographic covariance matrix (3x3)

            - "T3": Pauli coherency matrix (3x3)

            - "C4": 4x4 covariance matrix

            - "T4": 4x4 coherency matrix

            The dataset must share the same spatial grid as
            ``incidence_angle``.
        incidence_angle (xr.DataArray): Incidence angle raster in radians.
            Values must be numeric and in the range ``(0, pi/2)``.
        freq_ghz (float): Radar center frequency in GHz.
        thresh1 (float): Maximum allowed ``HV / VV`` ratio in dB for the
            Dubois validity mask.
        thresh2 (float): Maximum allowed ``HH / VV`` ratio in dB for the
            Dubois validity mask.
        calibration_coeff (float or None, optional): Optional multiplicative
            calibration coefficient. When provided, the HH, VV, and HV
            channels are scaled by ``sin(theta) / calibration_coeff`` before
            the inversion is applied. Defaults to None.

    Returns:
        xr.Dataset: Dataset containing the Dubois estimates and masks:
            ``dubois_ks`` (surface roughness parameter), ``dubois_er``
            (relative dielectric constant), ``dubois_mv`` (volumetric moisture
            estimate), ``dubois_mask_in`` (input validity mask),
            ``dubois_mask_out`` (model validity mask), and
            ``dubois_mask_valid_in_out`` (combined mask).

    Notes:
        The returned dataset preserves the input coordinates and uses the
        input spatial dimensions. Output variables are stored as ``float32``
        arrays.
    """
    if not isinstance(incidence_angle, xr.DataArray):
        raise TypeError("incidence_angle must be an xarray.DataArray.")
    if not np.issubdtype(incidence_angle.dtype, np.number):
        raise TypeError("incidence_angle must contain numeric values.")

    if not isinstance(freq_ghz, (int, float, np.number)):
        raise TypeError(f"freq_ghz must be a number, got {type(freq_ghz).__name__}.")
    if freq_ghz <= 0:
        raise ValueError(f"freg_ghz must be strictly positive, got {freq_ghz}.")

    if not isinstance(thresh1, (int, float, np.number)):
        raise TypeError(f"thresh1 must be a number, got {type(thresh1).__name__}.")
    if not isinstance(thresh2, (int, float, np.number)):
        raise TypeError(f"thresh2 must be a number, got {type(thresh2).__name__}.")

    if calibration_coeff is not None:
        if not isinstance(calibration_coeff, (int, float, np.number)):
            raise TypeError(
                "calibration_coeff must be a number or None, "
                f"got {type(calibration_coeff).__name__}."
            )
        if calibration_coeff <= 0:
            raise ValueError(
                f"calibration_coeff must be strictly positive, got {calibration_coeff}."
            )

    allowed_poltypes = ("S", "C3", "T3", "C4", "T4")
    poltype = validate_dataset(input_data, allowed_poltypes=allowed_poltypes)

    converters = {
        "C3": lambda ds: ds,
        "T3": T3_to_C3,
        "C4": C4_to_C3,
        "T4": T4_to_C3,
        "S": S_to_C3,
    }
    C3 = converters[poltype](input_data)

    out = _apply_dubois_inversion(
        theta=incidence_angle,
        f0=freq_ghz,
        hh=C3.m11,
        vv=C3.m33,
        hv=C3.m22 / 2.0,
        calib=calibration_coeff,
        thresh1=thresh1,
        thresh2=thresh2,
    )
    original_non_nan_mask = input_data.to_array().notnull().all("variable")
    out["dubois_mask_valid_in_out"] = (
        out["dubois_mask_in"]
        * out["dubois_mask_out"]
        * original_non_nan_mask.astype(np.float32, copy=False)
    ).astype(np.float32, copy=False)

    return xr.Dataset(
        {k: (tuple(input_data.dims), v.data) for k, v in out.items()},
        attrs={
            "poltype": "dubois_surface_inversion",
            "description": "Results of the Dubois surface inversion.",
        },
        coords=input_data.coords,
    )


# helper function, do not use directly
def _apply_dubois_inversion(theta, f0, hh, vv, hv, calib, thresh1, thresh2):

    scale = np.sin(theta) / calib if calib is not None else 1.0
    hh = hh * scale
    vv = vv * scale
    hv = hv * scale

    lambd = 100 * 0.3 / f0
    eps = np.finfo(np.float32).eps

    theta_valid = np.isfinite(theta) & (theta > 0) & (theta < (np.pi / 2.0))
    hh_pos = hh > 0
    vv_pos = vv > 0
    base_valid = theta_valid & hh_pos & vv_pos

    vv_safe = xr.where(vv_pos, vv, 1.0)
    hh_safe = xr.where(hh_pos, hh, 1.0)
    theta_safe = xr.where(theta_valid, theta, 1.0)

    msk_valid = (
        base_valid
        & (hv / vv_safe < 10 ** (thresh1 / 10.0))
        & (hh / vv_safe < 10 ** (thresh2 / 10.0))
    )

    ks_inv = np.exp(
        1.36905 * np.log(hh_safe)
        - 0.83333 * np.log(vv_safe)
        + 0.446425 * np.log(np.cos(theta_safe))
        + 3.34525 * np.log(np.sin(theta_safe))
        - 0.375 * np.log(lambd)
        + 1.78989 * np.log(10)
    )
    msk_ks = base_valid & (ks_inv >= 0) & (ks_inv <= np.pi)
    ks_dub = xr.where(msk_valid & msk_ks, ks_inv, 0.0)

    ks_safe = xr.where(ks_dub > 0, ks_dub, np.nan)
    er_inv = (
        np.log10(hh_safe)
        + np.log10(vv_safe)
        + 5.12
        + np.log10(np.cos(theta_safe))
        - 2.5 * np.log10(ks_safe + eps)
        - 1.4 * np.log10(lambd)
    ) / (0.074 * np.tan(theta_safe) + eps)
    msk_er = base_valid & (er_inv > 0) & (er_inv <= 100)
    er_dub = xr.where(msk_valid & msk_ks & msk_er, er_inv, 0.0)

    er_safe = xr.where(er_dub > 0, er_dub, np.nan)
    mv_inv = (
        -5.3e-2 + 2.92e-2 * er_safe - 5.5e-4 * er_safe**2 + 4.3e-6 * er_safe**3
    ) * 100
    msk_mv = base_valid & (mv_inv > 0) & (mv_inv <= 100)
    mv_dub = xr.where(msk_valid & msk_ks & msk_er & msk_mv, mv_inv, 0.0)

    msk_out = (msk_valid & msk_ks & msk_er & msk_mv).astype(np.float32)

    return {
        "dubois_ks": ks_dub.astype(np.float32, copy=False),
        "dubois_er": er_dub.astype(np.float32, copy=False),
        "dubois_mv": mv_dub.astype(np.float32, copy=False),
        "dubois_mask_out": msk_out.astype(np.float32, copy=False),
        "dubois_mask_in": msk_valid.astype(np.float32, copy=False),
    }
