# tanml/ui/reports.py
"""
Report generation logic for TanML UI.

Extracts the Word document generation functions from app.py.
"""

from __future__ import annotations

import io
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional
from importlib.resources import files


from docx import Document
from docx.shared import Inches, RGBColor
from docx.oxml.ns import nsdecls
from docx.oxml import parse_xml

# TanML UI Helpers
from tanml.ui.glossary import GLOSSARY
from tanml.ui.narratives import (
    story_performance as _story_performance,
    story_features as _story_features,
    story_overfitting as _story_overfitting,
    story_drift as _story_drift,
    story_stress as _story_stress,
    story_shap as _story_shap,
)
import matplotlib.colors as mcolors


def _choose_report_template(task_type: str) -> Path:
    """Return the correct .docx template for 'regression' or 'classification'."""
    # packaged location (recommended): tanml/report/templates/*.docx
    try:
        templates_pkg = files("tanml.report.templates")
        name = "report_template_reg.docx" if task_type == "regression" else "report_template_cls.docx"
        p = templates_pkg / name
        if p.is_file():
            return Path(str(p))
    except Exception:
        pass

    # repo fallback
    repo_guess = Path(__file__).resolve().parents[1] / "report" / "templates"
    p2 = repo_guess / ("report_template_reg.docx" if task_type == "regression" else "report_template_cls.docx")
    if p2.exists():
        return p2

    # cwd fallback
    return Path.cwd() / ("report_template_reg.docx" if task_type == "regression" else "report_template_cls.docx")


def _filter_metrics_for_task(summary: Dict[str, Any]) -> Dict[str, Any]:
    """Keep only metrics relevant to the task_type inside summary."""
    if not isinstance(summary, dict):
        return summary or {}

    task = summary.get("task_type")
    cls_keys = {"auc", "ks", "f1", "pr_auc", "rules_failed", "task_type"}
    reg_keys = {"rmse", "mae", "r2", "rules_failed", "task_type"}  

    if task == "classification":
        return {k: v for k, v in summary.items() if k in cls_keys}
    if task == "regression":
        return {k: v for k, v in summary.items() if k in reg_keys}
    return summary


def _fmt2(v, *, decimals=2, dash="—"):
    if v is None:
        return dash
    try:
        if isinstance(v, float):
            return f"{v:.{decimals}f}"
        if isinstance(v, int):
            return str(v)
        return str(v)
    except Exception:
        return dash


def _add_md_paragraph(doc, text: str, style=None):
    """
    Add a paragraph to doc, parsing bold markdown (**text**).
    
    Args:
        doc: docx Document object
        text: String with **bold** markers
        style: Optional style name
    """
    p = doc.add_paragraph(style=style)
    # Split by double asterisks
    parts = text.split("**")
    
    for i, part in enumerate(parts):
        run = p.add_run(part)
        # Even indices (0, 2, 4...) are normal text
        # Odd indices (1, 3, 5...) are wrapped in ** and should be bold
        if i % 2 == 1:
            run.bold = True
    return p
    


