r"""
Prepare the final interpretation package
========================================

The last step of an inversion workflow is not only making one attractive
figure.  A useful handoff package should let another geophysicist answer:

* Which data and processing choices produced this model?
* Which inversion run was selected, and why?
* How well did the model fit the observations?
* Which figures/tables support the interpretation?
* What caveats should be carried into the report?

This example builds a report-ready interpretation package from the artefacts
created by the previous inversion-gallery examples.  It is intentionally
solver-neutral in structure, but it also knows how to include the bundled
ModEM result sample when available.

The package is written under ``docs/examples/inversion/workspaces`` so it can
be regenerated by the gallery without polluting the source tree.
"""

# %%
# 1. Imports and package layout
# -----------------------------

import json
import os
import shutil
import sys
from datetime import datetime, timezone
from pathlib import Path

# sphinx-gallery executes examples without __file__ (the gallery
# runner sets the working directory to this example's folder).
try:
    EXAMPLE_DIR = Path(__file__).resolve().parent
except NameError:
    EXAMPLE_DIR = Path.cwd()

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def repo_root():
    root = os.environ.get("PYCSAMT_DOCS_REPO_ROOT")
    return Path(root) if root else EXAMPLE_DIR.parents[2]


ROOT = repo_root()
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from pycsamt.models.modem import InversionResult

example_dir = EXAMPLE_DIR
workspace = example_dir / "workspaces" / "l18_prepared_workspace"
modem_sample = ROOT / "data" / "modem" / "willy_27freq_watex_line02_sample"

source_table_dirs = [
    workspace / "02_tables",
    example_dir / "workspaces" / "modem_result_tables",
]
source_figure_dirs = [
    workspace / "05_figures",
    example_dir / "workspaces" / "modem_result_figures",
]

package_root = example_dir / "workspaces" / "final_interpretation_package"
package_dirs = {
    "figures": package_root / "figures",
    "tables": package_root / "tables",
    "models": package_root / "models",
    "logs": package_root / "logs",
    "metadata": package_root / "metadata",
    "report": package_root / "report",
}
for path in package_dirs.values():
    path.mkdir(parents=True, exist_ok=True)

# %%
# 2. Helper functions
# -------------------
# A final package should be deterministic.  These helpers copy selected files,
# record their sizes, and write a manifest that can be reviewed later.


def copy_if_exists(src, dst_dir, label, required=False):
    src = Path(src)
    dst_dir = Path(dst_dir)
    record = {
        "label": label,
        "source": str(src),
        "destination": "",
        "exists": src.exists(),
        "required": bool(required),
        "size_bytes": 0,
    }
    if src.exists() and src.is_file():
        dst = dst_dir / src.name
        shutil.copy2(src, dst)
        record["destination"] = str(dst)
        record["size_bytes"] = dst.stat().st_size
    elif required:
        print(f"WARNING: required artefact missing: {src}")
    return record


def collect_by_suffix(src_dirs, suffixes, dst_dir, label, limit=None):
    records = []
    for src_dir in src_dirs:
        if not src_dir.exists():
            continue
        candidates = []
        for suffix in suffixes:
            candidates.extend(sorted(src_dir.glob(f"*{suffix}")))
        if limit is not None:
            candidates = candidates[:limit]
        for src in candidates:
            records.append(copy_if_exists(src, dst_dir, label))
    return records


def write_text(path, lines):
    path = Path(path)
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return path


def safe_float(value):
    try:
        value = float(value)
    except Exception:
        return None
    if not np.isfinite(value):
        return None
    return value


# %%
# 3. Collect preparation and validation artefacts
# -----------------------------------------------
# These files document the input contract: station coverage, validation checks,
# data/error choices, model grid, and run-folder staging.  If a report includes
# only the final model image but none of these files, it is hard to audit.

manifest_records = []

