from __future__ import annotations

import math
from typing import Optional


def print_summary(categories, original_results, finetuned_results,
                  kme_metrics: Optional[dict] = None,
                  relation_config=None):
    if relation_config is not None:
        locality_cats = list(relation_config.locality_categories)
        efficacy_cat = relation_config.efficacy_category
    else:
        locality_cats = ["bathroom_with_toilet", "non_bathroom_with_toilet", "unrelated"]
        efficacy_cat = "bathroom_no_toilet"

    print("\nKeyword mention rate:")
    header = f"{'Category':<35} {'Orig':>7} {'FT':>7} {'Delta':>7}"
    print(header)
    print("-" * len(header))
    for cat in categories:
        orig = original_results[cat]["mention_rate_keyword"]
        ft = finetuned_results[cat]["mention_rate_keyword"]
        delta = ft - orig
        print(f"{cat:<35} {orig:>6.1%} {ft:>6.1%} {delta:>+7.1%}")

    print("\nVLM-guided mention rate (LLM judge == YES):")
    header_llm = f"{'Category':<35} {'Orig':>7} {'FT':>7} {'Delta':>7}"
    print(header_llm)
    print("-" * len(header_llm))
    for cat in categories:
        orig = original_results[cat]["mention_rate_llm"]
        ft = finetuned_results[cat]["mention_rate_llm"]
        delta = ft - orig
        print(f"{cat:<35} {orig:>6.1%} {ft:>6.1%} {delta:>+7.1%}")

    print("\nAlignment score (image-text model):")
    header2 = f"{'Category':<35} {'Orig':>7} {'FT':>7} {'Delta':>7}"
    print(header2)
    print("-" * len(header2))
    for cat in categories:
        orig_clip = original_results[cat]["avg_clip_score"]
        ft_clip = finetuned_results[cat]["avg_clip_score"]
        delta_clip = (ft_clip - orig_clip) if not (math.isnan(orig_clip) or math.isnan(ft_clip)) else float("nan")
        orig_s = "N/A" if math.isnan(orig_clip) else f"{orig_clip:.3f}"
        ft_s = "N/A" if math.isnan(ft_clip) else f"{ft_clip:.3f}"
        delta_s = "N/A" if math.isnan(delta_clip) else f"{delta_clip:+.3f}"
        print(f"{cat:<35} {orig_s:>7} {ft_s:>7} {delta_s:>7}")

    unrelated_orig = original_results.get("unrelated", {}).get("avg_clip_score", float("nan"))
    unrelated_ft = finetuned_results.get("unrelated", {}).get("avg_clip_score", float("nan"))
    if not (math.isnan(unrelated_orig) or math.isnan(unrelated_ft)):
        regression = unrelated_ft - unrelated_orig
        if regression < -0.05:
            print(f"\n  *** REGRESSION WARNING: alignment delta on 'unrelated' = {regression:+.3f} (<-0.05) ***")

    print("\nQuality-gated mention rate (caption_quality >= 3):")
    header3 = f"{'Category':<35} {'Orig':>7} {'FT':>7}"
    print(header3)
    print("-" * len(header3))
    for cat in categories:
        orig = original_results[cat]["quality_gated_mention_rate"]
        ft = finetuned_results[cat]["quality_gated_mention_rate"]
        print(f"{cat:<35} {orig:>6.1%} {ft:>6.1%}")

    print("\nCaption quality (LLM judge, 1-5):")
    header4 = f"{'Category':<35} {'Orig':>7} {'FT':>7} {'Delta':>7}"
    print(header4)
    print("-" * len(header4))
    for cat in categories:
        orig = original_results[cat]["avg_caption_quality"]
        ft = finetuned_results[cat]["avg_caption_quality"]
        print(f"{cat:<35} {orig:>7.2f} {ft:>7.2f} {(ft - orig):>+7.2f}")

    if kme_metrics is None:
        return

    # ------------------------------------------------------------------
    # KME metrics (comparative)
    # ------------------------------------------------------------------
    print("\n" + "=" * 70)
    print("Knowledge Editing Metrics")
    print("=" * 70)

    def _fmt(v: float) -> str:
        return "N/A" if (isinstance(v, float) and math.isnan(v)) else f"{v:.3f}"

    # Efficacy
    print(f"\nEfficacy (1 - mention_rate on {efficacy_cat}):")
    print(f"  efficacy_keyword:  {_fmt(kme_metrics.get('efficacy_keyword', float('nan')))}")

    # Generality
    print(f"\nGenerality (efficacy split by prompt type):")
    print(f"  seen prompts:      {_fmt(kme_metrics.get('generality_seen', float('nan')))}")
    print(f"  unseen prompts:    {_fmt(kme_metrics.get('generality_unseen', float('nan')))}")

    # Locality per category
    sim_keys = ["exact_match", "rouge_l", "bert_score_f1"]

    print(f"\nLocality (text similarity: original vs. finetuned on preserved categories):")
    loc_header = f"  {'Category':<30} {'ExactMatch':>11} {'ROUGE-L':>9} {'BERTScore':>10}"
    print(loc_header)
    print("  " + "-" * (len(loc_header) - 2))
    for cat in locality_cats:
        vals = [_fmt(kme_metrics.get(f"locality/{cat}/{k}", float("nan"))) for k in sim_keys]
        print(f"  {cat:<30} {vals[0]:>11} {vals[1]:>9} {vals[2]:>10}")

    # Consistency (aggregate)
    print(f"\nConsistency (aggregated across all preserved categories):")
    for k in sim_keys:
        print(f"  {k:<20} {_fmt(kme_metrics.get(f'consistency/{k}', float('nan')))}")