"""
Per-habitat radiomics
=====================

**Background.** Classic radiomics computes intensity and texture features
over the whole tumour. Per-habitat radiomics computes the same PyRadiomics
features separately inside each habitat, so every sub-region gets its own
intensity and texture description.

**Purpose.** You get one row per subject with columns
``habitat_{id}_{feature}_of_{modality}`` (plus ``has_habitat_{id}``), maps
of voxel-wise GLCM Id inside each habitat, and a bar chart comparing the
habitats feature by feature.

**Key terms.**

* **first-order feature** -- a statistic of the voxel intensities alone
  (e.g. ``Mean``, ``Energy``), ignoring their spatial arrangement.
* **GLCM feature** -- texture from the grey-level co-occurrence matrix, i.e.
  how often pairs of neighbouring voxels share grey levels (e.g. ``Id``).
* **binWidth** -- width of the intensity bins used before texture features;
  bins are set per habitat, as in PyRadiomics ``execute(label=id)``.
* **one-step habitats** -- see
  :doc:`/auto_examples/04_designs/plot_02_inside_each_subject`; match ids
  before comparing ``habitat_1_*`` across patients
  (:doc:`/auto_examples/06_matching/plot_07_match_labels`).

Extract first-order and GLCM features **within each habitat subregion**
using :class:`~habit.habitat_features.EachHabitatRadiomicsFeatures`.
"""

# %%
# One-step habitats, then per-habitat PyRadiomics on the intensity image.
# sphinx_gallery_thumbnail_number = 1
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

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

from habit.contracts import MaskVolume, cohort_from_directory
from habit.datasets import fetch_demo
from habit.habitat_features import EachHabitatRadiomicsFeatures
from habit.recipes import one_step_habitat
from habit.voxel_features import extract_voxel_texture
from habit.viz import plot_voxel_texture_slice

DATA = fetch_demo()
MODALITIES = ("LAP",)
ROI = "LAP"
cohort = cohort_from_directory(DATA, modalities=MODALITIES, roi=ROI)[:1]
result = one_step_habitat(
    modalities=MODALITIES, n_habitats=3, random_seed=0, roi=ROI
).fit_predict(cohort)
subject = cohort[0]
habitat_map = result.habitat_maps[0]

# Narrow PyRadiomics params keep the gallery table readable.
params: Dict[str, Any] = {
    "imageType": {"Original": {}},
    "featureClass": {
        "firstorder": ["Mean", "Energy"],
        "glcm": ["Autocorrelation", "Id"],
    },
    "setting": {"binWidth": 25, "voxelArrayShift": 0},
}
# One PyRadiomics pass per habitat id, using the habitat map as the mask.
table = EachHabitatRadiomicsFeatures(params=params)(subject, habitat_map)
row = table.frame.iloc[0]
display_cols = [
    col
    for col in table.feature_columns
    if "firstorder" in col or "glcm" in col
][:8]
print("Per-habitat radiomics (sample columns):")
print(row[display_cols].to_string())
row[display_cols]

# %%
# Voxel-wise GLCM Id inside each habitat. Same feature name and
# ``binWidth=25`` as the table; the table is one ROI-level scalar per
# habitat, these maps use a 3x3x3 kernel around each voxel.
Path("out").mkdir(exist_ok=True)
image_vol = subject.image(ROI)
glcm_field = extract_voxel_texture(
    image_vol,
    subject.mask(ROI),
    kernel_radius=1,
    bin_width=25,
    feature_classes={"glcm": ["Id"]},
)
glcm_id_feature = next(
    name for name in glcm_field.feature_names if "glcm_Id" in name
)
habitat_ids: List[int] = [
    hid
    for hid in habitat_map.habitat_ids
    if bool(np.any(habitat_map.label_array == hid))
]
for hid in habitat_ids:
    habitat_roi = MaskVolume.from_geometry(
        (habitat_map.label_array == hid).astype(np.uint8),
        habitat_map.geometry,
        roi_name=ROI,
        labels=(int(hid),),
    )
    fig_tex = plot_voxel_texture_slice(
        glcm_field,
        feature=glcm_id_feature,
        anatomy=image_vol,
        roi_mask=habitat_roi,
        axis=0,
        crop_to="roi",
        title=f"habitat {hid} GLCM Id",
        feature_label="GLCM Id",
    )
    fig_tex.savefig(
        f"out/each_habitat_{hid}_glcm_id.png",
        dpi=150,
        bbox_inches="tight",
    )
    plt.show()

# %%
# One panel per feature so scales stay honest (Mean / Energy / GLCM Id
# must not share a single y-axis — Energy dominates and hides the rest).
feature_specs: List[Tuple[str, str]] = [
    ("Mean", "firstorder_Mean"),
    ("Energy", "firstorder_Energy"),
    ("GLCM Id", "glcm_Id"),
]


def _column_for(habitat_id: int, suffix: str) -> Optional[str]:
    """Return the first feature column matching habitat id + suffix."""
    prefix = f"habitat_{habitat_id}_"
    for col in table.feature_columns:
        if col.startswith(prefix) and suffix in col:
            return col
    return None


panel_rows: List[Dict[str, Any]] = []
for hid in habitat_ids:
    for label, suffix in feature_specs:
        col = _column_for(hid, suffix)
        panel_rows.append(
            {
                "habitat": f"H{hid}",
                "feature": label,
                "value": float(row[col]) if col is not None else float("nan"),
            }
        )
panel = pd.DataFrame(panel_rows)
print(panel.to_string(index=False))
panel

fig, axes = plt.subplots(1, len(feature_specs), figsize=(9.5, 3.2), sharey=False)
if len(feature_specs) == 1:
    axes = [axes]
x = np.arange(len(habitat_ids))
xticklabels = [f"H{hid}" for hid in habitat_ids]
for ax, (label, suffix) in zip(axes, feature_specs):
    values: List[float] = []
    for hid in habitat_ids:
        col = _column_for(hid, suffix)
        values.append(float(row[col]) if col is not None else float("nan"))
    ax.bar(x, values, color="#0072B2", width=0.65)
    ax.set_xticks(x)
    ax.set_xticklabels(xticklabels)
    ax.set_ylabel(label)
    ax.set_title(label)
fig.suptitle("Per-habitat radiomics (one scale per feature)", y=1.02)
fig.tight_layout()
fig.savefig("out/each_habitat_radiomics_bar.png", dpi=150, bbox_inches="tight")
plt.show()