// aux_loss.cu — class-weighted BCE + NaN-masked Huber for the A+B
// paired prof-binary + size-regression aux supervision (SDD-3 CB1/CB3/CB4).
//
// Per CB1's relabel: each (direction, horizon) emits TWO labels:
//   y_prof ∈ {0, 1} or NaN  — did the trade hit profit before stop?
//   y_size ∈ ℝ    or NaN  — σ-normalised realised return (only when prof=1)
//
// Two loss kernels here:
//
//   aux_bce_loss_fwd_bwd     — sigmoid+BCE on (prof_logit, y_prof), with
//                              per-horizon `pos_weight` for class balance
//                              (`pos_weight = clamp(n_neg/n_pos, 1.0, 50.0)`
//                              computed once per epoch on the loader side).
//                              NaN in y_prof zeroes the gradient and
//                              skips the loss contribution.
//
//   aux_huber_masked_fwd_bwd — Huber on (y_size_hat, y_size_true). The
//                              loader emits y_size = NaN whenever
//                              y_prof = 0 (or the sample is invalid),
//                              which gives the conditional-Huber
//                              semantics for free: the NaN mask both
//                              zeroes the gradient and skips the loss
//                              contribution.
//
// Formulation:
//
//   BCE (sigmoid fused, class-weighted):
//     p          = 1 / (1 + exp(-z))
//     L_BCE      = -[pos_weight * y * log(p) + (1-y) * log(1-p)]
//     dL_BCE/dz  = (y > 0.5) ? pos_weight * (p - 1)
//                            : p
//     (clamp log(·) at log(1e-7) to guard against p ∈ {0, 1} blowup;
//     `--use_fast_math` already produces approximate sigmoid so the
//     clamp is a defence-in-depth.)
//
//   Huber (δ = 1.0 default):
//     r          = y_hat - y_true
//     L_Huber    = (|r| <= δ) ? 0.5 * r²        : δ * (|r| - 0.5 * δ)
//     dL/dy_hat  = (|r| <= δ) ? r                : δ * sign(r)
//
// NaN-mask discipline: y_true may be NaN at masked positions; we MUST
// zero the gradient AND skip the loss accumulator at those slots so
// downstream `aux_heads_bwd` propagates zero through the head and the
// `aux_trunk_bwd` sees no spurious gradient. The optimiser cannot eat
// a NaN scalar — a single masked label without the mask would
// poison the entire training step.
//
// Reduction discipline (per `feedback_no_atomicadd.md`):
//   * Per-thread strided accumulation into registers.
//   * Warp-shuffle reduce (`__shfl_xor_sync`) to one float per warp.
//   * One `__syncthreads` to gather warp partials in shared memory.
//   * Final warp-0 reduction to total loss / total valid count.
//   * No atomicAdd anywhere.
//
// Per-direction split: each kernel processes ONE direction × ONE head's
// tensor at a time. The trainer calls each twice (once for long, once
// for short) and sums the four scalars into the total aux loss. Keeping
// the kernels per-(direction, head) keeps the block layout simple and
// lets the per-direction telemetry (mean loss, valid count) emerge
// for free.
//
// Block layout: grid = (1), block = (BLOCK = 256). Single block handles
// the full [B × N_AUX_HORIZONS] grid (max ~32 × 3 = 96, easily within
// stride reach). Mirrors the pattern in `bce_loss_multi_horizon.cu`.

#define AUX_LOSS_BLOCK   256
#define AUX_LOSS_WARPS   (AUX_LOSS_BLOCK / 32)   // == 8
#define N_AUX_HORIZONS   3


__device__ __forceinline__ float warp_reduce_sum_f32(float v) {
    #pragma unroll
    for (int s = 16; s > 0; s >>= 1) {
        v += __shfl_xor_sync(0xffffffff, v, s);
    }
    return v;
}

__device__ __forceinline__ int warp_reduce_sum_i32(int v) {
    #pragma unroll
    for (int s = 16; s > 0; s >>= 1) {
        v += __shfl_xor_sync(0xffffffff, v, s);
    }
    return v;
}


