"""Per-query prediction heads: lane/background classification + cubic curve regression.

Phase 1 note: curve coefficients are regressed directly (unconstrained) rather than
with LSTR's full numerical-reparameterization trick -- kept simple for the first
trainable pass; revisit if training shows y^3-term gradient instability.
"""
from __future__ import annotations

import torch
import torch.nn as nn


class MLP(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int = 2):
        super().__init__()
        dims = [in_dim] + [hidden_dim] * (num_layers - 1) + [out_dim]
        layers = []
        for i in range(len(dims) - 1):
            layers.append(nn.Linear(dims[i], dims[i + 1]))
            if i < len(dims) - 2:
                layers.append(nn.ReLU(inplace=True))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class PredictionHeads(nn.Module):
    """x = k*y^3 + m*y^2 + n*y + b, plus (y_start, y_end), plus lane/background logits."""

    def __init__(self, d_model: int, hidden_dim: int = 128):
        super().__init__()
        self.cls_head = MLP(d_model, hidden_dim, 2, num_layers=2)
        self.curve_head = MLP(d_model, hidden_dim, 4, num_layers=3)     # k, m, n, b
        self.endpoint_head = MLP(d_model, hidden_dim, 2, num_layers=2)  # y_start, y_end (pre-sigmoid)

    def forward(self, queries: torch.Tensor) -> dict:
        """queries: (B, N_lanes, d_model)"""
        cls_logits = self.cls_head(queries)              # (B, N, 2)
        curve_coeffs = self.curve_head(queries)           # (B, N, 4)
        endpoints = torch.sigmoid(self.endpoint_head(queries))  # (B, N, 2) in [0,1]
        return {"cls_logits": cls_logits, "curve_coeffs": curve_coeffs, "endpoints": endpoints}


def eval_curve(curve_coeffs: torch.Tensor, ys: torch.Tensor) -> torch.Tensor:
    """curve_coeffs: (..., 4) = [k,m,n,b]; ys: (...,S) or (S,) broadcastable normalized y.
    Returns x(y) with shape (..., S).
    """
    k, m, n, b = curve_coeffs.unbind(dim=-1)  # each (...,)
    k, m, n, b = k.unsqueeze(-1), m.unsqueeze(-1), n.unsqueeze(-1), b.unsqueeze(-1)
    y = ys
    return k * y**3 + m * y**2 + n * y + b