from __future__ import annotations

import os
from pathlib import Path
from typing import Any

MPLCONFIGDIR = Path(__file__).resolve().parents[2] / "outputs" / "logs" / "matplotlib"
MPLCONFIGDIR.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("MPLCONFIGDIR", str(MPLCONFIGDIR))

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib import patches
import numpy as np

from preprocessing.image_io import load_rgb_image
from utils.project import ensure_dir


DPI = 220
BLAST_COLOR = "#d62728"
NON_BLAST_COLOR = "#1f77b4"
TEXT_BOX_EDGE = "#c7c7c7"


def _short_text(value: Any, max_length: int = 26) -> str:
    text = str(value)
    if len(text) <= max_length:
        return text
    head = max(6, max_length // 2 - 2)
    tail = max(6, max_length - head - 3)
    return f"{text[:head]}...{text[-tail:]}"


def _binary_label(value: int | None) -> str:
    if value is None:
        return "n/a"
    return "blast" if int(value) == 1 else "non-blast"


def _save(fig: plt.Figure, output_path: Path) -> str:
    ensure_dir(output_path.parent)
    fig.tight_layout()
    fig.savefig(output_path, dpi=DPI, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    return str(output_path)


def _add_panel_badge(ax: plt.Axes, label: str) -> None:
    ax.text(
        0.02,
        0.98,
        label,
        transform=ax.transAxes,
        ha="left",
        va="top",
        fontsize=12,
        fontweight="bold",
        color="black",
        bbox={"facecolor": "white", "edgecolor": "black", "pad": 2.5},
    )


def _draw_text_box(ax: plt.Axes, title: str, lines: list[str]) -> None:
    ax.axis("off")
    ax.set_facecolor("white")
    ax.add_patch(
        patches.FancyBboxPatch(
            (0.0, 0.0),
            1.0,
            1.0,
            boxstyle="round,pad=0.02,rounding_size=0.02",
            linewidth=1.0,
            edgecolor=TEXT_BOX_EDGE,
            facecolor="white",
            transform=ax.transAxes,
        )
    )
    ax.text(0.04, 0.96, title, transform=ax.transAxes, ha="left", va="top", fontsize=10, fontweight="bold")
    ax.text(0.04, 0.88, "\n".join(lines), transform=ax.transAxes, ha="left", va="top", fontsize=8.5)


def _draw_image_panel(ax: plt.Axes, image_path: str, title: str = "", subtitle: str | None = None) -> None:
    image = load_rgb_image(image_path)
    ax.imshow(image)
    ax.axis("off")
    if title:
        ax.set_title(title, fontsize=10)
    if subtitle:
        ax.text(0.5, -0.08, subtitle, transform=ax.transAxes, ha="center", va="top", fontsize=8)


def _draw_roi_panel(
    ax: plt.Axes,
    row: dict[str, Any],
    *,
    image_column: str = "bbox_crop_path",
    title: str = "",
    subtitle: str | None = None,
) -> None:
    image = load_rgb_image(str(row[image_column]))
    ax.imshow(image)
    ax.axis("off")
    if title:
        ax.set_title(title, fontsize=9)
    if subtitle:
        ax.text(0.5, -0.08, subtitle, transform=ax.transAxes, ha="center", va="top", fontsize=7.5)


def _draw_smear_with_boxes(
    ax: plt.Axes,
    image_path: str,
    object_rows: list[dict[str, Any]],
    *,
    title: str,
    labels: list[str] | None = None,
    show_scores: bool = False,
) -> None:
    image = load_rgb_image(image_path)
    ax.imshow(image)
    ax.axis("off")
    ax.set_title(title, fontsize=10)
    for index, row in enumerate(object_rows, start=1):
        color = BLAST_COLOR if int(row.get("target_binary", row.get("y_true", 0))) == 1 else NON_BLAST_COLOR
        x1 = float(row["bbox_x1"])
        y1 = float(row["bbox_y1"])
        width = float(row["bbox_width_px"])
        height = float(row["bbox_height_px"])
        ax.add_patch(
            patches.Rectangle((x1, y1), width, height, linewidth=1.6, edgecolor=color, facecolor="none")
        )
        label = labels[index - 1] if labels and index - 1 < len(labels) else f"#{index}"
        if show_scores and "y_score" in row:
            label = f"{label} {float(row['y_score']):.2f}"
        ax.text(
            x1,
            max(2.0, y1 - 4.0),
            label,
            fontsize=7,
            color="white",
            bbox={"facecolor": color, "edgecolor": color, "pad": 1.4},
        )


def _score_lines(method_rows: list[dict[str, Any]]) -> list[str]:
    lines: list[str] = []
    for row in method_rows:
        lines.append(f"{row['aggregation_method']}: {float(row['score']):.3f}")
    return lines


def _render_a01_f1(spec: dict[str, Any], output_path: Path) -> str:
    class_rows = spec["class_rows"]
    fig, axes = plt.subplots(len(class_rows), 4, figsize=(11.8, 15.2))
    if len(class_rows) == 1:
        axes = np.asarray([axes])
    for row_index, group in enumerate(class_rows):
        for col_index in range(4):
            row = group["rows"][col_index]
            title = group["class_name"] if col_index == 0 else ""
            _draw_roi_panel(axes[row_index, col_index], row, title=title, subtitle=str(row["class_name"]))
    fig.suptitle(spec["title"], fontsize=13)
    return _save(fig, output_path)


def _render_a01_f2(spec: dict[str, Any], output_path: Path) -> str:
    object_rows = spec["object_rows"]
    fig = plt.figure(figsize=(13.2, 10.6))
    grid = fig.add_gridspec(3, 4, height_ratios=[1.2, 1.0, 1.0])

    ax_a = fig.add_subplot(grid[0, :2])
    ax_b = fig.add_subplot(grid[0, 2:])
    _draw_image_panel(
        ax_a,
        spec["image_row"]["image_path"],
        title=f"Smear image | patient {_short_text(spec['image_row']['patient_id'], 18)}",
    )
    _add_panel_badge(ax_a, "A")
    _draw_smear_with_boxes(
        ax_b,
        spec["image_row"]["image_path"],
        object_rows,
        title=f"Selected bounding boxes | image {_short_text(spec['image_row']['image_id'], 18)}",
        labels=[str(index) for index in range(1, len(object_rows) + 1)],
    )
    _add_panel_badge(ax_b, "B")

    roi_positions = [(1, 0), (1, 1), (1, 2), (1, 3), (2, 0), (2, 1)]
    for slot, (row_index, col_index) in enumerate(roi_positions, start=1):
        ax = fig.add_subplot(grid[row_index, col_index])
        if slot <= len(object_rows):
            row = object_rows[slot - 1]
            _draw_roi_panel(ax, row, title=f"#{slot} {_short_text(row['class_name'], 18)}")
            if slot == 1:
                _add_panel_badge(ax, "C")
        else:
            ax.axis("off")
    for col_index in [2, 3]:
        ax = fig.add_subplot(grid[2, col_index])
        ax.axis("off")

    fig.suptitle(spec["title"], fontsize=13)
    return _save(fig, output_path)


def _render_a01_f3(spec: dict[str, Any], output_path: Path) -> str:
    object_rows = spec["object_rows"]
    fig, axes = plt.subplots(len(object_rows), 3, figsize=(12.4, 15.6))
    if len(object_rows) == 1:
        axes = np.asarray([axes])
    column_titles = ["bbox crop", "mask crop", "mask-tight crop"]
    for column_index, column_title in enumerate(column_titles):
        axes[0, column_index].set_title(column_title, fontsize=10)
    for row_index, row in enumerate(object_rows):
        score_line = (
            f"bbox={float(row['bbox_only_y_score']):.3f} | "
            f"mask={float(row['mask_only_y_score']):.3f} | "
            f"combined={float(row['bbox_plus_mask_y_score']):.3f}"
        )
        _draw_roi_panel(
            axes[row_index, 0],
            row,
            image_column="bbox_crop_path",
            title=f"{_short_text(row['class_name'], 20)} | {_binary_label(int(row['target_binary']))}",
            subtitle=score_line,
        )
        _draw_roi_panel(axes[row_index, 1], row, image_column="mask_crop_path")
        _draw_roi_panel(axes[row_index, 2], row, image_column="mask_tight_crop_path")
    fig.suptitle(spec["title"], fontsize=13)
    return _save(fig, output_path)


def _render_a01_f4(spec: dict[str, Any], output_path: Path) -> str:
    grid_rows = spec["grid_rows"]
    fig, axes = plt.subplots(len(grid_rows), 3, figsize=(12.2, 13.8))
    if len(grid_rows) == 1:
        axes = np.asarray([axes])
    column_titles = [
        "clean_cell references",
        "objects removed by strict_morphology",
        "full_realistic difficult cases",
    ]
    for column_index, column_title in enumerate(column_titles):
        axes[0, column_index].set_title(column_title, fontsize=10)
    for row_index, row_group in enumerate(grid_rows):
        for column_index, key in enumerate(["clean", "strict_removed", "full_problem"]):
            row = row_group[key]
            _draw_roi_panel(
                axes[row_index, column_index],
                row,
                title=_short_text(row["class_name"], 22),
                subtitle=_binary_label(int(row["target_binary"])) if row.get("target_binary") is not None else None,
            )
    fig.suptitle(spec["title"], fontsize=13)
    return _save(fig, output_path)


def _render_a01_f5(spec: dict[str, Any], output_path: Path) -> str:
    outcome_rows = spec["outcome_rows"]
    fig, axes = plt.subplots(4, 4, figsize=(12.8, 12.8))
    for row_index, outcome_label in enumerate(["TP", "TN", "FP", "FN"]):
        rows = outcome_rows[outcome_label]
        for col_index in range(4):
            ax = axes[row_index, col_index]
            row = rows[col_index]
            subtitle = (
                f"true={row['true_label']} | pred={row['predicted_label']}\n"
                f"score={float(row['y_score']):.3f}"
            )
            title = f"{outcome_label} | {_short_text(row['class_name'], 18)}" if col_index == 0 else _short_text(row["class_name"], 18)
            _draw_roi_panel(ax, row, title=title, subtitle=subtitle)
    fig.suptitle(spec["title"], fontsize=13)
    return _save(fig, output_path)


def _render_a01_f6(spec: dict[str, Any], output_path: Path) -> str:
    confident_rows = spec["confident_rows"]
    disagreement_rows = spec["disagreement_rows"]
    fig, axes = plt.subplots(4, 2, figsize=(12.2, 15.6))
    axes[0, 0].set_title("handcrafted confident and correct", fontsize=10)
    axes[0, 1].set_title("model disagreement on the same cell", fontsize=10)
    for row_index in range(4):
        confident = confident_rows[row_index]
        disagreement = disagreement_rows[row_index]
        confident_subtitle = (
            f"HCF={float(confident['handcrafted_y_score']):.3f} | "
            f"BBOX={float(confident['cnn_bbox_y_score']):.3f}\n"
            f"MASK={float(confident['cnn_mask_y_score']):.3f} | "
            f"FUSION={float(confident['cnn_late_fusion_y_score']):.3f}"
        )
        disagreement_subtitle = (
            f"HCF={float(disagreement['handcrafted_y_score']):.3f} | "
            f"BBOX={float(disagreement['cnn_bbox_y_score']):.3f}\n"
            f"MASK={float(disagreement['cnn_mask_y_score']):.3f} | "
            f"FUSION={float(disagreement['cnn_late_fusion_y_score']):.3f}"
        )
        _draw_roi_panel(
            axes[row_index, 0],
            confident,
            title=f"{_short_text(confident['class_name'], 18)} | {_binary_label(int(confident['y_true']))}",
            subtitle=confident_subtitle,
        )
        _draw_roi_panel(
            axes[row_index, 1],
            disagreement,
            title=f"{_short_text(disagreement['class_name'], 18)} | {_binary_label(int(disagreement['y_true']))}",
            subtitle=disagreement_subtitle,
        )
    fig.suptitle(spec["title"], fontsize=13)
    return _save(fig, output_path)


def _render_a02_f1(spec: dict[str, Any], output_path: Path) -> str:
    object_rows = spec["top_object_rows"]
    fig = plt.figure(figsize=(13.6, 10.8))
    grid = fig.add_gridspec(4, 4, height_ratios=[1.15, 1.15, 1.0, 1.0])
    ax_a = fig.add_subplot(grid[:2, :2])
    ax_b = fig.add_subplot(grid[:2, 2:])
    _draw_image_panel(
        ax_a,
        spec["image_row"]["image_path"],
        title=f"Original image | {_short_text(spec['image_row']['patient_id'], 18)}",
        subtitle=_short_text(spec["image_row"]["image_id"], 22),
    )
    _add_panel_badge(ax_a, "A")
    _draw_smear_with_boxes(
        ax_b,
        spec["image_row"]["image_path"],
        object_rows,
        title="Top object-level detections on the same image",
        labels=[f"#{index}" for index in range(1, len(object_rows) + 1)],
        show_scores=True,
    )
    _add_panel_badge(ax_b, "B")
    for slot in range(6):
        row_index = 2 + slot // 3
        col_index = slot % 3
        ax = fig.add_subplot(grid[row_index, col_index])
        row = object_rows[slot]
        _draw_roi_panel(
            ax,
            row,
            title=f"#{slot + 1} {_short_text(row['class_name'], 16)}",
            subtitle=f"score={float(row['y_score']):.3f}",
        )
        if slot == 0:
            _add_panel_badge(ax, "C")
    ax_d = fig.add_subplot(grid[2:, 3])
    _draw_text_box(
        ax_d,
        "Image-level aggregation",
        _score_lines(spec["method_rows"])
        + ["", f"n_objects={int(spec['image_row']['n_objects'])}", f"spread={float(spec['image_row']['method_spread']):.3f}"],
    )
    _add_panel_badge(ax_d, "D")
    fig.suptitle(spec["title"], fontsize=13)
    return _save(fig, output_path)


def _render_a02_f2(spec: dict[str, Any], output_path: Path) -> str:
    cases = spec["cases"]
    fig, axes = plt.subplots(4, 4, figsize=(13.4, 15.6))
    for row_index, case in enumerate(cases):
        image_row = case["image_row"]
        top_object_rows = case["top_object_rows"]
        _draw_image_panel(
            axes[row_index, 0],
            image_row["image_path"],
            title=f"{case['outcome_label']} | score={float(image_row['score']):.3f}",
            subtitle=_short_text(image_row["image_id"], 22),
        )
        for col_index in range(1, 4):
            ax = axes[row_index, col_index]
            if col_index - 1 < len(top_object_rows):
                row = top_object_rows[col_index - 1]
                _draw_roi_panel(
                    ax,
                    row,
                    title=f"#{col_index} {_short_text(row['class_name'], 14)}",
                    subtitle=f"score={float(row['y_score']):.3f}",
                )
            else:
                ax.axis("off")
    fig.suptitle(f"{spec['title']} | best method: {spec['best_method']}", fontsize=13)
    return _save(fig, output_path)


def _render_a02_f3(spec: dict[str, Any], output_path: Path) -> str:
    object_rows = spec["top_object_rows"]
    fig = plt.figure(figsize=(13.2, 10.8))
    grid = fig.add_gridspec(4, 4, height_ratios=[1.15, 1.15, 1.15, 0.95])
    ax_a = fig.add_subplot(grid[:3, :2])
    _draw_smear_with_boxes(
        ax_a,
        spec["image_row"]["image_path"],
        object_rows,
        title=f"Same image under four aggregators | {_short_text(spec['image_row']['image_id'], 22)}",
        labels=[f"#{index}" for index in range(1, len(object_rows) + 1)],
        show_scores=True,
    )
    _add_panel_badge(ax_a, "A")
    positions = [(0, 2), (0, 3), (1, 2), (1, 3)]
    for slot, (row_index, col_index) in enumerate(positions, start=1):
        row = object_rows[slot - 1]
        ax = fig.add_subplot(grid[row_index, col_index])
        _draw_roi_panel(
            ax,
            row,
            title=f"#{slot} {_short_text(row['class_name'], 14)}",
            subtitle=f"score={float(row['y_score']):.3f}",
        )
        if slot == 1:
            _add_panel_badge(ax, "B")
    method_titles = {
        "max_object_score": "max object score",
        "mean_top3_scores": "mean top-3 scores",
        "count_objects_above_threshold": "count > threshold",
        "soft_noisy_or": "soft noisy-or",
    }
    for index, method_row in enumerate(spec["method_rows"]):
        ax = fig.add_subplot(grid[3, index])
        _draw_text_box(
            ax,
            method_titles.get(method_row["aggregation_method"], method_row["aggregation_method"]),
            [f"score = {float(method_row['score']):.3f}"],
        )
        if index == 0:
            _add_panel_badge(ax, "C")
    ax_blank = fig.add_subplot(grid[2, 2:])
    ax_blank.axis("off")
    fig.suptitle(spec["title"], fontsize=13)
    return _save(fig, output_path)


def _render_a02_f4(spec: dict[str, Any], output_path: Path) -> str:
    image_panels = spec["image_panels"]
    n_rows = len(image_panels)
    fig = plt.figure(figsize=(13.4, 3.1 * n_rows + 1.2))
    grid = fig.add_gridspec(n_rows, 4, width_ratios=[1.6, 1.0, 1.0, 1.1])
    summary_ax = fig.add_subplot(grid[:, 3])
    summary_lines = [
        f"patient: {_short_text(spec['patient_row']['patient_id'], 22)}",
        f"image method: {spec['best_image_method']}",
        f"patient method: {spec['best_patient_method']}",
        f"patient score: {float(spec['patient_row']['score']):.3f}",
        f"n_images shown: {len(image_panels)}",
        f"n_objects total: {int(spec['patient_row']['n_objects'])}",
    ]
    _draw_text_box(summary_ax, "Patient aggregation", summary_lines)
    _add_panel_badge(summary_ax, "D")
    for row_index, panel in enumerate(image_panels):
        image_row = panel["image_row"]
        top_object_rows = panel["top_object_rows"]
        ax_image = fig.add_subplot(grid[row_index, 0])
        _draw_image_panel(
            ax_image,
            image_row["image_path"],
            title=f"{_short_text(image_row['image_id'], 18)} | {float(image_row['score']):.3f}",
            subtitle=f"contains blast={int(image_row['image_contains_blast'])}",
        )
        if row_index == 0:
            _add_panel_badge(ax_image, "A")
        for object_slot, col_index in enumerate([1, 2], start=0):
            ax = fig.add_subplot(grid[row_index, col_index])
            if object_slot < len(top_object_rows):
                row = top_object_rows[object_slot]
                _draw_roi_panel(
                    ax,
                    row,
                    title=f"#{object_slot + 1} {_short_text(row['class_name'], 14)}",
                    subtitle=f"score={float(row['y_score']):.3f}",
                )
            else:
                ax.axis("off")
            if row_index == 0 and object_slot == 0:
                _add_panel_badge(ax, "B")
            if row_index == 0 and object_slot == 1:
                _add_panel_badge(ax, "C")
    fig.suptitle(spec["title"], fontsize=13)
    return _save(fig, output_path)


def _render_a02_f5(spec: dict[str, Any], output_path: Path) -> str:
    case_panels = spec["case_panels"]
    max_rows = max(len(case["image_panels"]) for case in case_panels)
    fig, axes = plt.subplots(max_rows, 6, figsize=(14.3, 4.1 * max_rows + 0.8))
    if max_rows == 1:
        axes = np.asarray([axes])

    for section_index, case in enumerate(case_panels):
        patient_row = case["patient_row"]
        start_col = 0 if case["case_label"] == "TP" else 3
        section_title = (
            f"{case['case_label']} patient | {_short_text(patient_row['patient_id'], 20)} | "
            f"{spec['best_patient_method']}={float(patient_row['score']):.3f}"
        )
        axes[0, start_col].set_title(section_title, fontsize=10)
        for row_index in range(max_rows):
            for col_offset in range(3):
                axes[row_index, start_col + col_offset].axis("off")
        for row_index, panel in enumerate(case["image_panels"]):
            image_row = panel["image_row"]
            top_object_rows = panel["top_object_rows"]
            image_ax = axes[row_index, start_col]
            _draw_image_panel(
                image_ax,
                image_row["image_path"],
                title=f"{_short_text(image_row['image_id'], 16)} | {float(image_row['score']):.3f}",
            )
            if row_index == 0:
                _add_panel_badge(image_ax, "A" if case["case_label"] == "TP" else "D")
            for object_slot, col_offset in enumerate([1, 2], start=0):
                ax = axes[row_index, start_col + col_offset]
                if object_slot < len(top_object_rows):
                    row = top_object_rows[object_slot]
                    _draw_roi_panel(
                        ax,
                        row,
                        title=f"#{object_slot + 1} {_short_text(row['class_name'], 13)}",
                        subtitle=f"score={float(row['y_score']):.3f}",
                    )
                else:
                    ax.axis("off")
                if row_index == 0 and object_slot == 0:
                    _add_panel_badge(ax, "B" if case["case_label"] == "TP" else "E")
                if row_index == 0 and object_slot == 1:
                    _add_panel_badge(ax, "C" if case["case_label"] == "TP" else "F")
    fig.suptitle(spec["title"], fontsize=13)
    return _save(fig, output_path)


def render_subject_figure(spec: dict[str, Any], output_path: Path) -> str:
    kind = spec["figure_kind"]
    if kind == "a01_f1_class_gallery":
        return _render_a01_f1(spec, output_path)
    if kind == "a01_f2_smear_to_objects":
        return _render_a01_f2(spec, output_path)
    if kind == "a01_f3_mask_representations":
        return _render_a01_f3(spec, output_path)
    if kind == "a01_f4_dataset_regimes":
        return _render_a01_f4(spec, output_path)
    if kind == "a01_f5_confusion_gallery":
        return _render_a01_f5(spec, output_path)
    if kind == "a01_f6_same_cells_model_comparison":
        return _render_a01_f6(spec, output_path)
    if kind == "a02_f1_cells_to_image_score":
        return _render_a02_f1(spec, output_path)
    if kind == "a02_f2_image_case_studies":
        return _render_a02_f2(spec, output_path)
    if kind == "a02_f3_same_image_four_aggregators":
        return _render_a02_f3(spec, output_path)
    if kind == "a02_f4_images_to_patient_score":
        return _render_a02_f4(spec, output_path)
    if kind == "a02_f5_patient_case_studies":
        return _render_a02_f5(spec, output_path)
    raise ValueError(f"Unsupported subject figure kind: {kind}")