# -*- coding: utf-8 -*-
"""TEASUR-BMC Software2.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1rCveSa_rJXjtxsA2-Mn2i56ftsdceUYC
"""

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
TEASUR Reproducible Bundle v1.1  (English version)
--------------------------------------------------
Purpose:
    End-to-end, reviewer-friendly pipeline that generates figures, CSV/JSON
    outputs, environment metadata, license/citation templates and optional SHAP
    explainability artifacts. Designed for BMC "Software article" and preprint.

Run (synthetic data):
    python teasur_reproducible.py --outdir out --no-shap --no-pdf

Run (with real CSV):
    python teasur_reproducible.py --csv patients.csv --outdir out --seed 42

Notebook (Colab/Jupyter):
    from teasur_reproducible import run_in_notebook
    run_in_notebook(outdir='out', no_shap=True, no_pdf=True, show_inline=True)

Main outputs in out/ :
    - Figures: fig_its_country.png, fig_mc_hist.png, fig_box_methods.png
    - SHAP (optional): shap_summary.png, shap_bar.png, shap_importance.csv
    - Results: Informe_TEASUR.json, cohorte_teasur.csv
    - Metadata/templates: environment.json, requirements.txt, README.md,
      LICENSE-ACADEMIC-NC.txt, CITATION.cff, MANIFEST_TEASUR.txt
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import textwrap
from pathlib import Path
from typing import Dict, Tuple, List

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.utils import resample

# Optional deps (handled gracefully)
try:
    import shap  # noqa: F401
except Exception:
    shap = None  # will be checked later

try:
    from fpdf import FPDF  # noqa: F401
except Exception:
    FPDF = None

__version__ = "1.1.0"

FEATURE_DIMS: List[str] = [
    "fisica",
    "psicologica",
    "social",
    "espiritual",
    "soporte_familiar",
    "nivel_educativo",
    "componente_psiquico",
]

COUNTRIES = ["Colombia", "Perú", "México", "Chile", "Brasil"]
GENDERS = ["masculino", "femenino"]


# ------------------------- Utilities -------------------------
def set_seed(seed: int = 42) -> None:
    np.random.seed(seed)


def ensure_outdir(path: str | Path) -> Path:
    p = Path(path)
    p.mkdir(parents=True, exist_ok=True)
    return p


def is_notebook() -> bool:
    try:
        from IPython import get_ipython  # type: ignore
        return get_ipython() is not None
    except Exception:
        return False


def display_inline_images(outdir: Path, names: List[str]) -> None:
    """Display images inline in notebooks, if available."""
    if not is_notebook():
        return
    try:
        from IPython.display import Image, display  # type: ignore
        for n in names:
            p = outdir / n
            if p.exists():
                display(Image(filename=str(p)))
    except Exception:
        pass


# ------------------------- Data: load / simulate -------------------------
def simulate_df(n: int = 300, seed: int = 42) -> pd.DataFrame:
    set_seed(seed)
    df = pd.DataFrame({
        "nombre": [f"Paciente_{i+1}" for i in range(n)],
        "edad": np.random.randint(20, 90, n),
        "pais": np.random.choice(COUNTRIES, size=n),
        "genero": np.random.choice(GENDERS, size=n),
        "acceso": np.random.choice([True, False], size=n, p=[0.6, 0.4]),
        "voluntad_anticipada": np.random.choice([True, False], size=n, p=[0.3, 0.7]),
        "fisica": np.random.uniform(0, 10, n),
        "psicologica": np.random.uniform(0, 10, n),
        "social": np.random.uniform(0, 10, n),
        "espiritual": np.random.uniform(0, 10, n),
        "soporte_familiar": np.random.uniform(0, 10, n),
        "nivel_educativo": np.random.uniform(0, 10, n),
        "componente_psiquico": np.random.uniform(0, 10, n),
    })
    return df


