src/evaluator.py

"""Evaluation: per-aspect metrics, aggregated overall comparison.

Two label encodings live in this project (keep them straight!):
  Per-aspect (ACSA):   0=Not_Mentioned, 1=Positive, 2=Negative
  Overall (3-class):   0=Negative,      1=Neutral,  2=Positive

`aspect_to_overall_sentiment` translates between them explicitly.
"""
import json
import logging
from pathlib import Path
from typing import Dict, Optional, Union

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
from sklearn.metrics import (
    f1_score, accuracy_score, confusion_matrix, classification_report,
)
from tqdm import tqdm

from . import config as cfg
from .dataset import ACSADataset, MetaACSADataset, OverallSentimentDataset
from .models import GatedAspectSemanticMetaFusionACSAModel, BertMetaFusionACSAModel, BertACSAModel, BertOverallModel
from .meta_encoder import MetaEncoder
from .trainer import get_device


logger = logging.getLogger(__name__)

# Label-code constants (DO NOT change carelessly; many tests depend on them).
ASPECT_NOT_MENTIONED = 0
ASPECT_POSITIVE = 1
ASPECT_NEGATIVE = 2

OVERALL_NEGATIVE = 0
OVERALL_NEUTRAL = 1
OVERALL_POSITIVE = 2


# ---------------------------------------------------------------------------
# Loaders (all use weights_only=False 閳?these are our own checkpoints)
# ---------------------------------------------------------------------------

def _load_ckpt(path: Path, device):
    return torch.load(path, map_location=device, weights_only=False)


def _move_model_to_device(model, device):
    """Move safely when Transformers creates meta tensors in Spaces."""
    try:
        return model.to(device)
    except NotImplementedError as exc:
        if "meta" not in str(exc).lower():
            raise
        logger.warning("Model contains meta tensors; using to_empty() before loading checkpoint weights.")
        return model.to_empty(device=device)
    except RuntimeError as exc:
        if "meta" not in str(exc).lower():
            raise
        logger.warning("Model contains meta tensors; using to_empty() before loading checkpoint weights.")
        return model.to_empty(device=device)


def _load_model_state(model, state_dict, strict=True):
    try:
        return model.load_state_dict(state_dict, strict=strict, assign=True)
    except TypeError:
        return model.load_state_dict(state_dict, strict=strict)


def _set_module_attr(root, dotted_name, value):
    module = root
    parts = dotted_name.split(".")
    for part in parts[:-1]:
        module = getattr(module, part)
    leaf = parts[-1]
    if isinstance(value, torch.nn.Parameter) and leaf in getattr(module, "_parameters", {}):
        module._parameters[leaf] = value
    elif leaf in getattr(module, "_buffers", {}):
        module._buffers[leaf] = value
    else:
        setattr(module, leaf, value)


def _has_meta_tensors(model) -> bool:
    for _, param in model.named_parameters():
        if getattr(param, "is_meta", False):
            return True
    for _, buf in model.named_buffers():
        if getattr(buf, "is_meta", False):
            return True
    return False


def _ensure_materialized_before_load(model, device):
    """Force a real tensor skeleton before checkpoint assignment in Spaces."""
    if _has_meta_tensors(model):
        logger.warning("Model skeleton contains meta tensors before checkpoint load; using to_empty().")
        model = model.to_empty(device=device)
    return model


def _materialize_meta_buffers(model, device):
    """Create real tensors for buffers that are not stored in older checkpoints."""
    for name, buf in list(model.named_buffers()):
        if not getattr(buf, "is_meta", False):
            continue
        shape = tuple(buf.shape)
        dtype = buf.dtype
        if name.endswith("position_ids") and len(shape) == 2:
            value = torch.arange(shape[1], dtype=dtype, device=device).expand(shape[0], -1)
        else:
            value = torch.zeros(shape, dtype=dtype, device=device)
        _set_module_attr(model, name, value)


def _materialize_meta_parameters(model, device):
    """Last-resort guard for Spaces when checkpoint loading leaves meta params."""
    made = []
    for name, param in list(model.named_parameters()):
        if not getattr(param, "is_meta", False):
            continue
        shape = tuple(param.shape)
        value = torch.empty(shape, dtype=param.dtype, device=device)
        if param.ndim >= 2:
            torch.nn.init.xavier_uniform_(value)
        else:
            torch.nn.init.zeros_(value)
        _set_module_attr(model, name, torch.nn.Parameter(value, requires_grad=param.requires_grad))
        made.append(name)
    if made:
        logger.warning("Materialized %d remaining meta parameters with fallback initialization: %s",
                       len(made), ", ".join(made[:10]))