def _generate_dev_report_docx(dev_data):
    doc = Document()
    doc.add_heading('Model Development Report', 0)
    doc.add_paragraph(f"Generated by TanML on: {datetime.now().strftime('%Y-%m-%d %H:%M')}")
    
    # 1. Executive Summary
    doc.add_heading('1. Executive Summary', level=1)
    p = doc.add_paragraph()
    if dev_data:
        task = dev_data.get("task_type", "Unknown")
        algo = dev_data.get("model_config", {}).get("algorithm", "Unknown")
        p.add_run(f"This report summarizes the development of a ").bold = False
        p.add_run(f"{task.title()}").bold = True
        p.add_run(f" model using ").bold = False
        p.add_run(f"{algo}").bold = True
        p.add_run(".")
        
        # Smart Summary
        mets = dev_data.get("metrics", {})
        if task == "classification":
            score = mets.get("accuracy", 0)
            p2 = doc.add_paragraph(f"The model achieved an Accuracy of {score:.4f}.")
            if score > 0.98: p2.add_run(" (Note: Extremely high accuracy may indicate overfitting or data leakage.)").bold = True
        else:
            score = mets.get("r2", 0)
            p2 = doc.add_paragraph(f"The model achieved an R² Score of {score:.4f}.")

        # Check CV Consistency
        cv_met = dev_data.get("cv_metrics", {})
        if cv_met:
            # Try multiple keys for robustness (simple keys first)
            cv_s = cv_met.get("accuracy", 
                   cv_met.get("r2", 
                   cv_met.get("R2",
                   cv_met.get("test_accuracy", 
                   cv_met.get("test_r2", 
                   cv_met.get("mean_test_score", 0))))))
                   
            diff = abs(score - cv_s)
            if diff > 0.05:
                _add_md_paragraph(doc, f"⚠️ **Note**: Validation score ({score:.2f}) differs from Cross-Validation mean ({cv_s:.2f}) by >5%.", style='List Bullet')
            else:
                 _add_md_paragraph(doc, "✅ **Stability**: Validation score aligns with Cross-Validation results.", style='List Bullet')
    else:
        p.add_run("No development data available.")
        
    # 2. Results
    if dev_data:
        doc.add_heading('2. Model Performance', level=1)
        
        # Narrative
        narrative = _story_performance(dev_data.get("metrics", {}), dev_data.get("task_type"))
        if narrative:
            doc.add_heading('Executive Insight', level=2)
            doc.add_paragraph(narrative)
            
        doc.add_heading('Final Model Metrics (Full Dataset)', level=2)
        table = doc.add_table(rows=1, cols=2)
        table.style = 'Table Grid'
        hdr = table.rows[0].cells
        hdr[0].text = 'Metric'; hdr[1].text = 'Score'
        for k, v in dev_data.get("metrics", {}).items():
            r = table.add_row().cells
            r[0].text = k.replace("_", " ").title()
            r[1].text = f"{v:.4f}" if isinstance(v, (float, int)) else str(v)
            
        # CV Metrics
        cv_met = dev_data.get("cv_metrics", {})
        if cv_met:
            doc.add_heading('Cross-Validation Metrics (Mean)', level=2)
            table2 = doc.add_table(rows=1, cols=2)
            table2.style = 'Table Grid'
            hdr2 = table2.rows[0].cells
            hdr2[0].text = 'Metric'; hdr2[1].text = 'Mean Score'
            for k, v in cv_met.items():
                r = table2.add_row().cells
                r[0].text = k.replace("_", " ").title()
                r[1].text = f"{v:.4f}" if isinstance(v, (float, int)) else str(v)

        # Images
        imgs = dev_data.get("images", {})
        cv_imgs = dev_data.get("cv_images", {})
        
        if cv_imgs:
            doc.add_heading('2.1 Cross-Validation Plots', level=2)
            for name, img_bytes in cv_imgs.items():
                    doc.add_paragraph(name.replace("_", " ").title())
                    doc.add_picture(io.BytesIO(img_bytes), width=Inches(4))
                    
        if imgs:
            doc.add_heading('2.2 Final Model Diagnostics', level=2)
            for name, img_bytes in imgs.items():
                    doc.add_paragraph(name.replace("_", " ").title())
                    doc.add_picture(io.BytesIO(img_bytes), width=Inches(4))

    # Appendix: Glossary
    doc.add_page_break()
    doc.add_heading('Appendix: Guide to Metrics', level=1)
    doc.add_paragraph("Key terms used in this report:")
    
    if dev_data.get("cv_metrics"): 
        doc.add_paragraph(f"**Cross-Validation**: {GLOSSARY['Cross-Validation']}")
    
    metrics_to_show = []
    if dev_data.get("task_type") == "classification":
        metrics_to_show = ["ROC AUC", "PR AUC", "F1 Score", "Accuracy", "Precision", "Recall", "Log Loss", "Brier Score", "MCC", "KS Statistic", "Confusion Matrix"]
    else:
        metrics_to_show = ["RMSE", "MAE", "R2 Score", "Median AE"]
        
    for m in metrics_to_show:
        if m in GLOSSARY:
             doc.add_paragraph(f"**{m}**: {GLOSSARY[m]}", style='List Bullet')



    # Save
    buffer = io.BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    return buffer


