"""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):
    # Always read checkpoints on CPU first. This avoids device/meta edge cases
    # in Hugging Face Spaces, then we move the fully materialized model later.
    return torch.load(path, map_location="cpu", weights_only=False)


def _construct_model_on_cpu(factory):
    """Build modules with real CPU tensors even if the runtime default device is meta."""
    old_device = None
    can_restore = False
    if hasattr(torch, "get_default_device") and hasattr(torch, "set_default_device"):
        try:
            old_device = torch.get_default_device()
            torch.set_default_device("cpu")
            can_restore = True
        except Exception:
            old_device = None
            can_restore = False
    try:
        try:
            with torch.device("cpu"):
                return factory()
        except Exception:
            return factory()
    finally:
        # Do not restore a meta default device inside Spaces; that would make
        # later tensors empty again. Restoring normal devices is safe.
        if can_restore and old_device is not None and str(old_device) != "meta":
            try:
                torch.set_default_device(old_device)
            except Exception:
                pass


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 tensor" not in str(exc):
            raise
        logger.warning("Model contains meta tensors; using to_empty() before loading checkpoint weights.")
        return model.to_empty(device=device)


def _has_meta_tensors(model) -> bool:
    return any(p.device.type == "meta" for p in model.parameters()) or any(
        b is not None and b.device.type == "meta" for b in model.buffers()
    )


def _reset_known_bert_buffers(model, device="cpu"):
    """Reset non-persistent BERT embedding buffers after to_empty()."""
    bert = getattr(model, "bert", None)
    emb = getattr(bert, "embeddings", None)
    if emb is None:
        return model
    device = torch.device(device)
    for name, buffer in list(emb._buffers.items()):
        if buffer is None or name not in {"position_ids", "token_type_ids"}:
            continue
        shape = tuple(buffer.shape)
        if not shape:
            continue
        if name == "position_ids":
            values = torch.arange(shape[-1], dtype=torch.long, device=device)
            emb._buffers[name] = values.view(1, -1).expand(shape).clone()
        elif name == "token_type_ids":
            emb._buffers[name] = torch.zeros(shape, dtype=torch.long, device=device)
    return model


def _load_model_state(model, state_dict, strict=True):
    """Load checkpoint weights and report mismatches instead of hiding them."""
    if _has_meta_tensors(model):
        logger.warning("Model was built with meta tensors; materializing empty CPU tensors before loading checkpoint.")
        model.to_empty(device="cpu")
    try:
        result = model.load_state_dict(state_dict, strict=strict, assign=True)
    except TypeError:
        result = model.load_state_dict(state_dict, strict=strict)
    _reset_known_bert_buffers(model, device="cpu")
    missing = list(getattr(result, "missing_keys", []))
    unexpected = list(getattr(result, "unexpected_keys", []))
    if missing:
        logger.warning("Missing checkpoint keys: %s", missing[:20])
    if unexpected:
        logger.warning("Unexpected checkpoint keys: %s", unexpected[:20])
    return result


def _real_buffer_like(buffer, name: str, device):
    shape = tuple(buffer.shape)
    dtype = buffer.dtype
    if "position_ids" in name and shape:
        values = torch.arange(shape[-1], device=device, dtype=torch.long)
        return values.view(1, -1).expand(shape).clone()
    return torch.zeros(shape, device=device, dtype=dtype)


def _fix_meta_buffers_only(model, device):
    """After assign=True loading, only non-checkpoint BERT buffers may remain meta."""
    device = torch.device(device)
    meta_params = [n for n, p in model.named_parameters() if p.device.type == "meta"]
    if meta_params:
        raise RuntimeError(
            "Checkpoint did not materialize model parameters. First meta params: "
            + ", ".join(meta_params[:20])
        )
    fixed = []
    for module_name, module in model.named_modules():
        prefix = f"{module_name}." if module_name else ""
        for name, buffer in list(module._buffers.items()):
            if buffer is not None and buffer.device.type == "meta":
                full_name = prefix + name
                module._buffers[name] = _real_buffer_like(buffer, full_name, device)
                fixed.append(full_name)
    if fixed:
        logger.warning("Materialized leftover meta buffers: %s", fixed[:20])
    remaining = [n for n, b in model.named_buffers() if b.device.type == "meta"]
    if remaining:
        raise RuntimeError(f"Model still has meta buffers after loading: {remaining[: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 = _construct_model_on_cpu(lambda: GatedAspectSemanticMetaFusionACSAModel(
            bert_name=bert_name,
            meta_in_dim=meta_in_dim,
        ))
    else:
        model = _construct_model_on_cpu(lambda: BertMetaFusionACSAModel(bert_name=bert_name, meta_in_dim=meta_in_dim))
    _load_model_state(model, ckpt["model_state_dict"], strict=False)
    model = _fix_meta_buffers_only(model, device)
    model = model.to(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 = _construct_model_on_cpu(lambda: BertACSAModel(bert_name=bert_name))
    _load_model_state(model, ckpt["model_state_dict"])
    model = _fix_meta_buffers_only(model, device)
    model = model.to(device)
    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 = _construct_model_on_cpu(lambda: BertOverallModel(bert_name=bert_name))
    _load_model_state(model, ckpt["model_state_dict"])
    model = _fix_meta_buffers_only(model, device)
    model = model.to(device)
    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)}