"""
Chinese RoBERTa Multi-Task Model for Labor Arbitration Element Extraction.

Architecture:
  Chinese RoBERTa (hfl/chinese-roberta-wwm-ext)
    ├── Token Classification Head → BIO NER (31 labels)
    ├── Sequence Classification Heads → case cause, contract type, etc.
    ├── Span QA Head → text span extraction
    └── Numeric Regression Head → amount prediction

Joint loss: L = λ₁×L_NER + λ₂×L_CLS + λ₃×L_QA + λ₄×L_REG
"""

from __future__ import annotations

from typing import Any, Optional

# NOTE: transformers must be imported BEFORE torch on Windows to avoid segfault
from transformers import AutoConfig, AutoModel, AutoTokenizer
from transformers.modeling_outputs import TokenClassifierOutput

import torch
import torch.nn as nn


class ChineseRobertaMultiTask(nn.Module):
    """Multi-task model based on hfl/chinese-roberta-wwm-ext for
    joint NER, classification, span extraction, and numeric regression."""

    def __init__(
        self,
        model_name: str = "hfl/chinese-roberta-wwm-ext",
        num_ner_labels: int = 31,
        num_cause_classes: int = 6,
        num_contract_classes: int = 4,
        num_employment_classes: int = 5,
        num_bool_fields: int = 9,
        hidden_dropout_prob: float = 0.1,
        task_weights: Optional[list[float]] = None,
    ):
        super().__init__()

        # Shared encoder
        self.config = AutoConfig.from_pretrained(model_name)
        self.bert = AutoModel.from_pretrained(model_name)
        self.hidden_size = self.config.hidden_size
        self.dropout = nn.Dropout(hidden_dropout_prob)

        # ---- NER Head (token-level BIO classification) ----
        self.ner_head = nn.Linear(self.hidden_size, num_ner_labels)
        self.num_ner_labels = num_ner_labels

        # ---- Classification Head: case cause (6 classes) ----
        self.cause_classifier = nn.Linear(self.hidden_size, num_cause_classes)
        self.num_cause_classes = num_cause_classes

        # ---- Classification Head: contract type (4 classes) ----
        self.contract_classifier = nn.Linear(self.hidden_size, num_contract_classes)
        self.num_contract_classes = num_contract_classes

        # ---- Classification Head: employment type (5 classes) ----
        self.employment_classifier = nn.Linear(self.hidden_size, num_employment_classes)
        self.num_employment_classes = num_employment_classes

        # ---- Classification Head: binary/ternary fields (9 fields × 3 classes) ----
        self.bool_classifiers = nn.ModuleList([
            nn.Linear(self.hidden_size, 3) for _ in range(num_bool_fields)
        ])
        self.num_bool_fields = num_bool_fields

        # ---- Span QA Head (start/end logits) ----
        self.qa_head = nn.Linear(self.hidden_size, 2)  # start, end

        # ---- Numeric Regression Head ----
        self.reg_head = nn.Sequential(
            nn.Linear(self.hidden_size, 128),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(128, 64),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(64, 1),  # single scalar output
        )

        # Task weights for joint training
        self.task_weights = task_weights or [0.40, 0.25, 0.15, 0.10, 0.10]  # ner, cause, contract, employment, qa

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        """Initialize classifier heads with normal distribution."""
        for module in [self.ner_head, self.cause_classifier, self.contract_classifier,
                       self.employment_classifier, self.qa_head]:
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        for classifier in self.bool_classifiers:
            classifier.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            classifier.bias.data.zero_()

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        token_type_ids: Optional[torch.Tensor] = None,
        ner_labels: Optional[torch.Tensor] = None,
        cause_label: Optional[torch.Tensor] = None,
        contract_label: Optional[torch.Tensor] = None,
        employment_label: Optional[torch.Tensor] = None,
        bool_labels: Optional[torch.Tensor] = None,  # shape: [batch, 9]
        qa_start_labels: Optional[torch.Tensor] = None,
        qa_end_labels: Optional[torch.Tensor] = None,
        reg_values: Optional[torch.Tensor] = None,
        reg_mask: Optional[torch.Tensor] = None,  # 1 where reg target is valid
    ) -> dict[str, Any]:
        """
        Forward pass with optional multi-task loss computation.
        When labels are provided, returns (loss, logits_dict).
        When labels are None, returns logits_dict only.
        """
        # Shared encoder
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        sequence_output = outputs.last_hidden_state  # [batch, seq_len, hidden]
        pooled_output = outputs.pooler_output         # [batch, hidden]

        loss = torch.tensor(0.0, device=input_ids.device, dtype=sequence_output.dtype)
        num_active_tasks = 0

        # ---- NER logits ----
        ner_logits = self.ner_head(self.dropout(sequence_output))  # [batch, seq, 31]

        if ner_labels is not None:
            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
            ner_loss = loss_fct(
                ner_logits.view(-1, self.num_ner_labels),
                ner_labels.view(-1),
            )
            loss = loss + self.task_weights[0] * ner_loss
            num_active_tasks += 1

        # ---- Case cause classification ----
        cause_logits = self.cause_classifier(self.dropout(pooled_output))  # [batch, 6]

        if cause_label is not None:
            cause_loss = nn.CrossEntropyLoss()(cause_logits, cause_label)
            loss = loss + self.task_weights[1] * cause_loss
            num_active_tasks += 1

        # ---- Contract type classification ----
        contract_logits = self.contract_classifier(self.dropout(pooled_output))

        if contract_label is not None:
            contract_loss = nn.CrossEntropyLoss()(contract_logits, contract_label)
            loss = loss + self.task_weights[2] * contract_loss
            num_active_tasks += 1

        # ---- Employment type classification ----
        employment_logits = self.employment_classifier(self.dropout(pooled_output))

        if employment_label is not None:
            employment_loss = nn.CrossEntropyLoss()(employment_logits, employment_label)
            loss = loss + self.task_weights[3] * employment_loss
            num_active_tasks += 1

        # ---- Boolean field classification ----
        bool_logits_list = []
        bool_loss_total = torch.tensor(0.0, device=input_ids.device, dtype=sequence_output.dtype)
        for i, classifier in enumerate(self.bool_classifiers):
            logits = classifier(self.dropout(pooled_output))  # [batch, 3]
            bool_logits_list.append(logits)
            if bool_labels is not None and i < bool_labels.size(1):
                labels_i = bool_labels[:, i].long()
                mask_i = (labels_i >= 0)
                if mask_i.any():
                    bool_loss_total = bool_loss_total + nn.CrossEntropyLoss()(
                        logits[mask_i], labels_i[mask_i]
                    ) / self.num_bool_fields

        if bool_labels is not None:
            loss = loss + self.task_weights[4] * bool_loss_total
            num_active_tasks += 1

        # ---- Span QA ----
        qa_logits = self.qa_head(self.dropout(sequence_output))  # [batch, seq, 2]
        start_logits, end_logits = qa_logits[..., 0], qa_logits[..., 1]

        if qa_start_labels is not None and qa_end_labels is not None:
            qa_loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
            start_loss = qa_loss_fct(start_logits, qa_start_labels)
            end_loss = qa_loss_fct(end_logits, qa_end_labels)
            qa_loss = (start_loss + end_loss) / 2.0
            loss = loss + 0.10 * qa_loss  # relatively small weight for QA

        # ---- Numeric regression ----
        reg_preds = self.reg_head(self.dropout(pooled_output)).squeeze(-1)  # [batch]

        if reg_values is not None:
            if reg_mask is not None:
                mask = reg_mask.bool()
                if mask.any():
                    reg_loss = nn.MSELoss()(reg_preds[mask], reg_values[mask])
                    loss = loss + 0.05 * reg_loss  # small weight for regression
            else:
                reg_loss = nn.MSELoss()(reg_preds, reg_values)
                loss = loss + 0.05 * reg_loss

        return {
            "loss": loss if num_active_tasks > 0 else None,
            "ner_logits": ner_logits,
            "cause_logits": cause_logits,
            "contract_logits": contract_logits,
            "employment_logits": employment_logits,
            "bool_logits": bool_logits_list,
            "qa_start_logits": start_logits,
            "qa_end_logits": end_logits,
            "reg_preds": reg_preds,
        }

    @classmethod
    def from_pretrained(cls, model_path: str) -> "ChineseRobertaMultiTask":
        """Load a trained model from disk."""
        state_dict = torch.load(
            f"{model_path}/pytorch_model.bin",
            map_location="cpu",
            weights_only=True,
        )
        config = json.loads(
            open(f"{model_path}/model_config.json", encoding="utf-8").read()
        )
        model = cls(
            model_name=config.get("model_name", "hfl/chinese-roberta-wwm-ext"),
            num_ner_labels=config.get("num_ner_labels", 31),
            num_cause_classes=config.get("num_cause_classes", 6),
            num_contract_classes=config.get("num_contract_classes", 4),
            num_employment_classes=config.get("num_employment_classes", 5),
            num_bool_fields=config.get("num_bool_fields", 9),
        )
        model.load_state_dict(state_dict, strict=False)
        return model

    def save_pretrained(self, save_path: str):
        """Save model weights and config."""
        import os, json
        os.makedirs(save_path, exist_ok=True)
        torch.save(self.state_dict(), f"{save_path}/pytorch_model.bin")
        config = {
            "model_name": "hfl/chinese-roberta-wwm-ext",
            "num_ner_labels": self.num_ner_labels,
            "num_cause_classes": self.num_cause_classes,
            "num_contract_classes": self.num_contract_classes,
            "num_employment_classes": self.num_employment_classes,
            "num_bool_fields": self.num_bool_fields,
        }
        with open(f"{save_path}/model_config.json", "w", encoding="utf-8") as f:
            json.dump(config, f, ensure_ascii=False, indent=2)

    @torch.no_grad()
    def predict(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        tokenizer: Optional[Any] = None,
    ) -> dict[str, Any]:
        """Inference mode: extract elements from text."""
        self.eval()
        outputs = self.forward(input_ids=input_ids, attention_mask=attention_mask)

        # NER spans
        ner_preds = torch.argmax(outputs["ner_logits"], dim=-1)  # [batch, seq]

        # Cause prediction
        cause_pred = torch.argmax(outputs["cause_logits"], dim=-1)  # [batch]

        # Contract prediction
        contract_pred = torch.argmax(outputs["contract_logits"], dim=-1)

        # Employment prediction
        employment_pred = torch.argmax(outputs["employment_logits"], dim=-1)

        # Boolean predictions
        bool_preds = torch.stack([
            torch.argmax(logits, dim=-1) for logits in outputs["bool_logits"]
        ], dim=-1)  # [batch, 9]

        # Regression predictions
        reg_vals = outputs["reg_preds"]  # [batch]

        return {
            "ner_predictions": ner_preds.cpu().tolist(),
            "cause_predictions": cause_pred.cpu().tolist(),
            "contract_predictions": contract_pred.cpu().tolist(),
            "employment_predictions": employment_pred.cpu().tolist(),
            "bool_predictions": bool_preds.cpu().tolist(),
            "reg_predictions": reg_vals.cpu().tolist(),
        }


import json