import nibabel as nib
import numpy as np
import os
import cv2

def run(files):
    data = []
    for f in files:
        img = nib.load(f).get_fdata()
        data.append(img)

    seg = np.zeros_like(data[0], dtype=np.uint8)

    # Segmentation fictive (tu remplaceras par ton modèle réel)
    seg[data[0] > np.percentile(data[0], 99)] = 4  # Enhancing Tumor (ET)
    seg[data[1] > np.percentile(data[1], 99)] = 1  # NCR/NET
    seg[data[2] > np.percentile(data[2], 99)] = 2  # Edema

    voxel_volume_ml = np.prod(nib.load(files[0]).header.get_zooms()) / 1000.0
    voxels = int((seg > 0).sum())
    volume_ml = voxels * voxel_volume_ml

    # Décomposition par classes
    ncr_voxels = int((seg == 1).sum())
    ed_voxels  = int((seg == 2).sum())
    et_voxels  = int((seg == 4).sum())

    ncr_vol = ncr_voxels * voxel_volume_ml
    ed_vol  = ed_voxels * voxel_volume_ml
    et_vol  = et_voxels * voxel_volume_ml

    # Rapport enrichi
    report_text = f"""
==== Compte-rendu automatique de segmentation ====

Modalités utilisées : FLAIR, T1, T1CE, T2

--- Volumétrie ---
Volume tumoral total : {volume_ml:.2f} ml ({voxels} voxels)
 - Nécrose / non rehaussé (NCR/NET) : {ncr_vol:.2f} ml ({ncr_voxels} voxels)
 - Œdème péri-tumoral (ED)         : {ed_vol:.2f} ml ({ed_voxels} voxels)
 - Masse rehaussée (ET)            : {et_vol:.2f} ml ({et_voxels} voxels)

--- Interprétation clinique ---
- Masse suspectée détectée.
{ "- Présence de zones nécrotiques." if ncr_voxels > 0 else "" }
{ "- Œdème détecté." if ed_voxels > 0 else "" }
{ "- Rehaussement détecté → suspicion de tumeur de haut grade." if et_voxels > 0 else "- Pas de rehaussement visible." }
- Volume > 100 ml → effet de masse probable.

--- Recommandations ---
1. Discussion multidisciplinaire (neuro-oncologie).
2. IRM avec injection + séquence de perfusion si dispo.
3. Biopsie/examen histologique recommandé.

⚠️ Rapport généré automatiquement, non destiné à remplacer un avis médical.
"""

    # Sauvegardes dans /tmp
    output_dir = "/tmp/output"
    os.makedirs(output_dir, exist_ok=True)
    nii_path = os.path.join(output_dir, "irm_image.nii.gz")
    report_path = os.path.join(output_dir, "rapport.txt")
    mask_path = os.path.join(output_dir, "mask_preview.png")

    nib.save(nib.Nifti1Image(data[0], np.eye(4)), nii_path)
    with open(report_path, "w") as f: f.write(report_text)

    # Sauvegarde du masque visible (couleur)
    slice_idx = seg.shape[2] // 2
    mask_rgb = cv2.applyColorMap((seg[:,:,slice_idx] * 60).astype(np.uint8), cv2.COLORMAP_JET)
    cv2.imwrite(mask_path, mask_rgb)

    return seg, report_text, (nii_path, report_path, mask_path)