# Copyright 2026 Impala Systems, Inc.
# SPDX-License-Identifier: AGPL-3.0-only

from __future__ import annotations

"""Canonical Cronformer model.

Cronformer is an encoder-only structured decoder: it reads a natural-language
schedule request and emits logits for the five cron fields rather than rendering
cron text token by token. The decoder uses learned field queries and learned
slot queries so each output decision has an explicit place in the model.
"""

import math

import torch
from torch import nn
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, PreTrainedModel

from .configuration_cronformer import CronformerConfig
from .cron import (
    COMPONENT_DIMS,
    COMPONENT_NAMES,
    MAX_LIST_ITEMS,
    CronComponentOutput,
    CronOutput,
    CronPatternType,
)


SLOT_NAMES = [
    "pattern",
    "values",
    "list_values",
    "list_items",
    "range_start",
    "range_end",
    "step_start",
    "step_size",
    "nth",
    "last_offset",
    "field_presence",
]


class FieldSlotDecoder(nn.Module):
    """Prediction heads for one cron field after field/slot attention.

    The field-slot states carry global semantic context for this field, while
    token-level heads preserve copy-like access to numbers and pattern words
    that appear directly in the prompt.
    """

    def __init__(self, hidden_size: int, output_dim: int):
        super().__init__()
        self.output_dim = output_dim
        self.token_value_norm = nn.LayerNorm(hidden_size)
        self.token_step_size_norm = nn.LayerNorm(hidden_size)
        self.token_pattern_norm = nn.LayerNorm(hidden_size)
        self.pattern_head = nn.Linear(hidden_size, len(CronPatternType))
        self.values_head = nn.Linear(hidden_size, output_dim)
        self.list_values_head = nn.Linear(hidden_size, output_dim)
        self.range_start_head = nn.Linear(hidden_size, output_dim)
        self.range_end_head = nn.Linear(hidden_size, output_dim)
        self.step_start_head = nn.Linear(hidden_size, output_dim)
        self.step_size_head = nn.Linear(hidden_size, output_dim)
        self.nth_head = nn.Linear(hidden_size, 6)
        self.last_offset_head = nn.Linear(hidden_size, output_dim)
        self.value_count_head = nn.Linear(hidden_size, output_dim + 1)
        self.list_item_head = nn.Linear(hidden_size, output_dim)
        self.field_presence_head = nn.Linear(hidden_size, 1)
        self.token_values_head = nn.Linear(hidden_size, output_dim)
        self.token_step_size_head = nn.Linear(hidden_size, output_dim)
        self.token_pattern_head = nn.Linear(hidden_size, len(CronPatternType))
        self.value_fusion_gate = nn.Linear(hidden_size, output_dim)
        self.list_value_fusion_gate = nn.Linear(hidden_size, output_dim)
        self.step_size_fusion_gate = nn.Linear(hidden_size, output_dim)

    @staticmethod
    def _slot(states: torch.Tensor, name: str) -> torch.Tensor:
        return states[:, SLOT_NAMES.index(name), :]

    @staticmethod
    def _masked_token_pool(token_logits: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
        # Max pooling keeps a strong local token signal for numbers without
        # forcing the structured slot state to memorize every literal value.
        if attention_mask is None:
            return token_logits.amax(dim=1)
        mask = attention_mask.bool().unsqueeze(-1)
        return token_logits.masked_fill(~mask, -1.0e4).amax(dim=1)

    def forward(
        self,
        slot_states: torch.Tensor,
        token_states: torch.Tensor,
        attention_mask: torch.Tensor | None,
    ) -> CronComponentOutput:
        value_slot = self._slot(slot_states, "values")
        list_value_slot = self._slot(slot_states, "list_values")
        step_size_slot = self._slot(slot_states, "step_size")
        token_values_logits = self.token_values_head(
            self.token_value_norm(token_states + value_slot.unsqueeze(1))
        )
        token_step_size_logits = self.token_step_size_head(
            self.token_step_size_norm(token_states + step_size_slot.unsqueeze(1))
        )
        token_pattern_logits = self.token_pattern_head(
            self.token_pattern_norm(token_states + self._slot(slot_states, "pattern").unsqueeze(1))
        )
        pooled_token_values = self._masked_token_pool(token_values_logits, attention_mask)
        pooled_token_step_size = self._masked_token_pool(token_step_size_logits, attention_mask)

        # Fuse slot-level classification with token-local evidence. The gate is
        # learned per output value, which lets the model decide when prompt
        # literals should dominate over the semantic field representation.
        value_logits = self.values_head(value_slot) + torch.sigmoid(self.value_fusion_gate(value_slot)) * pooled_token_values
        list_values_logits = self.list_values_head(list_value_slot) + torch.sigmoid(
            self.list_value_fusion_gate(list_value_slot)
        ) * pooled_token_values
        step_size_logits = self.step_size_head(step_size_slot) + torch.sigmoid(
            self.step_size_fusion_gate(step_size_slot)
        ) * pooled_token_step_size

        list_item_state = self._slot(slot_states, "list_items")
        list_item_logits = self.list_item_head(list_item_state).unsqueeze(1).expand(
            -1,
            MAX_LIST_ITEMS,
            -1,
        )
        return CronComponentOutput(
            pattern_logits=self.pattern_head(self._slot(slot_states, "pattern")),
            values_logits=value_logits,
            range_start_logits=self.range_start_head(self._slot(slot_states, "range_start")),
            range_end_logits=self.range_end_head(self._slot(slot_states, "range_end")),
            step_start_logits=self.step_start_head(self._slot(slot_states, "step_start")),
            step_size_logits=step_size_logits,
            nth_logits=self.nth_head(self._slot(slot_states, "nth")),
            last_offset_logits=self.last_offset_head(self._slot(slot_states, "last_offset")),
            token_values_logits=token_values_logits,
            token_step_size_logits=token_step_size_logits,
            token_pattern_logits=token_pattern_logits,
            value_count_logits=self.value_count_head(value_slot),
            list_values_logits=list_values_logits,
            list_item_logits=list_item_logits,
            constrained_field_logits=self.field_presence_head(self._slot(slot_states, "field_presence")),
        )


class CronformerModel(PreTrainedModel):
    """Field-query Cronformer.

    Forward flow:
    1. Encode prompt tokens with a small Hugging Face encoder.
    2. Cross-attend five learned field queries over token states.
    3. Run field self-attention so minute/hour/dom/month/dow can coordinate.
    4. For each field, cross-attend learned slot queries over token states.
    5. Decode slot states into the CronOutput contract used by training/eval.
    """

    config_class = CronformerConfig

    def __init__(self, config: CronformerConfig):
        super().__init__(config)
        self.config = config
        self.encoder = AutoModel.from_config(config.to_encoder_config())
        encoder_dim = self.encoder.config.hidden_size
        hidden_size = config.hidden_size
        self.token_projection = nn.Sequential(
            nn.Linear(encoder_dim, hidden_size),
            nn.GELU(),
            nn.LayerNorm(hidden_size),
        )
        self.field_queries = nn.Parameter(torch.empty(len(COMPONENT_NAMES), hidden_size))
        self.slot_queries = nn.Parameter(torch.empty(len(SLOT_NAMES), hidden_size))
        self.field_attention = nn.MultiheadAttention(
            hidden_size,
            config.num_attention_heads,
            dropout=config.dropout,
            batch_first=True,
        )
        field_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=config.num_attention_heads,
            dim_feedforward=hidden_size * 4,
            dropout=config.dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.field_interaction = nn.TransformerEncoder(
            field_layer,
            num_layers=config.num_field_layers,
        )
        self.slot_attention = nn.MultiheadAttention(
            hidden_size,
            config.num_attention_heads,
            dropout=config.dropout,
            batch_first=True,
        )
        self.slot_norm = nn.LayerNorm(hidden_size)
        self.heads = nn.ModuleDict(
            {
                name: FieldSlotDecoder(hidden_size, output_dim)
                for name, output_dim in zip(COMPONENT_NAMES, COMPONENT_DIMS)
            }
        )
        self.dom_dow_pool = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.GELU(),
            nn.LayerNorm(hidden_size),
        )
        self.dom_dow_intersect_head = nn.Linear(hidden_size, 2)
        self.post_init()
        nn.init.normal_(self.field_queries, mean=0.0, std=0.02)
        nn.init.normal_(self.slot_queries, mean=0.0, std=0.02)

    def _init_weights(self, module: nn.Module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def _key_padding_mask(self, attention_mask: torch.Tensor | None) -> torch.Tensor | None:
        if attention_mask is None:
            return None
        return ~attention_mask.bool()

    def _token_states(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None,
    ) -> torch.Tensor:
        encoder_output = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        return self.token_projection(encoder_output.last_hidden_state)

    def _field_states(
        self,
        token_states: torch.Tensor,
        attention_mask: torch.Tensor | None,
    ) -> torch.Tensor:
        batch_size = token_states.shape[0]
        queries = self.field_queries.unsqueeze(0).expand(batch_size, -1, -1)
        # Each learned field query asks the token sequence for the evidence
        # relevant to one cron field. This is the Direct+V2 merge point: output
        # remains structured, but field semantics are learned by attention.
        field_states, _ = self.field_attention(
            queries,
            token_states,
            token_states,
            key_padding_mask=self._key_padding_mask(attention_mask),
            need_weights=False,
        )
        # Field interaction is the only learned global coordination layer. It
        # replaces hand-coded semantic corrections while still allowing fields
        # such as day-of-month and day-of-week to influence each other.
        return self.field_interaction(field_states)

    def _slot_states(
        self,
        field_states: torch.Tensor,
        token_states: torch.Tensor,
        attention_mask: torch.Tensor | None,
    ) -> dict[str, torch.Tensor]:
        slot_states = {}
        slot_offsets = self.slot_queries.unsqueeze(0)
        key_padding_mask = self._key_padding_mask(attention_mask)
        for index, name in enumerate(COMPONENT_NAMES):
            # Slot queries specialize a field state into concrete decisions:
            # pattern, values, ranges, steps, list items, and field presence.
            queries = field_states[:, index : index + 1, :] + slot_offsets
            attended, _ = self.slot_attention(
                queries,
                token_states,
                token_states,
                key_padding_mask=key_padding_mask,
                need_weights=False,
            )
            slot_states[name] = self.slot_norm(attended + queries)
        return slot_states

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor | None = None,
        token_type_ids: torch.Tensor | None = None,
        **_: object,
    ) -> CronOutput:
        token_states = self._token_states(input_ids, attention_mask)
        field_states = self._field_states(token_states, attention_mask)
        slot_states = self._slot_states(field_states, token_states, attention_mask)
        component_outputs = {}
        count_decode_components = set(getattr(self.config, "value_count_decoding_components", [0, 1]))
        token_decode_components = set(getattr(self.config, "token_value_decoding_components", [0, 1]))
        list_item_components = getattr(self.config, "list_item_components", None)
        list_item_component_set = None if list_item_components is None else set(list_item_components)
        for component_index, name in enumerate(COMPONENT_NAMES):
            output = self.heads[name](slot_states[name], token_states, attention_mask)
            # These flags are decoder policy switches consumed by cron.py. The
            # model always emits the full CronComponentOutput; configs choose
            # which learned logits participate in final cron reconstruction.
            output.use_value_count_decoding = (
                self.config.use_value_count_decoding and component_index in count_decode_components
            )
            output.use_token_value_decoding = (
                self.config.use_token_value_decoding and component_index in token_decode_components
            )
            output.use_path_score_decoding = self.config.use_path_score_decoding
            output.use_list_item_decoding = self.config.use_list_item_decoding and (
                list_item_component_set is None or component_index in list_item_component_set
            )
            output.use_list_item_scoring = self.config.use_list_item_scoring and (
                list_item_component_set is None or component_index in list_item_component_set
            )
            output.list_item_score_scale = self.config.list_item_score_scale
            component_outputs[name] = output

        pooled_fields = field_states.mean(dim=1)
        dom_dow_logits = self.dom_dow_intersect_head(self.dom_dow_pool(pooled_fields))
        return CronOutput(
            minute=component_outputs["minute"],
            hour=component_outputs["hour"],
            dom=component_outputs["dom"],
            month=component_outputs["month"],
            dow=component_outputs["dow"],
            dom_dow_intersect_logits=dom_dow_logits,
        )

    @classmethod
    def from_encoder(cls, encoder_model: str = "google/bert_uncased_L-2_H-128_A-2") -> "CronformerModel":
        config = CronformerConfig(encoder_model=encoder_model)
        model = cls(config)
        encoder_config = AutoConfig.from_pretrained(encoder_model)
        architectures = getattr(encoder_config, "architectures", None) or []
        # Some small language models expose their reusable token backbone behind
        # AutoModelForCausalLM rather than AutoModel. Keep this loader narrow:
        # Cronformer still uses hidden states as an encoder, never generation.
        if any(str(architecture).endswith("ForCausalLM") for architecture in architectures):
            causal_lm = AutoModelForCausalLM.from_pretrained(encoder_model)
            if not hasattr(causal_lm, "model"):
                raise ValueError(
                    f"{encoder_model} advertises a CausalLM architecture but does not expose a .model backbone"
                )
            model.encoder = causal_lm.model
        else:
            model.encoder = AutoModel.from_pretrained(encoder_model)
        return model