def load_or_simulate(csv_path: str | None, n: int, seed: int) -> pd.DataFrame:
    if csv_path and Path(csv_path).exists():
        df = pd.read_csv(csv_path)
        missing = [c for c in FEATURE_DIMS if c not in df.columns]
        if missing:
            raise ValueError(f"CSV is missing required columns: {missing}")
        # Fill optional basics if missing
        for c in ["nombre", "edad", "pais", "genero", "acceso", "voluntad_anticipada"]:
            if c not in df.columns:
                if c == "nombre":
                    df["nombre"] = [f"Registro_{i+1}" for i in range(len(df))]
                elif c == "edad":
                    df["edad"] = np.random.randint(20, 90, len(df))
                elif c == "pais":
                    df["pais"] = np.random.choice(COUNTRIES, size=len(df))
                elif c == "genero":
                    df["genero"] = np.random.choice(GENDERS, size=len(df))
                elif c == "acceso":
                    df["acceso"] = np.random.choice([True, False], size=len(df), p=[0.6, 0.4])
                elif c == "voluntad_anticipada":
                    df["voluntad_anticipada"] = np.random.choice([True, False], size=len(df), p=[0.3, 0.7])
        return df
    else:
        return simulate_df(n=n, seed=seed)


# ------------------------- Core TEASUR -------------------------
def simulate_outcome(df: pd.DataFrame) -> pd.Series:
    """Synthetic 'need' outcome for calibration of weights."""
    true_w = np.array([0.15, 0.20, 0.15, 0.20, 0.10, 0.10, 0.10])
    X = df[FEATURE_DIMS].values
    y = X.dot(true_w) + np.random.normal(0, 0.5, len(df))
    return (y - y.min()) / (y.max() - y.min())


def calibrate_weights(df: pd.DataFrame, outcome_col: str) -> Tuple[Dict[str, float], LinearRegression]:
    X = df[FEATURE_DIMS].values
    y = df[outcome_col].values
    model = LinearRegression(fit_intercept=False)
    model.fit(X, y)
    coef = np.clip(model.coef_, 0, None)
    coef /= coef.sum()
    weights = dict(zip(FEATURE_DIMS, np.round(coef, 4)))
    return weights, model


def compute_ITS(df: pd.DataFrame, weights: Dict[str, float]) -> pd.Series:
    return df[FEATURE_DIMS].multiply(pd.Series(weights)).sum(axis=1)


def adjust_weights(row: pd.Series, base_w: Dict[str, float]) -> Dict[str, float]:
    """Contextual fairness-aware adjustments (illustrative)."""
    w = base_w.copy()
    # Barrier 1: lack of access
    if not row["acceso"]:
        d = 0.1 * w["social"]
        w["social"] += d
        for k in w:
            if k != "social":
                w[k] -= d / (len(w) - 1)
    # Barrier 2: gender
    if str(row["genero"]).lower() == "femenino":
        for dim in ["psicologica", "espiritual"]:
            w[dim] += 0.05
        dec = 0.05 * 2 / (len(w) - 2)
        for k in w:
            if k not in ["psicologica", "espiritual"]:
                w[k] -= dec
    # Barrier 3: no advance directives
    if not row["voluntad_anticipada"]:
        d = 0.05
        w["espiritual"] += d
        for k in w:
            if k != "espiritual":
                w[k] -= d / (len(w) - 1)
    total = sum(w.values())
    return {k: max(0.0, v / total) for k, v in w.items()}


def compute_ITS_barriers(df: pd.DataFrame, base_w: Dict[str, float]) -> pd.Series:
    return df.apply(lambda r: sum(r[d] * adjust_weights(r, base_w)[d] for d in FEATURE_DIMS), axis=1)


def generate_recommendation(row: pd.Series) -> str:
    rec = []
    rec.append("Facilitate access to palliative care." if not row["acceso"] else "Adequate access.")
    rec.append("Encourage advance care planning." if not row["voluntad_anticipada"] else "Advance directives documented.")
    its = row["ITS_barrier"]
    if its > 0.75:
        rec.append("Urgent palliative intervention.")
    elif its > 0.5:
        rec.append("Intense psychosocial support.")
    else:
        rec.append("Standard follow-up.")
    return " ".join(rec)


# ------------------------- Stats / Figures -------------------------
def bootstrap_country_ci(df: pd.DataFrame, B: int = 1000) -> Tuple[List[str], np.ndarray, np.ndarray, np.ndarray]:
    countries = sorted(df["pais"].unique())
    means, lo, hi = [], [], []
    for c in countries:
        vals = df[df["pais"] == c]["ITS_barrier"].values
        if len(vals) == 0:
            means.append(np.nan); lo.append(np.nan); hi.append(np.nan); continue
        boots = [resample(vals, replace=True).mean() for _ in range(B)]
        means.append(vals.mean())
        lo.append(np.percentile(boots, 2.5))
        hi.append(np.percentile(boots, 97.5))
    return countries, np.array(means), np.array(lo), np.array(hi)