def _generate_eval_report_docx(buf):
    
    # Defensive type check - ensure buf is a dict
    if isinstance(buf, list):
        buf = buf[0] if buf else {}
    if not isinstance(buf, dict):
        buf = {}
    
    doc = Document()
    doc.add_heading('Model Evaluation Report', 0)
    doc.add_paragraph(f"Generated by TanML on: {datetime.now().strftime('%Y-%m-%d %H:%M')}")

    
    # Determine task type from evaluation data
    ev = buf.get("evaluation", {})
    task_type = ev.get("task_type", "classification")

    # 1. Executive Summary (NEW)
    doc.add_heading('1. Executive Summary', level=1)
    
    warnings = []
    passing = []
    
    # Check Drift
    drift = buf.get("drift", [])
    if isinstance(drift, dict): drift = [drift]
    if drift:
        n_crit = sum(1 for d in drift if "Critical" in str(d.get("PSI Status","")) or "Critical" in str(d.get("KS Status","")))
        if n_crit > 0:
             _add_md_paragraph(doc, f"⚠️ **Summary**: {n_crit} features show Critical Drift.", style='List Bullet')
        else:
             _add_md_paragraph(doc, "✅ **Summary**: No critical data drift detected.", style='List Bullet')
    
    # Check Stress
    stress = buf.get("stress", [])
    if isinstance(stress, dict): stress = [stress]
    if stress and stress[0]:
        # Simple heuristic: if delta > 0.1
        s0 = stress[0]
        drop = s0.get("delta_accuracy", s0.get("delta_rmse", 0))
        if abs(drop) > 0.1:
             _add_md_paragraph(doc, f"⚠️ **Summary**: Performance drops significantly ({abs(drop):.2%}) under noise.", style='List Bullet')
        else:
             _add_md_paragraph(doc, "✅ **Summary**: Model is robust to noise (Stress Test passed).", style='List Bullet')
            
    # Check Overfitting
    m_tr = ev.get("metrics_train", {})
    m_te = ev.get("metrics_test", {})
    if task_type == "classification":
        tr_s = m_tr.get("Accuracy", 0)
        te_s = m_te.get("Accuracy", 0)
    else:
        tr_s = m_tr.get("R2", 0)
        te_s = m_te.get("R2", 0)
        
    if (tr_s - te_s) > 0.15:
        warnings.append(f"Potential Overfitting: Train score ({tr_s:.2f}) is much higher than Test score ({te_s:.2f}).")
    
    # Check Cluster Coverage (Heuristic)
    cluster_data = buf.get("cluster_coverage")
    if isinstance(cluster_data, list): cluster_data = cluster_data[0] if cluster_data else {}
    if not isinstance(cluster_data, dict): cluster_data = {}
    
    if cluster_data:
        ood_pct = cluster_data.get("ood_pct", 0)
        cov_pct = cluster_data.get("coverage_pct", 0)
        if ood_pct > 10:
            warnings.append(f"Significant Out-of-Distribution data detected ({ood_pct:.1f}% samples).")
        if cov_pct < 80:
             warnings.append(f"Poor Input Space Coverage: only {cov_pct:.1f}% of clusters covered.")
             
    if warnings:
        p = doc.add_paragraph()
        run = p.add_run("⚠️ Attention Required:")
        run.bold = True
        run.font.color.rgb = RGBColor(255, 0, 0)
        for w in warnings:
            _add_md_paragraph(doc, w, style='List Bullet')
    else:
        p = doc.add_paragraph()
        run = p.add_run("✅ Executive Assessment: Model looks healthy.")
        run.bold = True
        run.font.color.rgb = RGBColor(0, 150, 0)
        for p_msg in passing:
            _add_md_paragraph(doc, p_msg, style='List Bullet')

    doc.add_paragraph(f"This report covers {task_type.title()} performance, stability, and explainability checks.")

    # 1.1 Key Narrative Insights (Promoted from detailed sections)
    doc.add_heading('Key Insights', level=2)
    
    # Stability Story
    m_tr = ev.get("metrics_train", {})
    m_te = ev.get("metrics_test", {})
    narrative_stability = _story_overfitting(m_tr, m_te)
    if narrative_stability:
        _add_md_paragraph(doc, f"**Stability**: {narrative_stability}", style='List Bullet')

    # Drift Story
    drift_data = buf.get("drift")
    narrative_drift = _story_drift(drift_data)
    if narrative_drift:
        _add_md_paragraph(doc, f"**Drift**: {narrative_drift}", style='List Bullet')

    # Stress Story
    stress_data = buf.get("stress")
    narrative_stress = _story_stress(stress_data)
    if narrative_stress:
        _add_md_paragraph(doc, f"**Robustness**: {narrative_stress}", style='List Bullet')
        
    # SHAP Story
    expl = buf.get("explainability", {})
    if expl and expl.get("status") == "ok":
         narrative_shap = _story_shap(expl)
         if narrative_shap:
             _add_md_paragraph(doc, f"**Drivers**: {narrative_shap}", style='List Bullet')

    # 2. Results Comparison
    if ev:
        doc.add_heading('2. Results Comparison (Train vs Test)', level=1)
        
        # Narrative
        m_tr = ev.get("metrics_train", {})
        m_te = ev.get("metrics_test", {})
        narrative = _story_overfitting(m_tr, m_te)
        if narrative:
            doc.add_heading('Stability Check', level=2)
            # Duplicate Overfitting Warning if present
            if (tr_s - te_s) > 0.15:
                doc.add_paragraph(f"⚠️ **Attention**: Specific check indicates potential overfitting (Train {tr_s:.2f} vs Test {te_s:.2f}).", style='List Bullet')
            _add_md_paragraph(doc, narrative)
        
        # Metrics Table
        doc.add_heading('Metrics', level=2)
        table = doc.add_table(rows=1, cols=3)
        table.style = 'Table Grid'
        h = table.rows[0].cells
        h[0].text = 'Metric'; h[1].text = 'Train'; h[2].text = 'Test'
        
        m_tr = ev.get("metrics_train", {})
        m_te = ev.get("metrics_test", {})
        all_keys = set(m_tr.keys()) | set(m_te.keys())
        for k in sorted(all_keys):
            if k in ["confusion_matrix", "curves", "threshold_info"]: continue
            r = table.add_row().cells
            r[0].text = k
            r[1].text = f"{m_tr.get(k, 0):.4f}"
            r[2].text = f"{m_te.get(k, 0):.4f}"
        
        # Diagnostic Plots
        ev_imgs = ev.get("images", {})
        if ev_imgs:
                doc.add_heading('Diagnostic Plots', level=2)
                
                # Identify unique plot basenames
                unique_bases = set()
                for name in ev_imgs.keys():
                    if name.startswith("train_"): unique_bases.add(name.replace("train_", ""))
                    elif name.startswith("test_"): unique_bases.add(name.replace("test_", ""))
                    else: unique_bases.add(name)
                
                sorted_bases = sorted(list(unique_bases))
                
                for base in sorted_bases:
                    # Check for pair
                    train_key = f"train_{base}"
                    test_key = f"test_{base}"
                    
                    if train_key in ev_imgs and test_key in ev_imgs:
                        doc.add_heading(base.replace("_", " ").title(), level=3)
                        table = doc.add_table(rows=1, cols=2)
                        table.autofit = True
                        r = table.rows[0].cells
                        
                        # Train Left
                        p1 = r[0].add_paragraph("Train")
                        p1.alignment = 1 # Center
                        r[0].add_paragraph().add_run().add_picture(io.BytesIO(ev_imgs[train_key]), width=Inches(3.0))
                        
                        # Test Right
                        p2 = r[1].add_paragraph("Test")
                        p2.alignment = 1 # Center
                        r[1].add_paragraph().add_run().add_picture(io.BytesIO(ev_imgs[test_key]), width=Inches(3.0))
                    
                    elif base in ev_imgs: # Standalone (no prefix)
                        doc.add_paragraph(base.replace("_", " ").title())
                        doc.add_picture(io.BytesIO(ev_imgs[base]), width=Inches(4))
                    
                    else: # One of the prefixed exists but not the other?
                         if train_key in ev_imgs:
                             doc.add_paragraph(f"Train {base}".replace("_", " ").title())
                             doc.add_picture(io.BytesIO(ev_imgs[train_key]), width=Inches(4))
                         if test_key in ev_imgs:
                             doc.add_paragraph(f"Test {base}".replace("_", " ").title())
                             doc.add_picture(io.BytesIO(ev_imgs[test_key]), width=Inches(4))
    else:
        doc.add_paragraph("No evaluation data found.")

    # 2. Risk (Drift/Stress)
    doc.add_heading('2. Risk Assessment', level=1)
    
    # Drift
    drift_data = buf.get("drift")
    drift_imgs = buf.get("drift_images", {})
    doc.add_heading('2.1 Drift Analysis', level=2)
    
    # Duplicate Drift Heuristic
    if drift:
        n_crit = sum(1 for d in drift if "Critical" in str(d.get("PSI Status","")) or "Critical" in str(d.get("KS Status","")))
        if n_crit > 0:
             _add_md_paragraph(doc, f"⚠️ **Summary**: {n_crit} features show Critical Drift.", style='List Bullet')
        else:
             _add_md_paragraph(doc, "✅ **Summary**: No critical drift detected.", style='List Bullet')

    # Narrative
    nar = _story_drift(drift_data)
    if nar:
        _add_md_paragraph(doc, nar)
    
    if drift_data:
        # Ensure drift_data is a list of dicts
        if isinstance(drift_data, dict):
            drift_data = [drift_data]
        if not isinstance(drift_data, list) or not drift_data:
            doc.add_paragraph("Drift analysis not run.")
        else:
            table = doc.add_table(rows=1, cols=5)

        table.style = 'Table Grid'
        h = table.rows[0].cells
        h[0].text = 'Feature'; h[1].text = 'PSI'; h[2].text = 'PSI Status'; h[3].text = 'KS Stat'; h[4].text = 'KS Status'
        for d in drift_data:
            r = table.add_row().cells
            r[0].text = str(d['Feature'])
            r[1].text = f"{d['PSI']:.4f}"
            r[2].text = str(d.get('PSI Status', d.get('Status', ''))).replace("🟢 ", "").replace("🟠 ", "").replace("🔴 ", "").replace("🟡 ", "").replace("● ", "")
            r[3].text = f"{d.get('KS Stat', 0):.4f}"
            r[4].text = str(d.get('KS Status', '')).replace("🟢 ", "").replace("🟠 ", "").replace("🔴 ", "").replace("🟡 ", "")
        
        if "top_distribution" in drift_imgs:
            doc.add_paragraph("Top Drifting Feature Distribution")
            doc.add_picture(io.BytesIO(drift_imgs["top_distribution"]), width=Inches(5))
    else:
        doc.add_paragraph("Drift analysis not run.")
    
    # Stress
    stress_data = buf.get("stress")
    doc.add_heading('2.2 Stress Testing', level=2)
    
    # Duplicate Stress Heuristic
    if stress and stress[0]:
        s0 = stress[0]
        drop = s0.get("delta_accuracy", s0.get("delta_rmse", 0))
        if abs(drop) > 0.1:
             _add_md_paragraph(doc, f"⚠️ **Summary**: Performance drops significantly ({abs(drop):.2%}) under stress.", style='List Bullet')
        else:
             _add_md_paragraph(doc, "✅ **Summary**: Model performance remains stable under stress.", style='List Bullet')

    # Narrative
    nar_stress = _story_stress(stress_data)
    if nar_stress:
        _add_md_paragraph(doc, nar_stress)
        
    if stress_data:
            # Ensure stress_data is a list of dicts
            if isinstance(stress_data, dict):
                stress_data = [stress_data]
            if not isinstance(stress_data, list) or not stress_data:
                doc.add_paragraph("Stress test not run.")
            elif not isinstance(stress_data[0], dict):
                doc.add_paragraph("Invalid stress test data format.")
            else:
                doc.add_paragraph("Stress Test Results (Perturbed Data)")
                table = doc.add_table(rows=1, cols=len(stress_data[0]))

            table.style = 'Table Grid'
            # simplistic table dump
            hdr = table.rows[0].cells
            for i, k in enumerate(stress_data[0].keys()): hdr[i].text = str(k)
            for row_dat in stress_data:
                r = table.add_row().cells
                for i, v in enumerate(row_dat.values()): r[i].text = f"{v:.4f}" if isinstance(v, float) else str(v)
    else:
            doc.add_paragraph("Stress test not run.")

    # Cluster Coverage Check
    cluster_data = buf.get("cluster_coverage")
    cluster_imgs = buf.get("cluster_images", {})
    doc.add_heading('2.3 Input Cluster Coverage Check', level=2)
    
    if cluster_data:
        # Ensure cluster_data is a dict (not a list)
        if isinstance(cluster_data, list):
            cluster_data = cluster_data[0] if cluster_data else {}
        if not isinstance(cluster_data, dict):
            cluster_data = {}
        
        # Summary narrative
        coverage_pct = cluster_data.get("coverage_pct", 0)

        ood_pct = cluster_data.get("ood_pct", 0)
        n_clusters = cluster_data.get("n_clusters", 0)
        uncovered = cluster_data.get("uncovered_clusters", 0)
        
        if coverage_pct >= 95:
            _add_md_paragraph(doc, f"✅ **Excellent Coverage**: Test data covers {coverage_pct:.1f}% of the {n_clusters} training input space clusters.")
        elif coverage_pct >= 80:
            _add_md_paragraph(doc, f"⚠️ **Good Coverage**: Test data covers {coverage_pct:.1f}% of clusters ({uncovered} uncovered).")
        else:
            _add_md_paragraph(doc, f"🚨 **Poor Coverage**: Only {coverage_pct:.1f}% of training clusters are covered by test data.")
        
        if ood_pct > 10:
            _add_md_paragraph(doc, f"⚠️ **OOD Alert**: {ood_pct:.1f}% of test samples appear to be out-of-distribution.")
        
        # Summary table
        doc.add_paragraph("")
        summary_table = doc.add_table(rows=5, cols=2)
        summary_table.style = 'Table Grid'
        summary_data = [
            ("Total Clusters", str(n_clusters)),
            ("Coverage", f"{coverage_pct:.1f}%"),
            ("Covered Clusters", str(cluster_data.get("covered_clusters", 0))),
            ("Uncovered Clusters", str(uncovered)),
            ("OOD Samples", f"{ood_pct:.1f}% ({cluster_data.get('ood_count', 0)} samples)")
        ]
        for i, (label, value) in enumerate(summary_data):
            summary_table.rows[i].cells[0].text = label
            summary_table.rows[i].cells[1].text = value
        
        # Cluster detail table
        cluster_summary = cluster_data.get("cluster_summary", [])
        if cluster_summary:
            doc.add_paragraph("")
            doc.add_paragraph("Cluster Distribution:")
            table = doc.add_table(rows=1, cols=6)
            table.style = 'Table Grid'
            hdr = table.rows[0].cells
            hdr[0].text = "Cluster"
            hdr[1].text = "Train Count"
            hdr[2].text = "Train %"
            hdr[3].text = "Test Count"
            hdr[4].text = "Test %"
            hdr[5].text = "Status"
            
            for row_d in cluster_summary:
                r = table.add_row().cells
                r[0].text = str(row_d.get("Cluster", ""))
                r[1].text = str(row_d.get("Train Count", ""))
                r[2].text = str(row_d.get("Train %", ""))
                r[3].text = str(row_d.get("Test Count", ""))
                r[4].text = str(row_d.get("Test %", ""))
                r[5].text = str(row_d.get("Status", "")).replace("✓", "(ok)").replace("✗", "(x)")
        
        # Cluster distribution chart
        if "distribution" in cluster_imgs:
            doc.add_paragraph("")
            doc.add_paragraph("Cluster Distribution Chart:")
            doc.add_picture(io.BytesIO(cluster_imgs["distribution"]), width=Inches(5))
        
        # PCA scatter plot
        if "pca_scatter" in cluster_imgs:
            doc.add_paragraph("")
            doc.add_paragraph("Cluster Space Visualization (PCA):")
            doc.add_picture(io.BytesIO(cluster_imgs["pca_scatter"]), width=Inches(5))
    else:
        doc.add_paragraph("Cluster coverage check not run.")

    # Benchmarking
    bench_data = buf.get("benchmark")
    bench_imgs = buf.get("benchmark_images", {})
    doc.add_heading('2.4 Benchmarking', level=2)
    
    if bench_data:
        your_model = bench_data.get("your_model", {})
        baselines = bench_data.get("baselines", {})
        task = bench_data.get("task_type", "classification")
        
        if baselines:
            doc.add_paragraph(f"Comparison of your model against {len(baselines)} baseline model(s) on test data.")
            
            # Build table with all models
            metrics = list(your_model.keys())
            n_cols = 2 + len(baselines)  # Metric + Your Model + baselines
            table = doc.add_table(rows=1, cols=n_cols)
            table.style = 'Table Grid'
            
            # Header
            hdr = table.rows[0].cells
            hdr[0].text = "Metric"
            hdr[1].text = "Your Model"
            for i, baseline_name in enumerate(baselines.keys()):
                hdr[2 + i].text = baseline_name[:20]  # Truncate long names
            
            # Data rows
            for metric in metrics:
                r = table.add_row().cells
                r[0].text = metric.upper()
                r[1].text = f"{your_model.get(metric, 0):.4f}"
                for i, baseline_name in enumerate(baselines.keys()):
                    r[2 + i].text = f"{baselines[baseline_name].get(metric, 0):.4f}"
            
            # Add chart if available
            if "comparison" in bench_imgs:
                doc.add_paragraph("")
                doc.add_paragraph("Performance Comparison Chart:")
                doc.add_picture(io.BytesIO(bench_imgs["comparison"]), width=Inches(6))
        else:
            doc.add_paragraph("No baseline models were compared.")
    else:
        doc.add_paragraph("Benchmarking not run.")

    # 3. Explainability
    doc.add_heading('3. Explainability', level=1)
    expl = buf.get("explainability", {})
    if expl and expl.get("status") == "ok":
        # Narrative
        nar_shap = _story_shap(expl)
        if nar_shap:
            doc.add_heading('Executive Insight', level=2)
            _add_md_paragraph(doc, nar_shap)
            
        plots = expl.get("plots", {})
        
        if "beeswarm" in plots:
            doc.add_heading('SHAP Beeswarm', level=2)
            import os
            p = plots["beeswarm"]
            if isinstance(p, str) and os.path.exists(p): doc.add_picture(p, width=Inches(5))
            elif isinstance(p, bytes): doc.add_picture(io.BytesIO(p), width=Inches(5))
                
        if "bar" in plots:
            doc.add_heading('SHAP Importance', level=2)
            p = plots["bar"]
            if isinstance(p, str) and os.path.exists(p): doc.add_picture(p, width=Inches(5))
            elif isinstance(p, bytes): doc.add_picture(io.BytesIO(p), width=Inches(5))
            
        # Top Features Table
        if "top_features" in expl:
             doc.add_heading('Top Drivers Table', level=2)
             tf = expl["top_features"]
             table = doc.add_table(rows=1, cols=2)
             table.style = 'Table Grid'
             table.rows[0].cells[0].text = "Feature"
             table.rows[0].cells[1].text = "Impact"
             for row in tf[:5]:
                 r = table.add_row().cells
                 k_feat = "feature" if "feature" in row else list(row.keys())[0]
                 k_imp = "importance" if "importance" in row else list(row.keys())[1]
                 r[0].text = str(row.get(k_feat, ""))
                 val = row.get(k_imp, 0)
                 r[1].text = f"{val:.4f}" if isinstance(val, float) else str(val)
    else:
        doc.add_paragraph("Explainability not run.")

    # Appendix: Glossary
    doc.add_page_break()
    doc.add_heading('Appendix: Guide to Metrics', level=1)
    
    # Performance Metrics
    if task_type == "classification":
        cls_metrics = ["ROC AUC", "PR AUC", "F1 Score", "Accuracy", "Precision", "Recall", "Log Loss", "Brier Score", "MCC", "KS Statistic"]
        for m in cls_metrics:
            if m in GLOSSARY:
                doc.add_paragraph(f"**{m}**: {GLOSSARY[m]}", style='List Bullet')
    else:
        reg_metrics = ["RMSE", "MAE", "R2 Score", "Median AE"]
        for m in reg_metrics:
            if m in GLOSSARY:
                doc.add_paragraph(f"**{m}**: {GLOSSARY[m]}", style='List Bullet')
    
    # Validation Concepts
    doc.add_paragraph(f"**PSI**: {GLOSSARY['PSI']}", style='List Bullet')
    doc.add_paragraph(f"**Stress Test**: {GLOSSARY['Stress Test']}", style='List Bullet')
    doc.add_paragraph(f"**SHAP**: {GLOSSARY['SHAP']}", style='List Bullet')
    doc.add_paragraph(f"**Cluster Coverage**: {GLOSSARY['Cluster Coverage']}", style='List Bullet')



    # Save
    buffer = io.BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    return buffer