core_tables = [
    "inversion_input_validation_report.json",
    "inversion_input_validation_table.csv",
    "inversion_data_error_policy.json",
    "inversion_impedance_complex_table.csv",
    "inversion_rho_phase_table.csv",
    "starting_model_policy.json",
    "run_folder_audit.csv",
    "convergence_diagnostics.csv",
    "convergence_history_normalized.csv",
]

for name in core_tables:
    manifest_records.append(
        copy_if_exists(
            workspace / "02_tables" / name,
            package_dirs["tables"],
            label="workspace_table",
            required=False,
        )
    )
    manifest_records.append(
        copy_if_exists(
            workspace / "04_run_files" / name,
            package_dirs["tables"],
            label="run_file_table",
            required=False,
        )
    )

manifest_records.extend(
    collect_by_suffix(
        source_table_dirs,
        suffixes=(".csv", ".json"),
        dst_dir=package_dirs["tables"],
        label="analysis_table",
    )
)

# %%
# 4. Collect figures for interpretation
# -------------------------------------
# The figure folder is deliberately broad: section plots, response diagnostics,
# scenario dashboards, convergence plots, and validation dashboards all belong
# in the final interpretation context.  A report writer can then choose the
# best subset, but the package preserves the full audit trail.

manifest_records.extend(
    collect_by_suffix(
        source_figure_dirs,
        suffixes=(".png", ".jpg", ".jpeg", ".svg", ".pdf"),
        dst_dir=package_dirs["figures"],
        label="figure",
    )
)

# %%
# 5. Include compact ModEM result evidence
# ----------------------------------------
# The package should not duplicate a complete production inversion directory.
# We include only small-to-moderate files that support interpretation:
# convergence log, controls, README, and the final representative response/model
# files already curated in ``data/modem``.

modem_records = []
if modem_sample.exists():
    for name in [
        "README.txt",
        "Modular_NLCG.log",
        "inv.ctrl",
        "fwd.ctrl",
        "CSUr2.err",
        "Modular_NLCG_073.dat",
        "Modular_NLCG_073.res",
        "Modular_NLCG_073.rho",
    ]:
        modem_records.append(
            copy_if_exists(
                modem_sample / name,
                package_dirs["models"]
                if name.endswith(".rho")
                else package_dirs["logs"],
                label="modem_result_evidence",
                required=name in {"README.txt", "Modular_NLCG.log"},
            )
        )
manifest_records.extend(modem_records)

# %%
# 6. Summarize the selected ModEM run
# -----------------------------------
# This section is intentionally compact.  It records headline numbers that a
# reviewer can compare against response plots and section figures.

run_summary = {
    "available": False,
    "sample_dir": str(modem_sample),
}

if modem_sample.exists():
    result = InversionResult(
        modem_sample,
        load_control=False,
        load_covariance=False,
        load_models=False,
        load_data=False,
    )
    model_keys = [
        path.stem.replace("Modular_NLCG_", "iter_")
        for path in sorted(modem_sample.glob("Modular_NLCG_*.rho"))
        if path.stat().st_size > 0
    ]
    run_summary.update(
        {
            "available": True,
            "mode": result.mode,
            "n_log_records": int(result.n_iter),
            "final_rms": safe_float(result.final_rms),
            "best_rms": safe_float(result.best_rms),
            "iteration_numbers": [int(v) for v in result.iteration_numbers],
            "model_keys": model_keys,
            "has_observed_data": (
                modem_sample / "27-freq-run-watex01.dat"
            ).exists(),
            "has_predicted_data": (
                modem_sample / "Modular_NLCG_073.dat"
            ).exists(),
            "observed_sites": "not loaded in fast package mode",
            "observed_periods": "not loaded in fast package mode",
        }
    )

summary_file = (
    package_dirs["metadata"] / "selected_inversion_run_summary.json"
)
summary_file.write_text(json.dumps(run_summary, indent=2), encoding="utf-8")
manifest_records.append(
    {
        "label": "run_summary",
        "source": "generated",
        "destination": str(summary_file),
        "exists": True,
        "required": True,
        "size_bytes": summary_file.stat().st_size,
    }
)

