import torch
import torch.nn as nn
from typing import List, Optional
from transformers import (
    AutoConfig,
    AutoModel,
    PreTrainedModel,
    PretrainedConfig,
)


class RQAModelConfig(PretrainedConfig):
    model_type = "rqa_v2_2"

    def __init__(
        self,
        base_model_name: str = "FacebookAI/xlm-roberta-large",
        encoder_config: Optional[dict] = None,
        error_types: Optional[List[str]] = None,
        schema_version: str = "rqa.v2.2",
        has_issue_projection_dim: int = 256,
        hidden_projection_dim: int = 256,
        errors_projection_dim: int = 512,
        has_issue_dropout: float = 0.25,
        hidden_dropout: float = 0.25,
        errors_dropout: float = 0.3,
        temperature_has_issue: float = 1.0,
        temperature_is_hidden: float = 1.0,
        temperature_errors: Optional[List[float]] = None,
        threshold_has_issue: float = 0.5,
        threshold_is_hidden: float = 0.5,
        threshold_error: float = 0.5,
        threshold_errors: Optional[List[float]] = None,
        **kwargs
    ):
        super().__init__(**kwargs)

        self.base_model_name = base_model_name
        self.encoder_config = encoder_config
        self.error_types = error_types or [
            "false_causality",
            "unsupported_claim",
            "overgeneralization",
            "missing_premise",
            "contradiction",
            "circular_reasoning",
        ]
        self.num_error_types = len(self.error_types)
        self.schema_version = schema_version

        self.has_issue_projection_dim = has_issue_projection_dim
        self.hidden_projection_dim = hidden_projection_dim
        self.errors_projection_dim = errors_projection_dim

        self.has_issue_dropout = has_issue_dropout
        self.hidden_dropout = hidden_dropout
        self.errors_dropout = errors_dropout

        self.temperature_has_issue = float(temperature_has_issue)
        self.temperature_is_hidden = float(temperature_is_hidden)
        self.temperature_errors = (
            temperature_errors
            if temperature_errors is not None
            else [1.0] * self.num_error_types
        )

        self.threshold_has_issue = float(threshold_has_issue)
        self.threshold_is_hidden = float(threshold_is_hidden)
        self.threshold_error = float(threshold_error)
        self.threshold_errors = (
            threshold_errors
            if threshold_errors is not None
            else [float(threshold_error)] * self.num_error_types
        )

        try:
            self._experts_implementation = "eager"
            self._experts_implementation_internal = "eager"
        except Exception:
            pass


class MeanPooling(nn.Module):
    def forward(self, last_hidden_state, attention_mask):
        mask = attention_mask.unsqueeze(-1).float()
        summed = torch.sum(last_hidden_state * mask, dim=1)
        denom = torch.clamp(mask.sum(dim=1), min=1e-9)
        return summed / denom


class RQAModelHF(PreTrainedModel):
    config_class = RQAModelConfig
    _supports_grouped_mm = False

    def __init__(self, config: RQAModelConfig):
        super().__init__(config)

        try:
            config._experts_implementation = "eager"
            config._experts_implementation_internal = "eager"
        except Exception:
            pass

        self.encoder = AutoModel.from_pretrained(config.base_model_name)
        hidden_size = self.encoder.config.hidden_size

        self.pooler = MeanPooling()

        self.has_issue_projection = nn.Sequential(
            nn.Linear(hidden_size, config.has_issue_projection_dim),
            nn.LayerNorm(config.has_issue_projection_dim),
            nn.GELU(),
            nn.Dropout(config.has_issue_dropout),
        )

        self.hidden_projection = nn.Sequential(
            nn.Linear(hidden_size, config.hidden_projection_dim),
            nn.LayerNorm(config.hidden_projection_dim),
            nn.GELU(),
            nn.Dropout(config.hidden_dropout),
        )

        self.errors_projection = nn.Sequential(
            nn.Linear(hidden_size, config.errors_projection_dim),
            nn.LayerNorm(config.errors_projection_dim),
            nn.GELU(),
            nn.Dropout(config.errors_dropout),
        )

        self.has_issue_head = nn.Linear(config.has_issue_projection_dim, 1)
        self.is_hidden_head = nn.Linear(config.hidden_projection_dim, 1)
        self.errors_head = nn.Linear(
            config.errors_projection_dim,
            config.num_error_types,
        )

        self.log_var_has_issue = nn.Parameter(torch.zeros(1))
        self.log_var_is_hidden = nn.Parameter(torch.zeros(1))
        self.log_var_errors = nn.Parameter(torch.zeros(1))

        self._init_custom_weights()

    def _init_custom_weights(self):
        for module in [
            self.has_issue_projection[0],
            self.hidden_projection[0],
            self.errors_projection[0],
            self.has_issue_head,
            self.is_hidden_head,
            self.errors_head,
        ]:
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

    def forward(self, input_ids=None, attention_mask=None, **kwargs):
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True,
        )

        pooled = self.pooler(outputs.last_hidden_state, attention_mask)

        has_issue_logits = self.has_issue_head(
            self.has_issue_projection(pooled)
        ).squeeze(-1)

        is_hidden_logits = self.is_hidden_head(
            self.hidden_projection(pooled)
        ).squeeze(-1)

        errors_logits = self.errors_head(
            self.errors_projection(pooled)
        )

        return {
            "has_issue_logits": has_issue_logits,
            "is_hidden_logits": is_hidden_logits,
            "errors_logits": errors_logits,
        }


AutoConfig.register("rqa_v2_2", RQAModelConfig)
AutoModel.register(RQAModelConfig, RQAModelHF)

print("✅ RQA-R2 зарегистрирован в Transformers")