import os
import sys
import torch
import nibabel as nib
import numpy as np
from pathlib import Path
import SimpleITK as sitk

# =======================
# CONFIG
# =======================
SAVE_DIR = "./results"
os.makedirs(SAVE_DIR, exist_ok=True)

MODEL_PATH = "model.best"   # ton modèle entraîné

# =======================
# CHARGEMENT DU MODÈLE
# =======================
def load_model():
    print("🔄 Chargement du modèle...")
    model = torch.load(MODEL_PATH, map_location="cpu")
    model.eval()
    return model

# =======================
# PRÉPROCESSING
# =======================
def load_image(file_path):
    ext = Path(file_path).suffix.lower()
    
    if ext in [".nii", ".gz"]:
        img = nib.load(file_path)
        data = img.get_fdata()
        affine = img.affine
        return data, affine
    
    elif ext == ".dcm":
        itk_img = sitk.ReadImage(file_path)
        data = sitk.GetArrayFromImage(itk_img)
        affine = np.eye(4)  # placeholder simple
        return data, affine
    
    else:
        raise ValueError("❌ Format non supporté (seulement NIfTI ou DICOM).")

# =======================
# INFERENCE
# =======================
def run_inference(model, data):
    data = (data - np.min(data)) / (np.max(data) - np.min(data) + 1e-8)
    data_tensor = torch.tensor(data, dtype=torch.float32).unsqueeze(0).unsqueeze(0)  # (B,C,H,W,D)
    
    with torch.no_grad():
        pred = model(data_tensor)
        pred = torch.argmax(pred, dim=1).squeeze().cpu().numpy()
    return pred

# =======================
# RAPPORT
# =======================
def generate_report(pred, save_path):
    volume_total = np.prod(pred.shape)
    volume_tumeur = np.sum(pred > 0)
    ratio = (volume_tumeur / volume_total) * 100
    
    report = [
        "📄 RAPPORT MÉDICAL AUTOMATISÉ",
        "-----------------------------------",
        f"Dimensions de l'image : {pred.shape}",
        f"Volume total voxels   : {volume_total}",
        f"Volume suspect tumeur : {volume_tumeur}",
        f"Ratio suspect         : {ratio:.2f} %",
        "",
        "⚠️ Rapport automatique à valider par un médecin."
    ]
    
    with open(save_path, "w", encoding="utf-8") as f:
        f.write("\n".join(report))
    
    print("\n".join(report))

# =======================
# MAIN
# =======================
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("❌ Utilisation : python test_model.py <chemin_fichier_IRM>")
        sys.exit(1)
    
    input_path = sys.argv[1]
    print(f"📂 Fichier en entrée : {input_path}")

    # Load
    model = load_model()
    data, affine = load_image(input_path)

    # Inference
    pred = run_inference(model, data)

    # Save segmentation
    out_seg_path = os.path.join(SAVE_DIR, "segmentation_pred.nii.gz")
    nib.save(nib.Nifti1Image(pred.astype(np.uint8), affine), out_seg_path)
    print(f"✅ Segmentation sauvegardée : {out_seg_path}")

    # Save report
    out_report_path = os.path.join(SAVE_DIR, "segmentation_report_detailed.txt")
    generate_report(pred, out_report_path)
    print(f"✅ Rapport sauvegardé : {out_report_path}")