from __future__ import annotations

import math
import tempfile
from pathlib import Path
from typing import Any

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import cm
from reportlab.platypus import Image, Paragraph, SimpleDocTemplate, Spacer, Table, TableStyle


def _safe_float(value: Any) -> float | None:
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    return number if math.isfinite(number) else None


def _prediction_columns(frame: pd.DataFrame) -> tuple[str, str] | None:
    for true_col, pred_col in (
        ("y_true_vunit", "y_pred_vunit"),
        ("y_true", "y_pred"),
    ):
        if {true_col, pred_col}.issubset(frame.columns):
            return true_col, pred_col
    return None


def compute_metrics(frame: pd.DataFrame) -> dict[str, float] | None:
    cols = _prediction_columns(frame)
    if cols is None or frame.empty:
        return None

    true_col, pred_col = cols
    work = frame[[true_col, pred_col]].apply(pd.to_numeric, errors="coerce").dropna()
    if work.empty:
        return None

    y_true = work[true_col].to_numpy(dtype=float)
    y_pred = work[pred_col].to_numpy(dtype=float)
    residual = y_pred - y_true
    ratio = np.divide(y_pred, y_true, out=np.full_like(y_pred, np.nan), where=y_true != 0)
    ratio = ratio[np.isfinite(ratio)]

    positive_mask = (y_true > 0) & (y_pred > 0)
    if positive_mask.any():
        log_true = np.log(y_true[positive_mask])
        log_pred = np.log(y_pred[positive_mask])
        log_residual = log_pred - log_true
        log_ss_tot = float(np.sum((log_true - np.mean(log_true)) ** 2))
        r2_log = float(1.0 - np.sum(log_residual**2) / log_ss_tot) if log_ss_tot > 0 else float("nan")
    else:
        r2_log = float("nan")

    mape_mask = y_true != 0
    mape = float(np.mean(np.abs(residual[mape_mask] / y_true[mape_mask])) * 100.0) if mape_mask.any() else float("nan")

    out = {
        "n_obs": float(len(work)),
        "R2": r2_log,
        "RMSE": float(np.sqrt(np.mean(residual**2))),
        "MAE": float(np.mean(np.abs(residual))),
        "MAPE": mape,
    }
    if ratio.size:
        median_ratio = float(np.median(ratio))
        mean_ratio = float(np.mean(ratio))
        out.update(
            {
                "COD": float(np.mean(np.abs(ratio - median_ratio)) / median_ratio * 100.0) if median_ratio != 0 else float("nan"),
                "PRD": float(mean_ratio / (float(np.sum(y_pred)) / float(np.sum(y_true)))) if float(np.sum(y_true)) != 0 else float("nan"),
                "Mediana": median_ratio,
            }
        )
    return out


def split_prediction_frames(holdout_predictions: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    if holdout_predictions.empty:
        return pd.DataFrame(), pd.DataFrame()
    if "split" not in holdout_predictions.columns:
        return pd.DataFrame(), holdout_predictions.copy()
    split = holdout_predictions["split"].astype("string").str.lower()
    train = holdout_predictions.loc[split.eq("train")].copy()
    test = holdout_predictions.loc[split.eq("test")].copy()
    return train, test


def build_metric_summary(
    holdout_predictions: pd.DataFrame,
    metadata: dict[str, Any],
) -> dict[str, dict[str, float] | None]:
    train_frame, test_frame = split_prediction_frames(holdout_predictions)
    train_metrics = compute_metrics(train_frame)
    test_metrics = compute_metrics(test_frame)
    if test_metrics is None:
        metadata_holdout = metadata.get("holdout_metrics")
        if isinstance(metadata_holdout, dict):
            test_metrics = {
                key: float(value)
                for key, value in metadata_holdout.items()
                if key != "R2" and _safe_float(value) is not None
            }
    return {"train": train_metrics, "test": test_metrics}


def _fmt(value: Any, digits: int = 4) -> str:
    number = _safe_float(value)
    if number is None:
        return "n/d"
    return f"{number:.{digits}f}"


def make_scatter_plot(frame: pd.DataFrame, title: str, output_path: Path) -> Path:
    cols = _prediction_columns(frame)
    fig, ax = plt.subplots(figsize=(5.6, 4.0), dpi=160)
    if cols is None or frame.empty:
        ax.axis("off")
        ax.text(
            0.5,
            0.5,
            "Predicoes nao disponiveis\nnos artefatos exportados",
            ha="center",
            va="center",
            fontsize=11,
        )
    else:
        true_col, pred_col = cols
        work = frame[[true_col, pred_col]].apply(pd.to_numeric, errors="coerce").dropna()
        work = work.loc[(work[true_col] > 0) & (work[pred_col] > 0)].copy()
        if work.empty:
            ax.axis("off")
            ax.text(0.5, 0.5, "Sem observacoes validas", ha="center", va="center", fontsize=11)
        else:
            x = work[true_col].to_numpy(dtype=float)
            y = work[pred_col].to_numpy(dtype=float)
            lim_min = float(min(np.min(x), np.min(y)))
            lim_max = float(max(np.max(x), np.max(y)))
            ax.scatter(x, y, s=14, alpha=0.45, color="#2563EB", edgecolors="none")
            ax.plot([lim_min, lim_max], [lim_min, lim_max], color="#DC2626", linewidth=1.2)
            ax.set_xscale("log")
            ax.set_yscale("log")
            ax.set_xlabel("Valor observado (escala log)")
            ax.set_ylabel("Valor inferido (escala log)")
            ax.grid(alpha=0.25, which="both")
    ax.set_title(title, fontsize=11)
    fig.tight_layout()
    fig.savefig(output_path, bbox_inches="tight")
    plt.close(fig)
    return output_path


def _table(data: list[list[Any]], widths: list[float] | None = None) -> Table:
    table = Table(data, colWidths=widths)
    table.setStyle(
        TableStyle(
            [
                ("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#1F2937")),
                ("TEXTCOLOR", (0, 0), (-1, 0), colors.white),
                ("FONTNAME", (0, 0), (-1, 0), "Helvetica-Bold"),
                ("GRID", (0, 0), (-1, -1), 0.25, colors.HexColor("#D1D5DB")),
                ("VALIGN", (0, 0), (-1, -1), "TOP"),
                ("FONTNAME", (0, 1), (-1, -1), "Helvetica"),
                ("FONTSIZE", (0, 0), (-1, -1), 8),
                ("BACKGROUND", (0, 1), (-1, -1), colors.HexColor("#F9FAFB")),
            ]
        )
    )
    return table