def _refresh_runtime_buffers(model, device):
    """Rebuild non-persistent buffers after a Spaces meta-tensor fallback."""
    for module in model.modules():
        if hasattr(module, "source_prior_bias") and hasattr(module, "num_aspects"):
            prior_rows = []
            prior_cfg = getattr(cfg, "CROSS_ATTN_SOURCE_PRIOR", {})
            for aspect in cfg.ASPECTS[: module.num_aspects]:
                row = list(prior_cfg.get(aspect, [1.0 / 3.0] * 3))
                if len(row) < 3:
                    row = row + [1.0 / 3.0] * (3 - len(row))
                prior_rows.append(row[:3])
            prior = torch.tensor(prior_rows, dtype=torch.float, device=device)
            prior = prior - prior.mean(dim=-1, keepdim=True)
            module.source_prior_bias = prior
    _materialize_meta_buffers(model, device)
    _materialize_meta_parameters(model, device)
    meta_params = [name for name, param in model.named_parameters() if getattr(param, "is_meta", False)]
    if meta_params:
        raise RuntimeError(
            "Checkpoint did not materialize model parameters. First meta parameters: "
            + ", ".join(meta_params[:20])
        )
    return model


def load_meta_acsa(checkpoint_dir: Path = None, meta_encoder: Optional[MetaEncoder] = None,
                   device=None):
    if checkpoint_dir is None:
        checkpoint_dir = cfg.CHECKPOINT_DIR / "meta_acsa"
    checkpoint_dir = Path(checkpoint_dir)
    if device is None:
        device = get_device()
    ckpt = _load_ckpt(checkpoint_dir / "best.pt", device)
    bert_name = ckpt.get("config", {}).get("bert_name", cfg.BERT_MODEL_NAME)
    meta_in_dim = ckpt.get("config", {}).get("meta_in_dim",
                                              cfg.META_TFIDF_DIM + cfg.META_NUM_DIM)
    architecture = ckpt.get("config", {}).get("architecture", "legacy_meta_acsa")
    if meta_encoder is None:
        meta_encoder = MetaEncoder.load()
    if meta_encoder.total_dim != meta_in_dim:
        raise ValueError(
            f"Meta encoder dim {meta_encoder.total_dim} != checkpoint dim {meta_in_dim}. "
            f"Re-fit the encoder on the same train split used for training."
        )
    if architecture == "gated_aspect_semantic_meta_acsa":
        model = GatedAspectSemanticMetaFusionACSAModel(
            bert_name=bert_name,
            meta_in_dim=meta_in_dim,
        )
    else:
        model = BertMetaFusionACSAModel(bert_name=bert_name, meta_in_dim=meta_in_dim)
    model = _move_model_to_device(model, device)
    model = _ensure_materialized_before_load(model, device)
    _load_model_state(model, ckpt["model_state_dict"], strict=False)
    model = _refresh_runtime_buffers(model, device)
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir / "tokenizer")
    return model, tokenizer, meta_encoder, device


def load_acsa(checkpoint_dir: Path = None, device=None):
    if checkpoint_dir is None:
        checkpoint_dir = cfg.CHECKPOINT_DIR / "acsa"
    checkpoint_dir = Path(checkpoint_dir)
    if device is None:
        device = get_device()
    ckpt = _load_ckpt(checkpoint_dir / "best.pt", device)
    bert_name = ckpt.get("config", {}).get("bert_name", cfg.BERT_MODEL_NAME)
    model = BertACSAModel(bert_name=bert_name)
    model = _move_model_to_device(model, device)
    _load_model_state(model, ckpt["model_state_dict"])
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir / "tokenizer")
    return model, tokenizer, device


def load_bert_overall(checkpoint_dir: Path = None, device=None):
    if checkpoint_dir is None:
        checkpoint_dir = cfg.CHECKPOINT_DIR / "bert_overall"
    checkpoint_dir = Path(checkpoint_dir)
    if device is None:
        device = get_device()
    ckpt = _load_ckpt(checkpoint_dir / "best.pt", device)
    bert_name = ckpt.get("config", {}).get("bert_name", cfg.BERT_MODEL_NAME)
    model = BertOverallModel(bert_name=bert_name)
    model = _move_model_to_device(model, device)
    _load_model_state(model, ckpt["model_state_dict"])
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir / "tokenizer")
    return model, tokenizer, device


# ---------------------------------------------------------------------------
# Inference loops (return predictions in the same row order as test_df)
# ---------------------------------------------------------------------------