# %%
# 7. Create a report-ready figure index
# -------------------------------------
# A package with 30 files is much easier to use when every figure has a short
# purpose.  We infer a first-pass purpose from the filename.  A human can edit
# this CSV before final report delivery.


def figure_purpose(name):
    lower = name.lower()
    if "validation" in lower:
        return "Input validation and data-readiness audit"
    if "convergence" in lower or "misfit" in lower:
        return "Inversion convergence / fit history"
    if "response" in lower or "pseudo" in lower or "residual" in lower:
        return "Observed-vs-predicted response diagnostic"
    if "section" in lower or "model" in lower:
        return "Model section or resistivity-structure diagnostic"
    if "scenario" in lower or "tradeoff" in lower:
        return "Scenario comparison and decision support"
    return "Supporting inversion figure"


figure_rows = []
for fig in sorted(package_dirs["figures"].glob("*")):
    if fig.is_file():
        figure_rows.append(
            {
                "file": fig.name,
                "purpose": figure_purpose(fig.name),
                "size_bytes": fig.stat().st_size,
                "include_in_report": "review",
            }
        )

figure_index = pd.DataFrame(figure_rows)
figure_index_file = package_dirs["metadata"] / "figure_index.csv"
figure_index.to_csv(figure_index_file, index=False)

# %%
# 8. Create an interpretation checklist
# -------------------------------------
# This checklist is not just administrative.  It prevents a common failure
# mode: exporting the final section without documenting why the model is
# defensible.

checklist = pd.DataFrame(
    [
        {
            "item": "Input validation has no blocking failures",
            "status": "review",
            "evidence": "inversion_input_validation_report.json",
        },
        {
            "item": "Error floors and masked data are documented",
            "status": "review",
            "evidence": "inversion_data_error_policy.json",
        },
        {
            "item": "Convergence history was inspected",
            "status": "review",
            "evidence": "convergence diagnostics / ModEM log",
        },
        {
            "item": "Observed vs predicted responses were compared",
            "status": "review",
            "evidence": "response panels, residual heatmaps, station summary",
        },
        {
            "item": "Scenario choice is justified beyond lowest RMS",
            "status": "review",
            "evidence": "inversion_scenario_recommendation.json",
        },
        {
            "item": "Final section caveats are written in the report memo",
            "status": "review",
            "evidence": "interpretation_memo.md",
        },
    ]
)
checklist_file = package_dirs["metadata"] / "interpretation_checklist.csv"
checklist.to_csv(checklist_file, index=False)

# %%
# 9. Write the interpretation memo
# --------------------------------
# The memo is a report skeleton.  It is intentionally conservative: it records
# what was packaged, what remains to review, and which caveats should be carried
# forward.  Users can edit it into a project-specific report section.

created_utc = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
memo_lines = [
    "# Final inversion interpretation package",
    "",
    f"Created: {created_utc}",
    "",
    "## Purpose",
    "",
    "This package collects inversion-preparation tables, validation outputs, "
    "convergence diagnostics, response checks, scenario-comparison artefacts, "
    "and selected ModEM result evidence for report preparation.",
    "",
    "## Selected inversion run",
    "",
]

if run_summary["available"]:
    memo_lines.extend(
        [
            f"- Source folder: `{run_summary['sample_dir']}`",
            f"- Detected mode: `{run_summary['mode']}`",
            f"- Parsed log records: `{run_summary['n_log_records']}`",
            f"- Final RMS: `{run_summary['final_rms']}`",
            f"- Best RMS: `{run_summary['best_rms']}`",
            f"- Model snapshots: `{', '.join(run_summary['model_keys'])}`",
            f"- Observed sites: `{run_summary['observed_sites']}`",
            f"- Observed periods: `{run_summary['observed_periods']}`",
        ]
    )
else:
    memo_lines.append("- No bundled ModEM result sample was available.")