// ─────────────────────────────────────────────────────────────────────
// aux_bce_loss_fwd_bwd:
//   Fused sigmoid + class-weighted BCE forward + per-element logit
//   gradient writeback for ONE direction (long or short) × the prof
//   head. Inputs/outputs are [B × N_AUX_HORIZONS] layout. NaN-masked
//   positions contribute zero loss and write zero gradient (no
//   contamination of downstream consumers).
//
// Inputs:
//   prof_logit       [B × N_AUX_HORIZONS]   — RAW logit z (no sigmoid)
//   y_prof_true      [B × N_AUX_HORIZONS]   — labels in {0, 1} or NaN
//   pos_weight       [N_AUX_HORIZONS]       — per-horizon positive-class
//                                              weight (uploaded host-side
//                                              from loader epoch stats).
//   n_total          = B × N_AUX_HORIZONS
//
// Outputs:
//   loss_out         [1]                    — Σ_valid L_BCE (UNREDUCED;
//                                              caller divides by valid_count
//                                              if a mean is desired).
//   grad_prof_logit  [B × N_AUX_HORIZONS]   — dL_BCE/dz per element.
//   valid_count_out  [1]                    — #{(b, k) : isfinite(y_prof_true)}.
//
// `loss_out` is intentionally a SUM, not a mean — caller controls the
// reduction policy without re-running the kernel.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void aux_bce_loss_fwd_bwd(
    const float* __restrict__ prof_logit,       // [B × N_AUX_HORIZONS]
    const float* __restrict__ y_prof_true,      // [B × N_AUX_HORIZONS]   (NaN-maskable)
    const float* __restrict__ pos_weight,       // [N_AUX_HORIZONS]
    int          n_total,                       // = B × N_AUX_HORIZONS
    float* __restrict__ loss_out,               // [1] — Σ L_BCE (unreduced)
    float* __restrict__ grad_prof_logit,        // [B × N_AUX_HORIZONS]
    int*   __restrict__ valid_count_out         // [1]
) {
    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;

    // Stage `pos_weight[N_AUX_HORIZONS]` into shared so every strided
    // iteration reads from smem rather than re-fetching from DRAM.
    __shared__ float s_pw[N_AUX_HORIZONS];
    if (tid < N_AUX_HORIZONS) {
        s_pw[tid] = pos_weight[tid];
    }
    __syncthreads();

    float local_loss  = 0.0f;
    int   local_valid = 0;

    for (int i = tid; i < n_total; i += AUX_LOSS_BLOCK) {
        const float y_t = y_prof_true[i];
        if (!isfinite(y_t)) {
            grad_prof_logit[i] = 0.0f;
            continue;
        }
        const int h = i % N_AUX_HORIZONS;
        const float pw  = s_pw[h];
        const float z   = prof_logit[i];
        // sigmoid(z) = 1/(1+exp(-z)). With --use_fast_math nvcc emits
        // the fast intrinsic; we clamp the log argument to avoid -inf
        // when p saturates to 0 or 1.
        const float p   = 1.0f / (1.0f + expf(-z));
        const float p_c     = fmaxf(p,        1e-7f);
        const float one_m_p = fmaxf(1.0f - p, 1e-7f);

        // Class-weighted BCE: positive class gets weight `pw`, negative
        // class gets weight 1. Equivalent to scaling the positive-class
        // term in the standard formula.
        float L;
        float grad_z;
        if (y_t > 0.5f) {
            L      = -pw * logf(p_c);
            grad_z = pw * (p - 1.0f);
        } else {
            L      = -logf(one_m_p);
            grad_z = p;
        }
        local_loss          += L;
        local_valid         += 1;
        grad_prof_logit[i]   = grad_z;
    }

    // Warp-shuffle reduce.
    float warp_loss  = warp_reduce_sum_f32(local_loss);
    int   warp_valid = warp_reduce_sum_i32(local_valid);

    __shared__ float s_loss[AUX_LOSS_WARPS];
    __shared__ int   s_valid[AUX_LOSS_WARPS];
    if (lane == 0) {
        s_loss[warp]  = warp_loss;
        s_valid[warp] = warp_valid;
    }
    __syncthreads();

    if (warp == 0) {
        float v_l = (lane < AUX_LOSS_WARPS) ? s_loss[lane]  : 0.0f;
        int   v_v = (lane < AUX_LOSS_WARPS) ? s_valid[lane] : 0;
        v_l = warp_reduce_sum_f32(v_l);
        v_v = warp_reduce_sum_i32(v_v);
        if (lane == 0) {
            loss_out[0]        = v_l;
            valid_count_out[0] = v_v;
        }
    }
}


