"""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 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,
        ).to(device)
    else:
        model = BertMetaFusionACSAModel(bert_name=bert_name, meta_in_dim=meta_in_dim).to(device)
    model.load_state_dict(ckpt["model_state_dict"], strict=False)
    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).to(device)
    model.load_state_dict(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).to(device)
    model.load_state_dict(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)}