def mc_sensitivity(df: pd.DataFrame, base_weights: Dict[str, float], R: int = 500) -> List[float]:
    W = np.array(list(base_weights.values()))
    out = []
    for _ in range(R):
        pert = W + np.random.normal(0, 0.02, len(W))
        pert = np.clip(pert, 0, None)
        pert /= pert.sum()
        out.append((df[FEATURE_DIMS] * pert).sum(axis=1).mean())
    return out


def save_fig(path: Path) -> None:
    plt.tight_layout()
    plt.savefig(path, dpi=300, bbox_inches="tight")
    plt.close()


def make_figures(outdir: Path, countries, means, lo, hi, mc_vals: List[float], df: pd.DataFrame) -> None:
    # 1) ITS by country with 95% CI
    plt.figure(figsize=(8, 5))
    x = np.arange(len(countries))
    plt.bar(x, means, yerr=[means - lo, hi - means], capsize=5)
    plt.xticks(x, countries)
    plt.ylabel("Adjusted ITS")
    plt.title("Adjusted ITS by Country (95% CI)")
    save_fig(outdir / "fig_its_country.png")

    # 2) Monte Carlo histogram
    plt.figure(figsize=(6, 4))
    plt.hist(mc_vals, bins=20)
    plt.xlabel("Mean ITS")
    plt.title("Monte Carlo Sensitivity")
    save_fig(outdir / "fig_mc_hist.png")

    # 3) Boxplot comparison (FACIT-Sp / ESAS-R proxies)
    df["ITS_FACIT"] = df["espiritual"] / 10.0
    df["ITS_ESAS"] = df[["fisica", "psicologica"]].mean(axis=1) / 10.0
    plt.figure(figsize=(7, 4))
    data = [df["ITS_base"], df["ITS_barrier"], df["ITS_FACIT"], df["ITS_ESAS"]]
    bp = plt.boxplot(data, patch_artist=True)  # avoid deprecated labels param
    # cross-version tick labels
    labels = ["Base", "With Barriers", "FACIT-Sp", "ESAS-R"]
    plt.xticks(np.arange(1, len(labels) + 1), labels)
    plt.ylabel("ITS")
    plt.title("Method Comparison")
    save_fig(outdir / "fig_box_methods.png")


# ------------------------- Exports -------------------------
def export_json_pdf(outdir: Path, df: pd.DataFrame, no_pdf: bool = False) -> None:
    report = df[["nombre", "edad", "pais", "ITS_base", "ITS_barrier", "recomendaciones"]].to_dict(orient="records")
    with open(outdir / "Informe_TEASUR.json", "w", encoding="utf-8") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)

    if (not no_pdf) and (FPDF is not None):
        pdf = FPDF()
        pdf.add_page()
        pdf.set_font("Arial", "B", 16)
        pdf.cell(0, 10, "TEASUR Report", ln=True)
        pdf.set_font("Arial", "", 12)
        for r in report[:500]:  # safety limit
            line = f"{r['nombre']} ({r['edad']} y) - ITS_base:{r['ITS_base']:.2f}, ITS_barrier:{r['ITS_barrier']:.2f}"
            try:
                pdf.multi_cell(0, 8, line)
                pdf.multi_cell(0, 8, r["recomendaciones"])
                pdf.ln(2)
            except Exception:
                # FPDF can be picky with unicode in some environments
                pass
        try:
            pdf.output(str(outdir / "Informe_TEASUR.pdf"))
        except Exception:
            pass