memo_lines.extend(
    [
        "",
        "## Recommended review order",
        "",
        "1. Read `metadata/interpretation_checklist.csv`.",
        "2. Inspect input validation and data/error policy tables.",
        "3. Inspect convergence and scenario-comparison figures.",
        "4. Inspect observed-vs-predicted response figures and residual tables.",
        "5. Only then select final model-section figures for the report.",
        "",
        "## Caveats to carry into interpretation",
        "",
        "- A low RMS does not guarantee a geologically meaningful model.",
        "- Coherent response residuals should be explained before interpreting "
        "nearby model features.",
        "- Structures that appear only in late iterations or aggressive-error "
        "scenarios require extra caution.",
        "- Topography, static shift, and error floors should be described in "
        "the final report.",
        "- This gallery package is a reproducible template; project teams should "
        "replace demonstration scenario rows with their real run folders.",
        "",
        "## Package contents",
        "",
        "- `figures/`: copied diagnostic and interpretation figures.",
        "- `tables/`: copied CSV/JSON data, validation, response, and scenario tables.",
        "- `models/`: selected representative model artefacts.",
        "- `logs/`: solver logs, control files, and run notes.",
        "- `metadata/`: manifests, figure index, checklist, and run summary.",
    ]
)

memo_file = write_text(
    package_dirs["report"] / "interpretation_memo.md", memo_lines
)

# %%
# 10. Write the package manifest and a package summary figure
# -----------------------------------------------------------

manifest = pd.DataFrame(manifest_records)
manifest = manifest.drop_duplicates(subset=["destination", "source", "label"])
manifest_file = package_dirs["metadata"] / "package_manifest.csv"
manifest.to_csv(manifest_file, index=False)

summary_counts = (
    manifest.assign(
        category=manifest["destination"].map(
            lambda p: (
                Path(p).parent.name if isinstance(p, str) and p else "missing"
            )
        )
    )
    .groupby("category")
    .agg(n_files=("destination", lambda s: int(np.sum([bool(v) for v in s]))))
    .reset_index()
)

fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.8), constrained_layout=True)
ax_count, ax_status = axes

ax_count.bar(
    summary_counts["category"], summary_counts["n_files"], color="tab:blue"
)
ax_count.set_title("Packaged artefacts by folder")
ax_count.set_ylabel("File count")
ax_count.tick_params(axis="x", rotation=30)
ax_count.grid(axis="y", alpha=0.25)

status_counts = (
    manifest["exists"]
    .value_counts()
    .rename(index={True: "copied", False: "missing"})
)
ax_status.pie(
    status_counts.to_numpy(),
    labels=status_counts.index,
    autopct="%1.0f%%",
    startangle=90,
    colors=["tab:green", "tab:red"][: len(status_counts)],
)
ax_status.set_title("Manifest status")

fig.suptitle("Final interpretation package summary")
summary_plot = package_dirs["metadata"] / "package_summary.png"
fig.savefig(summary_plot, dpi=120)
plt.show()

print(f"Package root: {package_root}")
print(f"Manifest: {manifest_file}")
print(f"Figure index: {figure_index_file}")
print(f"Checklist: {checklist_file}")
print(f"Run summary: {summary_file}")
print(f"Interpretation memo: {memo_file}")
print(f"Package summary figure: {summary_plot}")

# %%
# 11. What to deliver
# -------------------
#
# For a real project handoff, zip or archive ``final_interpretation_package``
# after reviewing the checklist.  The recommended minimum deliverables are:
#
# * one final model-section figure plus one scenario-comparison figure;
# * observed-vs-predicted response panels for representative stations;
# * residual heatmap or station/component misfit table;
# * convergence log/plot;
# * data/error policy and validation report;
# * a short memo explaining caveats and why the selected scenario is preferred.
#
# This structure helps future users reproduce the interpretation path instead
# of only seeing the final picture.

# sphinx_gallery_thumbnail_number = 1