// ─────────────────────────────────────────────────────────────────────
// aux_huber_masked_fwd_bwd:
//   Fused Huber-loss forward + per-element gradient writeback for ONE
//   direction (long or short) × the size head. NaN-masked positions
//   contribute zero loss and write zero gradient.
//
// Conditional-Huber semantics: the loader (CB1) sets `y_size = NaN`
// whenever `y_prof = 0` (no profit hit so the realised-return target is
// undefined). The NaN-mask path here therefore implements the
// conditional-Huber loss `L_Huber if y_prof=1 else 0` without needing a
// separate prof-mask buffer — the loader has already encoded the
// condition into the label tensor.
//
// Inputs:
//   y_size_hat       [B × N_AUX_HORIZONS]   — raw linear regression output
//   y_size_true      [B × N_AUX_HORIZONS]   — σ-normalised return or NaN
//   delta            — Huber transition point (default 1.0)
//   n_total          = B × N_AUX_HORIZONS
//
// Outputs:
//   loss_out         [1]                    — Σ_valid L_Huber (unreduced)
//   grad_y_size_hat  [B × N_AUX_HORIZONS]   — dL/dy_size_hat per element
//   valid_count_out  [1]                    — #{(b, k) : isfinite(y_size_true)}.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void aux_huber_masked_fwd_bwd(
    const float* __restrict__ y_size_hat,       // [B × N_AUX_HORIZONS]
    const float* __restrict__ y_size_true,      // [B × N_AUX_HORIZONS]   (NaN-maskable)
    float        delta,                          // Huber transition point
    int          n_total,                        // = B × N_AUX_HORIZONS
    float* __restrict__ loss_out,                // [1] — Σ L_Huber (unreduced)
    float* __restrict__ grad_y_size_hat,         // [B × N_AUX_HORIZONS]
    int*   __restrict__ valid_count_out          // [1]
) {
    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;

    float local_loss  = 0.0f;
    int   local_valid = 0;

    for (int i = tid; i < n_total; i += AUX_LOSS_BLOCK) {
        const float y_t = y_size_true[i];
        if (!isfinite(y_t)) {
            grad_y_size_hat[i] = 0.0f;
            continue;
        }
        const float y_p = y_size_hat[i];
        const float r   = y_p - y_t;
        const float ar  = fabsf(r);

        float L, g;
        if (ar <= delta) {
            L = 0.5f * r * r;
            g = r;
        } else {
            L = delta * (ar - 0.5f * delta);
            g = copysignf(delta, r);
        }
        local_loss          += L;
        local_valid         += 1;
        grad_y_size_hat[i]   = g;
    }

    // Warp-shuffle reduce.
    float warp_loss  = warp_reduce_sum_f32(local_loss);
    int   warp_valid = warp_reduce_sum_i32(local_valid);

    __shared__ float s_loss[AUX_LOSS_WARPS];
    __shared__ int   s_valid[AUX_LOSS_WARPS];
    if (lane == 0) {
        s_loss[warp]  = warp_loss;
        s_valid[warp] = warp_valid;
    }
    __syncthreads();

    if (warp == 0) {
        float v_l = (lane < AUX_LOSS_WARPS) ? s_loss[lane]  : 0.0f;
        int   v_v = (lane < AUX_LOSS_WARPS) ? s_valid[lane] : 0;
        v_l = warp_reduce_sum_f32(v_l);
        v_v = warp_reduce_sum_i32(v_v);
        if (lane == 0) {
            loss_out[0]        = v_l;
            valid_count_out[0] = v_v;
        }
    }
}