def export_environment(outdir: Path) -> None:
    try:
        import importlib.metadata as md  # Python 3.8+
    except Exception:
        md = None  # type: ignore
    meta = {"python": sys.version}
    for pkg in ["numpy", "pandas", "scikit-learn", "matplotlib", "shap", "fpdf2"]:
        try:
            if md:
                meta[pkg] = md.version(pkg)  # type: ignore
        except Exception:
            pass
    with open(outdir / "environment.json", "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)


def write_requirements(outdir: Path) -> None:
    req = textwrap.dedent(
        """
        numpy>=1.23
        pandas>=1.5
        scikit-learn>=1.2
        matplotlib>=3.6
        # Optional
        shap>=0.45
        fpdf2>=2.7
        """
    ).strip()
    (outdir / "requirements.txt").write_text(req, encoding="utf-8")


def write_license_nc(outdir: Path) -> None:
    lic = textwrap.dedent(
        """
        TEASUR Academic Non-Commercial License v1.0 (Template)
        ------------------------------------------------------
        Copyright (c) 2025 TEASUR Authors.

        Permission is hereby granted, free of charge, to any person obtaining a
        copy of this software and associated documentation files (the "Software"),
        to use, reproduce, and modify the Software for non-commercial academic
        research and teaching purposes, subject to the following conditions:

        1) The above copyright notice and this permission notice shall be
           included in all copies or substantial portions of the Software.
        2) Commercial use, including but not limited to sale, licensing,
           or use in a product or service for which a fee is charged, is
           NOT permitted without prior written permission from the Authors.
        3) The Software is provided "as is", without warranty of any kind.

        For commercial licensing, contact: contacto@teasur.org (example)
        """
    ).strip()
    (outdir / "LICENSE-ACADEMIC-NC.txt").write_text(lic, encoding="utf-8")


def write_citation(outdir: Path, doi_placeholder: str = "10.5281/zenodo.TBD") -> None:
    cff = textwrap.dedent(
        f"""
        cff-version: 1.2.0
        message: "Please cite this software."
        title: "TEASUR Reproducible Bundle"
        version: "{__version__}"
        doi: "{doi_placeholder}"
        authors:
          - family-names: "Díaz Pérez"
            given-names: "Anderson"
        date-released: "2025-08-12"
        """
    )
    (outdir / "CITATION.cff").write_text(cff.strip() + "\n", encoding="utf-8")


def write_readme(outdir: Path) -> None:
    md = textwrap.dedent(
        f"""
        # TEASUR Reproducible Bundle v{__version__}

        Minimal steps:

        ```bash
        python teasur_reproducible.py --outdir out --no-shap --no-pdf
        ```

        Parameters:
        - `--csv`: path to a real CSV (optional). If omitted, synthetic data are generated.
        - `--n`: number of synthetic records (default 300).
        - `--seed`: random seed (default 42).
        - `--outdir`: output folder.
        - `--no-shap`: skip SHAP artifacts (faster).
        - `--no-pdf`: skip PDF generation.
        - `--country-ci`: number of bootstrap resamples per country for 95% CI.
        - `--shap-samples`: limit samples used to compute SHAP (default 1000).
        - `--show-inline`: display figures inline if running in a notebook.

        Outputs:
        - Figures (*.png), JSON with recommendations, CSV cohort, environment,
          requirements, license template, CITATION.cff, MANIFEST.
        """
    )
    (outdir / "README.md").write_text(md.strip() + "\n", encoding="utf-8")


# ------------------------- SHAP artifacts -------------------------
def save_shap_artifacts(
    df: pd.DataFrame,
    model: LinearRegression,
    feature_dims: List[str],
    outdir: Path,
    max_samples: int = 1000,
    show_inline: bool = False,
) -> None:
    """Compute SHAP summary + bar chart + CSV table (English)."""
    if shap is None:
        print("[WARN] SHAP artifacts skipped: shap not installed.")
        return

    try:
        X = df[feature_dims].values
        if len(X) > max_samples:
            X = X[:max_samples]

        explainer = shap.Explainer(model.predict, X)
        sv = explainer(X)  # Explanation
        shap_values = np.array(sv.values)

        # Importance table (mean |SHAP|)
        imp = np.abs(shap_values).mean(axis=0)
        imp_rel = imp / (imp.sum() if imp.sum() > 0 else 1.0)
        importance_df = (
            pd.DataFrame({"feature": feature_dims, "mean_abs_shap": imp, "relative_importance": imp_rel})
            .sort_values("relative_importance", ascending=False)
            .reset_index(drop=True)
        )
        importance_df.to_csv(outdir / "shap_importance.csv", index=False)

        # Summary (beeswarm)
        plt.figure()
        shap.summary_plot(sv, features=df[feature_dims], feature_names=feature_dims, show=False)
        plt.title("SHAP Summary (beeswarm)", fontsize=12)
        save_fig(outdir / "shap_summary.png")

        # Bar chart of relative importance
        plt.figure(figsize=(8, 5))
        ordered = importance_df.sort_values("relative_importance", ascending=True)
        plt.barh(ordered["feature"], ordered["relative_importance"])
        plt.xlabel("Relative importance (mean |SHAP|)")
        plt.title("Feature importance by SHAP")
        save_fig(outdir / "shap_bar.png")

        if show_inline:
            display_inline_images(outdir, ["shap_summary.png", "shap_bar.png"])

    except Exception as e:
        print("[WARN] SHAP artifacts skipped due to error:", e)


# ------------------------- Main pipeline -------------------------
def run(args: argparse.Namespace) -> None:
    outdir = ensure_outdir(args.outdir)

    # 1) Load/simulate data
    df = load_or_simulate(args.csv, n=args.n, seed=args.seed)

    # 2) Outcome + calibration
    if "necesidad_intervencion" not in df.columns:
        df["necesidad_intervencion"] = simulate_outcome(df)
    base_w, model = calibrate_weights(df, "necesidad_intervencion")

    # 3) ITS base and with contextual barriers
    df["ITS_base"] = compute_ITS(df, base_w)
    df["ITS_barrier"] = compute_ITS_barriers(df, base_w)

    # 4) Recommendations
    df["recomendaciones"] = df.apply(generate_recommendation, axis=1)

    # 5) SHAP artifacts (optional)
    if not args.no_shap:
        save_shap_artifacts(
            df, model, FEATURE_DIMS, outdir,
            max_samples=args.shap_samples, show_inline=args.show_inline
        )

    # 6) Stats and figures
    countries, means, lo, hi = bootstrap_country_ci(df, B=args.country_ci)
    mc_vals = mc_sensitivity(df, base_w, R=500)
    make_figures(outdir, countries, means, lo, hi, mc_vals, df)

    # 7) Exports
    export_json_pdf(outdir, df, no_pdf=args.no_pdf)
    df.to_csv(outdir / "cohorte_teasur.csv", index=False)

    # 8) Metadata and templates
    export_environment(outdir)
    write_requirements(outdir)
    write_license_nc(outdir)
    write_citation(outdir)
    write_readme(outdir)

    # 9) Manifest
    manifest = [
        "fig_its_country.png", "fig_mc_hist.png", "fig_box_methods.png",
        "shap_summary.png", "shap_bar.png", "shap_importance.csv",
        "Informe_TEASUR.json", "cohorte_teasur.csv",
        "environment.json", "requirements.txt", "README.md",
        "LICENSE-ACADEMIC-NC.txt", "CITATION.cff",
    ]
    (outdir / "MANIFEST_TEASUR.txt").write_text("\n".join(manifest), encoding="utf-8")

    # 10) Inline display (if requested)
    if args.show_inline:
        display_inline_images(outdir, ["fig_its_country.png", "fig_mc_hist.png", "fig_box_methods.png"])

    print("\n✅ Done. Files generated in:", outdir)


# ------------------------- CLI helpers -------------------------
def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="TEASUR reproducible bundle (English)")
    p.add_argument("--csv", type=str, default=None, help="Path to real CSV (optional)")
    p.add_argument("--n", type=int, default=300, help="N synthetic records if no CSV")
    p.add_argument("--seed", type=int, default=42, help="Random seed")
    p.add_argument("--outdir", type=str, default="out", help="Output folder")
    p.add_argument("--no-shap", action="store_true", help="Disable SHAP explainability artifacts")
    p.add_argument("--no-pdf", action="store_true", help="Skip PDF export")
    p.add_argument("--country-ci", type=int, default=1000, help="Bootstrap resamples per country for 95% CI")
    p.add_argument("--shap-samples", type=int, default=1000, help="Max samples for SHAP computation")
    p.add_argument("--show-inline", action="store_true", help="Display figures inline if in a notebook")
    # Accept unknown args (e.g., -f from Jupyter kernels) to avoid crashes
    args, _ = p.parse_known_args()
    return args


def run_in_notebook(
    csv=None, n=300, seed=42, outdir="out",
    no_shap=False, no_pdf=False, country_ci=1000,
    shap_samples=1000, show_inline=True
):
    """Convenience for Colab/Jupyter (avoids argv conflicts)."""
    ns = argparse.Namespace(
        csv=csv, n=n, seed=seed, outdir=outdir, no_shap=no_shap, no_pdf=no_pdf,
        country_ci=country_ci, shap_samples=shap_samples, show_inline=show_inline
    )
    run(ns)


if __name__ == "__main__":
    args = parse_args()
    run(args)