def predict_per_aspect(model, tokenizer, test_df, device,
                       meta_encoder: Optional[MetaEncoder] = None,
                       batch_size: int = 32):
    """Run a per-aspect model over test_df. Returns:
        all_preds  : list of NUM_ASPECTS lists, each length N
        all_labels : same shape, ground-truth aspect labels (if present)
    """
    test_df = test_df.reset_index(drop=True)
    if meta_encoder is not None:
        ds = MetaACSADataset(test_df, tokenizer, meta_encoder)
        loader = DataLoader(ds, batch_size=batch_size, shuffle=False)
        with_meta = True
    else:
        ds = ACSADataset(test_df, tokenizer)
        loader = DataLoader(ds, batch_size=batch_size, shuffle=False)
        with_meta = False

    all_preds = [[] for _ in range(cfg.NUM_ASPECTS)]
    all_labels = [[] for _ in range(cfg.NUM_ASPECTS)]
    with torch.no_grad():
        for batch in tqdm(loader, desc="predict per-aspect"):
            batch = {k: v.to(device) for k, v in batch.items()}
            if with_meta:
                out = model(batch["input_ids"], batch["attention_mask"],
                            batch["meta_features"])
            else:
                out = model(batch["input_ids"], batch["attention_mask"])
            preds = out["logits"].argmax(dim=-1).cpu().numpy()
            labels = batch["labels"].cpu().numpy()
            for i in range(cfg.NUM_ASPECTS):
                all_preds[i].extend(preds[:, i].tolist())
                all_labels[i].extend(labels[:, i].tolist())
    return all_preds, all_labels


def predict_overall_from_proposed(model, tokenizer, test_df, device,
                                  meta_encoder: MetaEncoder,
                                  batch_size: int = 32):
    """Use the Proposed model's overall_head to predict 3-class overall sentiment.

    Returns (preds_list, labels_list) where labels come from 'overall_label' column.
    """
    test_df = test_df.reset_index(drop=True)
    ds = MetaACSADataset(test_df, tokenizer, meta_encoder)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False)

    all_preds, all_labels = [], []
    with torch.no_grad():
        for batch in tqdm(loader, desc="predict overall (proposed head)"):
            batch = {k: v.to(device) for k, v in batch.items()}
            out = model(batch["input_ids"], batch["attention_mask"],
                        batch["meta_features"])
            preds = out["overall_logits"].argmax(dim=-1).cpu().numpy()
            all_preds.extend(preds.tolist())
            if "overall_labels" in batch:
                all_labels.extend(batch["overall_labels"].cpu().numpy().tolist())
            else:
                all_labels.extend(test_df["overall_label"].iloc[
                    len(all_labels):len(all_labels)+len(preds)].astype(int).tolist())
    return all_preds, all_labels


def evaluate_per_aspect(all_preds, all_labels):
    results = {"per_aspect": {}}
    f1s, accs = [], []
    for i, aspect in enumerate(cfg.ASPECTS):
        y_true, y_pred = all_labels[i], all_preds[i]
        report = classification_report(
            y_true, y_pred,
            labels=list(range(cfg.NUM_CLASSES)),
            target_names=cfg.LABEL_NAMES, output_dict=True, zero_division=0,
        )
        cm = confusion_matrix(y_true, y_pred, labels=list(range(cfg.NUM_CLASSES))).tolist()
        f1 = f1_score(y_true, y_pred, average="macro", zero_division=0)
        acc = accuracy_score(y_true, y_pred)
        results["per_aspect"][aspect] = {
            "macro_f1": float(f1), "accuracy": float(acc),
            "confusion_matrix": cm, "report": report,
        }
        f1s.append(f1); accs.append(acc)
    results["overall"] = {
        "mean_macro_f1": float(np.mean(f1s)) if f1s else 0.0,
        "mean_accuracy": float(np.mean(accs)) if accs else 0.0,
    }
    return results


def predict_overall(model, tokenizer, test_df, device, batch_size: int = 32):
    test_df = test_df.reset_index(drop=True)
    ds = OverallSentimentDataset(test_df, tokenizer)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False)
    all_preds, all_labels = [], []
    with torch.no_grad():
        for batch in loader:
            batch = {k: v.to(device) for k, v in batch.items()}
            out = model(**batch)
            all_preds.extend(out["logits"].argmax(dim=-1).cpu().numpy().tolist())
            all_labels.extend(batch["labels"].cpu().numpy().tolist())
    return all_preds, all_labels


def overall_metrics(y_true, y_pred):
    return {
        "test_macro_f1":    float(f1_score(y_true, y_pred, average="macro",    zero_division=0)),
        "test_accuracy":    float(accuracy_score(y_true, y_pred)),
        "test_weighted_f1": float(f1_score(y_true, y_pred, average="weighted", zero_division=0)),
    }


# ---------------------------------------------------------------------------
# Aggregation: per-aspect predictions -> single overall label
# ---------------------------------------------------------------------------

