#!/usr/bin/env python3
"""model.py: PyTorch model definition for TokenLedgerForecaster with PyTorchModelHubMixin."""

import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin


class TokenLedgerForecaster(
    nn.Module,
    PyTorchModelHubMixin,
    library_name="tokenledger",
    repo_url="https://huggingface.co/ngc1024/tokenledger-forecast",
    pipeline_tag="tabular-regression",
):
    """TokenLedgerForecaster: Predicts output token quantiles (p10, p50, p90)

    Given context window signals (prompt size, previous turn output, elapsed idle gap,
    cache hit ratio, and window fill percentage), this lightweight model forecasts
    the expected output length and flags the tail blowout risk.
    """

    def __init__(self, n_features: int = 6, hidden_dim: int = 32, n_quantiles: int = 3):
        super().__init__()
        self.n_features = n_features
        self.hidden_dim = hidden_dim
        self.n_quantiles = n_quantiles

        self.net = nn.Sequential(
            nn.Linear(n_features, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, n_quantiles),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass. Input shape: (batch_size, n_features)
        Returns: (batch_size, 3) representing predicted [p10, p50, p90] in token counts.
        """
        # The network predicts in log1p space to guarantee positive predictions
        log_preds = self.net(x)
        # Clamped to exp(11) ~ 59,874 tokens
        return torch.expm1(torch.clamp(log_preds, min=0.0, max=11.0))