# distutils: language = c++
# cython: boundscheck = False
# cython: wraparound = False
# cython: cdivision = True

import numpy as np
cimport numpy as cnp
from libc.math cimport exp, cos, tanh, log10, pow,atan2,sqrt
from cython.parallel import prange

cdef double windspeedcmd5n(double angle, double incidenceangle, double sigma0) nogil:
    """
    CMOD5N模型计算风速 - Cython优化版本
    
    输入参数：
      angle - 风向角度(度)
      incidenceangle - 入射角(度)  
      sigma0 - 后向散射系数(dB)
    输出参数：
      U - 计算得到的风速(m/s)
    """
    # 声明C类型变量
    # cdef double U=0
    cdef int i=0
    cdef double y0, n, a, b, x, a0, a1, a2, gamma, s0, s, g, alpha, f
    cdef double b0, b1, b2, angle_rad, sigma0_model, v0, d1, d2, y, v2
    
    # CMOD5N模型系数 - 使用C数组存储以提高性能
    cdef double[29] cmod5_c = [
        0, -0.6878, -0.7957, 0.3380, -0.1728, 0.0000, 0.0040, 0.1103, 0.0159,
        6.7329, 2.7713, -2.2885, 0.4971, -0.7250, 0.0450,
        0.0066, 0.3222, 0.0120, 22.7000, 2.0813, 3.0000, 8.3659,
        -3.3428, 1.3236, 6.2437, 2.3893, 0.3249, 4.1590, 1.6930
    ]
    
    # 确保入射角不小于16度
    if incidenceangle <= 16.0:
        incidenceangle = 16.0

    i=0
    cdef int start = 0
    cdef int end = 2601
    # 在0-26m/s范围内搜索风速
    for i in range(start, end):  # 0到26.0，步长0.01，共2601次迭代
        U_actual = i * 0.01
        
        # 计算归一化参数
        y0 = cmod5_c[19]  # 索引从0开始
        n = cmod5_c[20]
        a = y0 - (y0 - 1.0) / n
        b = 1.0 / (n * pow(y0 - 1.0, n - 1.0))
        
        # 归一化入射角
        x = (incidenceangle - 40.0) / 25.0
        
        # 计算基础后向散射系数
        a0 = (cmod5_c[1] + cmod5_c[2] * x + cmod5_c[3] * x * x + 
              cmod5_c[4] * x * x * x)
        a1 = cmod5_c[5] + cmod5_c[6] * x
        a2 = cmod5_c[7] + cmod5_c[8] * x
        gamma = (cmod5_c[9] + cmod5_c[10] * x + 
                cmod5_c[11] * x * x)
        s0 = cmod5_c[12] + cmod5_c[13] * x
        s = a2 * U_actual
        
        # 计算饱和函数
        if s < s0:
            g = 1.0 / (1.0 + exp(-s0))
            alpha = s0 * (1.0 - g)
            f = pow(s / s0, alpha) * g
        else:
            f = 1.0 / (1.0 + exp(-s))
        
        # 计算各向同性部分
        b0 = pow(f, gamma) * pow(10.0, a0 + a1 * U_actual)
        
        # 计算一阶调和系数
        b1 = (cmod5_c[15] * U_actual *
              (0.5 + x - tanh(4.0 * (x + cmod5_c[16] + cmod5_c[17] * U_actual))))
        b1 = ((cmod5_c[14] * (1.0 + x) - b1) /
              (exp(0.34 * (U_actual - cmod5_c[18])) + 1.0))
        
        # 计算二阶调和系数
        v0 = cmod5_c[21] + cmod5_c[22] * x + cmod5_c[23] * x * x
        d1 = cmod5_c[24] + cmod5_c[25] * x + cmod5_c[26] * x * x
        d2 = cmod5_c[27] + cmod5_c[28] * x
        y = (U_actual / v0 + 1.0)
        
        if y < y0:
            v2 = a + b * pow(y - 1.0, n)
        else:
            v2 = y

        b2 = (-d1 + d2 * v2) * exp(-v2)
        
        # 计算模拟后向散射系数并与实测值比较
        angle_rad = angle * 3.141592653589793 / 180.0
        sigma0_model = (10.0 * log10(b0 * 
                        pow(1.0 + b1 * cos(angle_rad) + b2 * cos(2.0 * angle_rad), 1.6)))
        
        if (sigma0_model - sigma0) >= 0.0:  # 函数为单调的
            return U_actual
    
    # 如果未找到匹配的风速，返回最大值
    return 26.0



cdef void windspeedCmod5n_parallel(float[:, :] wind_dir,
                                  float[:, :] inc_angle,
                                  float[:, :] sigma0,
                                  double[:, :] out_wind_speed) nogil:
    cdef int rowcount = wind_dir.shape[0]
    cdef int colcount = wind_dir.shape[1]
    cdef int rowid, colid

    for rowid in prange(rowcount, nogil=True, schedule='static'):  # 使用prange
        for colid in range(colcount):
            # 通过内存视图访问数据是纯C操作，无需GIL
            out_wind_speed[rowid, colid] = windspeedcmd5n(
                wind_dir[rowid, colid],
                inc_angle[rowid, colid],
                sigma0[rowid, colid]
            )

cdef void ERA5_windAngle_parallel(float[:, :] wind_U10,
                                  float[:, :] wind_V10,
                                  float[:, :] out_wind_angle) nogil:
    cdef int rowcount = wind_U10.shape[0]
    cdef int colcount = wind_U10.shape[1]
    cdef int rowid, colid
    cdef double V=0;
    cdef double sinValue=0;
    cdef double cosValue=0;
    cdef double U10=0
    cdef double V10=0
    for rowid in prange(rowcount, nogil=True, schedule='static'):
        for colid in range(colcount):
            U10=wind_U10[rowid,colid]
            V10=wind_V10[rowid,colid]
            V=sqrt(U10*U10+V10*V10)
            sinValue=-U10/V
            cosValue=-V10/V
            out_wind_angle[rowid, colid] =<float>atan2(sinValue,cosValue)





def python_ERA5_windAngle(cnp.ndarray[float,ndim=2] wind_U10, cnp.ndarray[float,ndim=2] wind_V10):
    cdef float[:, :] out_wind_angle = np.empty((wind_U10.shape[0], wind_U10.shape[1]), dtype=np.float32)
    cdef float[:, :] wind_U10_view = wind_U10
    cdef float[:, :] wind_V10_view = wind_V10
    with nogil:
        ERA5_windAngle_parallel(wind_U10_view, wind_V10_view,out_wind_angle)
    return np.asarray(out_wind_angle)


# 4. 提供一个Python包装函数（如果需要从Python调用）
def python_windspeedCmod5n(cnp.ndarray[float,ndim=2] wind_dir, cnp.ndarray[float,ndim=2] inc_angle, cnp.ndarray[float,ndim=2] sigma0):
    # 1. 创建输出数组
    cdef double[:, :] out_speed = np.empty((wind_dir.shape[0], wind_dir.shape[1]), dtype=np.float64)

    # 2. 将输入的numpy数组转换为内存视图
    cdef float[:, :] wind_dir_view = wind_dir
    cdef float[:, :] inc_angle_view = inc_angle
    cdef float[:, :] sigma0_view = sigma0

    # 3. 在nogil上下文中执行核心计算
    with nogil:
        windspeedCmod5n_parallel(wind_dir_view, inc_angle_view, sigma0_view, out_speed)

    # 4. 将结果内存视图转换回numpy数组并返回
    return np.asarray(out_speed)