def aspect_to_overall_sentiment(aspect_preds) -> np.ndarray:
    """Improved voting rule for aggregating per-aspect 閳?overall.

    Rules (applied per review):
      1. Count n_pos and n_neg among MENTIONED aspects (skip Not_Mentioned).
      2. If no aspect is mentioned at all 閳?NEUTRAL (conservative default).
      3. If n_neg 閳?2 閳?NEGATIVE (multiple negative aspects = clearly unhappy).
      4. If n_neg == 1 and n_pos == 0 閳?NEGATIVE.
      5. If n_pos 閳?2 and n_neg == 0 閳?POSITIVE.
      6. If n_pos > n_neg (but not all positive) 閳?POSITIVE.
      7. If n_neg > n_pos 閳?NEGATIVE.
      8. Otherwise (tied, or only 1 pos and 0 neg) 閳?NEUTRAL.

    Input  (per-aspect):  0=NM, 1=Pos, 2=Neg
    Output (overall):     0=Neg, 1=Neu, 2=Pos
    """
    preds = np.asarray(aspect_preds)
    if preds.ndim == 1:
        preds = preds[None, :]

    n_pos = (preds == ASPECT_POSITIVE).sum(axis=1)
    n_neg = (preds == ASPECT_NEGATIVE).sum(axis=1)
    n_mentioned = n_pos + n_neg

    out = np.full(preds.shape[0], OVERALL_NEUTRAL, dtype=np.int64)

    # No aspects mentioned 閳?Neutral
    # n_neg 閳?2 閳?Negative (strong signal)
    out[n_neg >= 2] = OVERALL_NEGATIVE
    # n_neg == 1, n_pos == 0 閳?Negative
    out[(n_neg == 1) & (n_pos == 0)] = OVERALL_NEGATIVE
    # n_pos 閳?2, n_neg == 0 閳?Positive (strong signal)
    out[(n_pos >= 2) & (n_neg == 0)] = OVERALL_POSITIVE
    # n_pos > n_neg and n_pos 閳?2 閳?Positive
    out[(n_pos > n_neg) & (n_pos >= 2)] = OVERALL_POSITIVE
    # n_neg > n_pos 閳?Negative (override the 閳? positive if more negatives)
    out[n_neg > n_pos] = OVERALL_NEGATIVE

    return out


# ---------------------------------------------------------------------------
# Drilldown: category x aspect aggregation (application-layer)
# ---------------------------------------------------------------------------

def aggregate_aspect_distribution_by_category(test_df, aspect_preds_per_aspect,
                                              category_col: str = "leaf_category",
                                              top_k_cats: int = 10):
    df = test_df.copy().reset_index(drop=True)
    preds = np.array(aspect_preds_per_aspect).T  # (N, num_aspects)
    if preds.shape[0] != len(df):
        logger.warning("preds rows (%d) != df rows (%d); skipping aggregation.",
                       preds.shape[0], len(df))
        return None
    for i, a in enumerate(cfg.ASPECTS):
        df[f"pred_{a}"] = preds[:, i]

    if category_col not in df.columns:
        logger.warning("No %s column for drilldown.", category_col)
        return None
    df = df[df[category_col].notna() &
            (df[category_col].astype(str).str.len() > 0)]
    if df.empty:
        return None

    top_cats = df[category_col].value_counts().head(top_k_cats).index.tolist()
    rows = []
    for cat in top_cats:
        sub = df[df[category_col] == cat]
        if len(sub) < 5:
            continue
        for a in cfg.ASPECTS:
            p = sub[f"pred_{a}"]
            mentioned = p[p != ASPECT_NOT_MENTIONED]
            if len(mentioned) == 0:
                pos_share, neg_share = 0.0, 0.0
            else:
                pos_share = float((mentioned == ASPECT_POSITIVE).mean())
                neg_share = float((mentioned == ASPECT_NEGATIVE).mean())
            rows.append({
                "category": cat, "aspect": a, "n_total": len(sub),
                "n_mentioned": int(len(mentioned)),
                "positive_share": pos_share, "negative_share": neg_share,
            })
    return pd.DataFrame(rows)


# ---------------------------------------------------------------------------
# Single-review formatted output (for the customer's stated need)
# ---------------------------------------------------------------------------

def format_aspect_summary(per_aspect_pred: np.ndarray) -> Dict[str, str]:
    """For one review: {'SIZE': 'Positive', 'MATERIAL': 'Not_Mentioned', ...}."""
    return {aspect: cfg.LABEL_NAMES[int(per_aspect_pred[i])]
            for i, aspect in enumerate(cfg.ASPECTS)}