import numpy as np
from scipy.special import exp1
from modules.calctav import calctav  # 你自己已有的函数

def prospect_d(N, Cab, Car, Ant, Brown, Cw, Cm, wl, nr, Kab, Kcar, Kant, Kbrown, Kw, Km):
    """
    Accurate PROSPECT-D implementation matching the original MATLAB model.

    Parameters:
        N       : Structure parameter
        Cab     : Chlorophyll content (µg/cm²)
        Car     : Carotenoid content (µg/cm²)
        Ant     : Anthocyanins content (µg/cm²)
        Brown   : Brown pigment content (a.u.)
        Cw      : Equivalent water thickness (cm)
        Cm      : Dry matter content (g/cm²)
        wl      : Wavelength array (nm)
        nr      : Real refractive index
        Kab,... : Absorption coefficients for components (same length as wl)

    Returns:
        LRT : ndarray of shape (n, 3) → [wavelength, reflectance, transmittance]
    """

    # 1. Compute total absorption coefficient
    Kall = (Cab * Kab + Car * Kcar + Ant * Kant + Brown * Kbrown + Cw * Kw + Cm * Km) / N

    # 2. Compute tau (internal transmittance) with expint for K > 0
    t1 = (1 - Kall) * np.exp(-Kall)
    t2 = (Kall ** 2) * exp1(Kall)
    tau = np.ones_like(Kall)
    j = Kall > 0
    tau[j] = t1[j] + t2[j]

    # 3. Fresnel interface properties
    talf = calctav(40, nr)
    ralf = 1 - talf
    t12 = calctav(90, nr)
    r12 = 1 - t12
    t21 = t12 / (nr ** 2)
    r21 = 1 - t21

    # 4. Top surface
    denom = 1 - (r21 ** 2) * (tau ** 2)
    Ta = talf * tau * t21 / denom
    Ra = ralf + r21 * tau * Ta

    # 5. Bottom surface
    t = t12 * tau * t21 / denom
    r = r12 + r21 * tau * t

    # 6. Stokes multilayer reflectance and transmittance
    D = np.sqrt((1 + r + t) * (1 + r - t) * (1 - r + t) * (1 - r - t))
    rq = r ** 2
    tq = t ** 2
    a = (1 + rq - tq + D) / (2 * r)
    b = (1 - rq + tq + D) / (2 * t)

    bNm1 = b ** (N - 1)
    bN2 = bNm1 ** 2
    a2 = a ** 2
    denom2 = a2 * bN2 - 1
    Rsub = a * (bN2 - 1) / denom2
    Tsub = bNm1 * (a2 - 1) / denom2

    # 7. Zero absorption edge-case correction
    j0 = (r + t >= 1)
    Tsub[j0] = t[j0] / (t[j0] + (1 - t[j0]) * (N - 1))
    Rsub[j0] = 1 - Tsub[j0]

    # 8. Final reflectance and transmittance
    denom3 = 1 - Rsub * r
    tran = Ta * Tsub / denom3
    refl = Ra + Ta * Rsub * t / denom3

    # 9. Clip values to [0, 1]
    refl = np.clip(refl, 0, 1)
    tran = np.clip(tran, 0, 1)

    return np.column_stack((wl, refl, tran))