def build_pdf_report(
    *,
    model_label: str,
    model_file: str,
    user_inputs: dict[str, Any],
    prepared_values: dict[str, Any],
    source_map: dict[str, str],
    prediction: dict[str, Any],
    metrics: dict[str, dict[str, float] | None],
    train_predictions: pd.DataFrame,
    test_predictions: pd.DataFrame,
    limitations: list[str],
) -> str:
    report_dir = Path(tempfile.mkdtemp(prefix="avm_report_"))
    train_plot = make_scatter_plot(train_predictions, "Treino: observado x inferido", report_dir / "scatter_train.png")
    test_plot = make_scatter_plot(test_predictions, "Teste: observado x inferido", report_dir / "scatter_test.png")
    pdf_path = report_dir / "relatorio_inferencia_avm.pdf"

    styles = getSampleStyleSheet()
    doc = SimpleDocTemplate(
        str(pdf_path),
        pagesize=A4,
        rightMargin=1.4 * cm,
        leftMargin=1.4 * cm,
        topMargin=1.2 * cm,
        bottomMargin=1.2 * cm,
    )
    story: list[Any] = [
        Paragraph("Relatorio de Inferencia Territorial", styles["Title"]),
        Spacer(1, 0.3 * cm),
        Paragraph(f"Modelo utilizado: <b>{model_label}</b> ({model_file})", styles["BodyText"]),
        Spacer(1, 0.25 * cm),
    ]

    story.append(Paragraph("Dados informados pelo usuario", styles["Heading2"]))
    user_rows = [["Campo", "Valor"]] + [[str(key), str(value)] for key, value in user_inputs.items()]
    story.append(_table(user_rows, [5.0 * cm, 11.0 * cm]))
    story.append(Spacer(1, 0.25 * cm))

    story.append(Paragraph("Predicao", styles["Heading2"]))
    pred_rows = [
        ["Indicador", "Valor"],
        ["Valor unitario estimado", f"R$ {prediction['valor_unitario']:,.2f} / m2"],
        ["Valor total estimado", f"R$ {prediction['valor_total']:,.2f}"],
        ["Modo da saida do modelo", str(prediction["target_mode"])],
    ]
    story.append(_table(pred_rows, [6.0 * cm, 10.0 * cm]))
    story.append(Spacer(1, 0.25 * cm))

    story.append(Paragraph("Variaveis efetivamente usadas na inferencia", styles["Heading2"]))
    feature_rows = [["Variavel", "Valor", "Fonte"]]
    for key in sorted(prepared_values):
        if key.startswith("_"):
            continue
        feature_rows.append([key, str(prepared_values.get(key)), str(source_map.get(key, "calculado/derivado"))])
    story.append(_table(feature_rows, [5.0 * cm, 5.2 * cm, 5.8 * cm]))
    story.append(Spacer(1, 0.25 * cm))

    story.append(Paragraph("Metricas do modelo", styles["Heading2"]))
    metrics_rows = [["Conjunto", "n", "R2 (log)", "RMSE", "MAE", "MAPE", "COD", "PRD", "Mediana"]]
    for split_name, metric in (("Treino", metrics.get("train")), ("Teste", metrics.get("test"))):
        metrics_rows.append(
            [
                split_name,
                _fmt(metric.get("n_obs") if metric else None, 0),
                _fmt(metric.get("R2") if metric else None),
                _fmt(metric.get("RMSE") if metric else None, 2),
                _fmt(metric.get("MAE") if metric else None, 2),
                _fmt(metric.get("MAPE") if metric else None, 2),
                _fmt(metric.get("COD") if metric else None, 2),
                _fmt(metric.get("PRD") if metric else None, 4),
                _fmt(metric.get("Mediana") if metric else None, 4),
            ]
        )
    story.append(_table(metrics_rows))
    story.append(Spacer(1, 0.25 * cm))

    story.append(Paragraph("Dispersao observado x inferido", styles["Heading2"]))
    story.append(Image(str(train_plot), width=7.8 * cm, height=5.6 * cm))
    story.append(Spacer(1, 0.15 * cm))
    story.append(Image(str(test_plot), width=7.8 * cm, height=5.6 * cm))
    story.append(Spacer(1, 0.25 * cm))

    story.append(Paragraph("Limitacoes e observacoes", styles["Heading2"]))
    for item in limitations:
        story.append(Paragraph(f"- {item}", styles["BodyText"]))
    doc.build(story)
    return str(pdf_path)