from dataclasses import dataclass

import torch
from torch import nn
from transformers import DistilBertConfig, DistilBertModel, PretrainedConfig, PreTrainedModel
from transformers.utils import ModelOutput


DEFAULT_DECISIONS = [
    {'name': 'task_complete', 'prompt': 'Has the task been successfully completed?', 'options': ['false', 'true']},
    {'name': 'should_retry', 'prompt': 'Should the same failed action be retried without changing it?', 'options': ['false', 'true']},
    {'name': 'failure_source', 'prompt': 'What is the source of the original reproduced failure?',
     'options': ['code', 'environment', 'test', 'dependency', 'unknown']},
    {'name': 'next_action', 'prompt': 'What action should the agent take next under the diagnostic policy?',
     'options': ['inspect_code', 'rerun_tests', 'inspect_environment', 'revert_change', 'escalate']},
    {'name': 'evidence_sufficient', 'prompt': 'Does exactly one controlled intervention identify the original failure source?',
     'options': ['false', 'true']},
]


class SystemOneConfig(PretrainedConfig):
    model_type = 'system_one'

    def __init__(self, encoder_config=None, max_length=256, head_depth=2,
                 representation_dropout=0.1, head_dropout=0.1,
                 decisions=None, temperatures=None, calibration=None, **kwargs):
        kwargs.setdefault('tie_word_embeddings', False)
        super().__init__(**kwargs)
        self.encoder_config = encoder_config or DistilBertConfig().to_dict()
        self.max_length = max_length
        self.head_depth = head_depth
        self.representation_dropout = representation_dropout
        self.head_dropout = head_dropout
        self.decisions = decisions or DEFAULT_DECISIONS
        self.num_labels_per_head = [len(item['options']) for item in self.decisions]
        self.temperatures = temperatures
        self.calibration = calibration
        if self.head_depth != 2:
            raise ValueError('This release class supports the published depth-2 checkpoint only')
        if len(self.decisions) != 5:
            raise ValueError('Expected exactly five fixed decisions')
        if temperatures is not None and (len(temperatures) != 5 or any(value <= 0 for value in temperatures)):
            raise ValueError('Expected five positive temperatures')


@dataclass
class SystemOneOutput(ModelOutput):
    logits: tuple[torch.Tensor, ...] | None = None
    hidden_states: tuple[torch.Tensor, ...] | None = None
    attentions: tuple[torch.Tensor, ...] | None = None


class SystemOneModel(PreTrainedModel):
    config_class = SystemOneConfig
    base_model_prefix = 'encoder'
    main_input_name = 'input_ids'
    _supports_sdpa = True

    def __init__(self, config):
        super().__init__(config)
        encoder_config = DistilBertConfig.from_dict(config.encoder_config)
        if config._attn_implementation is not None:
            encoder_config._attn_implementation = config._attn_implementation
        self.encoder = DistilBertModel(encoder_config)
        self.dropout = nn.Dropout(config.representation_dropout)
        self.heads = nn.ModuleList([
            nn.Sequential(
                nn.Linear(encoder_config.dim, encoder_config.dim),
                nn.GELU(),
                nn.Dropout(config.head_dropout),
                nn.Linear(encoder_config.dim, outputs),
            )
            for outputs in config.num_labels_per_head
        ])
        self.post_init()

    def forward(self, input_ids=None, attention_mask=None, head_mask=None, inputs_embeds=None,
                output_attentions=None, output_hidden_states=None, return_dict=None):
        return_dict = self.config.use_return_dict if return_dict is None else return_dict
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )
        representation = outputs.last_hidden_state[:, 0]
        logits = tuple(head(self.dropout(representation)) for head in self.heads)
        if not return_dict:
            return logits, outputs.hidden_states, outputs.attentions
        return SystemOneOutput(logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)