// bce_loss_multi_horizon.cu
//
// Fused multi-horizon BCE forward + backward with optional per-horizon
// loss weighting. Also emits per-horizon UNWEIGHTED mean BCE — that
// is the signal an ISV-driven horizon weighting layer EMA-tracks to
// produce dynamic per-horizon weights (high per-horizon BCE → horizon
// is under-trained → boost its weight). The total scalar loss_out
// remains the externally-weighted aggregate used by callers for
// reporting and early-stopping.
//
// Input layout (row-major):
//   probs   [n_pos, n_horizons]  — model outputs in (0, 1)
//   labels  [n_pos, n_horizons]  — binary {0.0, 1.0}; NaN = mask (drop)
//   loss_weights [n_horizons]    — per-horizon weight (nullptr = all 1.0)
// Output:
//   loss_out[1]                       — weighted mean BCE over valid entries
//   loss_per_horizon[n_horizons]      — UNWEIGHTED mean BCE per horizon (ISV signal)
//   grad_probs[n_pos, n_horizons]     — d loss / d prob (per-element scaled)
//   valid_count_out[1]                — number of non-NaN labels (unweighted)
//
//   per element (mask m = !isnan(y), weight w = loss_weights[h]):
//     p = clamp(probs[i], 1e-6, 1 - 1e-6)
//     L_i = -[y log p + (1-y) log(1-p)]   (unweighted; weighted form scales L_i by w)
//     dL/dp = m * w * (p - y) / (p (1-p))
//
//   Normaliser: W_valid = sum_i m_i * w_i
//   total_loss = (1/W_valid) * sum_i (m_i * w_i * L_i)
//   grad_probs[i] = m_i * w * (p - y) / (p (1-p)) / W_valid
//   loss_per_horizon[h] = (sum_{i: m_i, h_i=h} L_i) / (sum_{i: m_i, h_i=h} 1)
//
// Caller contract: n_horizons MUST equal N_HORIZONS_BCE (5). Hardcoded
// rather than dynamic-shared-mem because the per-horizon accumulator
// matrix has a fixed shape (5 × 256) at our trainer's scale.

#define N_HORIZONS_BCE 5

extern "C" __global__ void bce_multi_horizon_forward_backward(
    const float* __restrict__ probs,         // [n_pos * n_horizons]
    const float* __restrict__ labels,        // [n_pos * n_horizons]
    const float* __restrict__ loss_weights,  // [n_horizons] (nullptr OK → 1.0)
    int n_pos,
    int n_horizons,
    float* __restrict__ loss_out,            // [1] — externally-weighted mean
    float* __restrict__ loss_per_horizon,    // [n_horizons] — UNWEIGHTED per-horizon mean
    float* __restrict__ grad_probs,          // [n_pos * n_horizons]
    int* __restrict__ valid_count_out        // [1] — number of non-NaN labels
) {
    const int tid = threadIdx.x;
    const int total = n_pos * n_horizons;

    __shared__ float sloss[256];     // weighted-loss numerator accumulator
    __shared__ float sw_valid[256];  // weighted-mask normaliser accumulator
    __shared__ int   svalid[256];    // unweighted valid-label counter (reported)
    // Per-horizon accumulators (unweighted): sloss_h[h*256 + tid],
    // svalid_h[h*256 + tid]. Hardcoded width N_HORIZONS_BCE=5 → 13 KiB
    // total shared usage including the three scalars above; comfortably
    // under any SM smem limit (48-100 KiB depending on arch).
    __shared__ float sloss_h[N_HORIZONS_BCE * 256];
    __shared__ float svalid_h[N_HORIZONS_BCE * 256];
    sloss[tid] = 0.0f;
    sw_valid[tid] = 0.0f;
    svalid[tid] = 0;
    #pragma unroll
    for (int h = 0; h < N_HORIZONS_BCE; ++h) {
        sloss_h[h * 256 + tid] = 0.0f;
        svalid_h[h * 256 + tid] = 0.0f;
    }

    // First pass: count valid entries and accumulate weighted aggregate
    // loss + unweighted per-horizon loss + weight sums.
    for (int i = tid; i < total; i += blockDim.x) {
        const float y = labels[i];
        if (isnan(y)) {
            grad_probs[i] = 0.0f;
            continue;
        }
        const int h = i % n_horizons;
        const float w = (loss_weights != nullptr) ? loss_weights[h] : 1.0f;
        const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
        const float L = -(y * logf(p) + (1.0f - y) * logf(1.0f - p));
        sloss[tid] += w * L;
        sloss_h[h * 256 + tid] += L;       // UNWEIGHTED — the ISV signal
        svalid_h[h * 256 + tid] += 1.0f;
        sw_valid[tid] += w;
        svalid[tid] += 1;
    }
    __syncthreads();

    // Block tree-reduce on all accumulators (incl. per-horizon arrays).
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sloss[tid] += sloss[tid + s];
            sw_valid[tid] += sw_valid[tid + s];
            svalid[tid] += svalid[tid + s];
            #pragma unroll
            for (int h = 0; h < N_HORIZONS_BCE; ++h) {
                sloss_h[h * 256 + tid] += sloss_h[h * 256 + tid + s];
                svalid_h[h * 256 + tid] += svalid_h[h * 256 + tid + s];
            }
        }
        __syncthreads();
    }

    if (tid == 0) {
        const float w_valid = sw_valid[0];
        loss_out[0] = (w_valid > 0.0f) ? sloss[0] / w_valid : 0.0f;
        valid_count_out[0] = svalid[0];
        #pragma unroll
        for (int h = 0; h < N_HORIZONS_BCE; ++h) {
            const float c_h = svalid_h[h * 256];
            loss_per_horizon[h] = (c_h > 0.0f) ? sloss_h[h * 256] / c_h : 0.0f;
        }
    }
    __syncthreads();

    // Second pass: emit per-element grad scaled by w_h / W_valid.
    const float w_valid = sw_valid[0];
    const float inv = (w_valid > 0.0f) ? (1.0f / w_valid) : 0.0f;
    for (int i = tid; i < total; i += blockDim.x) {
        const float y = labels[i];
        if (isnan(y)) {
            continue;
        }
        const int h = i % n_horizons;
        const float w = (loss_weights != nullptr) ? loss_weights[h] : 1.0f;
        const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
        grad_probs[i] = w * (p - y) / (p * (1.0f - p)) * inv;
    }
}