def _generate_ranking_report_docx(metrics_df, corr_df, method, target, task_type, X=None, y=None):
    doc = Document()
    import matplotlib.pyplot as plt
    import seaborn as sns
    doc.add_heading('Feature Power Ranking Report', 0)
    
    doc.add_paragraph(f"Generated by TanML on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    # 1. Executive Summary
    doc.add_heading('1. Executive Summary', level=1)
    doc.add_paragraph(f"Target Variable: {target}")
    doc.add_paragraph(f"Task Type: {task_type}")
    doc.add_paragraph(f"Ranking Method: {method}")
    
    
    
    # Automated Narrative (Removed as per user request)
    # narrative = _story_features(metrics_df)
    # if narrative:
    #     doc.add_heading('Key Insights', level=2)
    #     doc.add_paragraph(narrative)
    
    doc.add_heading('Top Influential Features:', level=2)
    top_5 = metrics_df.head(5)
    for _, row in top_5.iterrows():
        doc.add_paragraph(f"{row['Feature']}: Power Score {row['Power']:.1f}", style='List Bullet')

    # 2. Detailed Metrics
    doc.add_heading('2. Feature Importance Metrics', level=1)
    
    # Convert DataFrame to Table
    t = doc.add_table(rows=1, cols=len(metrics_df.columns))
    t.style = 'Table Grid'
    
    # Header
    hdr_cells = t.rows[0].cells
    for i, col_name in enumerate(metrics_df.columns):
        hdr_cells[i].text = str(col_name)
        
    # Body
    for _, row in metrics_df.iterrows():
        row_cells = t.add_row().cells
        for i, val in enumerate(row):
            if isinstance(val, float):
                row_cells[i].text = f"{val:.3f}"
            else:
                row_cells[i].text = str(val)

    # 3. Visuals (Distribution Overlay)
    if X is not None and y is not None:
        doc.add_heading('3. Key Feature Distributions', level=1)
        doc.add_paragraph(f"Distribution overlay for top {len(top_5)} features relative to target '{target}'.")
        
        # Prepare plot style
        sns.set_theme(style="whitegrid")
        
        for idx, row in top_5.iterrows():
            feat = row['Feature']
            if feat not in X.columns: continue
            
            try:
                fig, ax = plt.subplots(figsize=(6, 4))
                
                if task_type == 'classification':
                    # KDE Plot with Hue
                    sns.kdeplot(data=X, x=feat, hue=y, fill=True, common_norm=False, palette="tab10", ax=ax)
                    ax.set_title(f"{feat} distribution by {target}")
                else:
                    # Scatter Plot for Regression
                    sns.scatterplot(x=X[feat], y=y, alpha=0.6, ax=ax)
                    ax.set_title(f"{feat} vs {target}")
                    ax.set_ylabel(target)
                
                plt.tight_layout()
                
                # Save to buffer
                img_buf = io.BytesIO()
                fig.savefig(img_buf, format='png', dpi=100)
                plt.close(fig)
                img_buf.seek(0)
                
                doc.add_heading(f"Feature: {feat}", level=3)
                doc.add_picture(img_buf, width=Inches(5.0))
                
            except Exception as e:
                doc.add_paragraph(f"Could not plot {feat}: {e}")

    # 4. Correlation Matrix
    if corr_df is not None and not corr_df.empty:
        doc.add_heading('4. Correlation Matrix', level=1)
        doc.add_paragraph("Pairwise correlation of numeric features (Pearson).")
        
        # Reset index to include feature names in table
        c_disp = corr_df.reset_index().rename(columns={"index": "Feature"})
        
        # Colormap setup (Use Data Min/Max to match Pandas 'background_gradient' defaults)
        cmap = plt.get_cmap("coolwarm")
        # Pandas default: vmin/vmax from data, unless specified
        d_min = corr_df.min().min()
        d_max = corr_df.max().max()
        norm = mcolors.Normalize(vmin=d_min, vmax=d_max)
        
        t2 = doc.add_table(rows=1, cols=len(c_disp.columns))
        t2.style = 'Table Grid'
        
        # Header
        hdr2 = t2.rows[0].cells
        for i, col_name in enumerate(c_disp.columns):
            hdr2[i].text = str(col_name)
            
        # Body
        for _, row in c_disp.iterrows():
            row_cells = t2.add_row().cells
            for i, val in enumerate(row):
                # Feature Name (Col 0)
                if i == 0:
                    row_cells[i].text = str(val)
                    continue

                # Correlation Value
                if isinstance(val, (int, float)):
                    row_cells[i].text = f"{val:.2f}"
                    
                    # Apply background gradient
                    try:
                        # Get RGBA from colormap
                        rgba = cmap(norm(val))
                        # Convert to Hex (ignore alpha)
                        hex_color = mcolors.to_hex(rgba, keep_alpha=False).lstrip('#')
                        
                        # XML Magic for shading
                        shading_elm = parse_xml(r'<w:shd {} w:fill="{}"/>'.format(nsdecls('w'), hex_color))
                        row_cells[i]._tc.get_or_add_tcPr().append(shading_elm)

                        # Text Contrast (White text on dark background)
                        r, g, b, _ = rgba
                        lum = 0.299*r + 0.587*g + 0.114*b
                        if lum < 0.5:
                           run = row_cells[i].paragraphs[0].runs[0]
                           run.font.color.rgb = RGBColor(255, 255, 255)
                           
                    except Exception:
                        pass
                else:
                    row_cells[i].text = str(val)

    # Appendix
    doc.add_page_break()
    doc.add_heading('Appendix: Guide to Metrics', level=1)
    _add_md_paragraph(doc, f"**Power Score**: {GLOSSARY.get('Power Score', 'Relative importance of a feature.')}", style='List Bullet')



    # Save
    buffer = io.BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    return buffer

