import gradio as gr
import os
import shutil
import joblib
import numpy as np
import nibabel as nib
import traceback
import tempfile
import torch
from scipy.ndimage import distance_transform_edt


# SET ENV VARIABLES *BEFORE* IMPORTING NNUNET

os.environ['nnUNet_raw'] = './tmp/nnUNet_raw'
os.environ['nnUNet_results'] = './model_weights' 
os.environ['nnUNet_preprocessed'] = './tmp/nnUNet_preprocessed'
os.makedirs('./tmp/nnUNet_raw', exist_ok=True)
os.makedirs('./tmp/nnUNet_preprocessed', exist_ok=True)

from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
from batchgenerators.utilities.file_and_folder_operations import join

# Model Initializations
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
try:
    predictor = nnUNetPredictor(
        tile_step_size=0.5, use_gaussian=True, use_mirroring=True,
        perform_everything_on_device=True, device=device,
        verbose=False, verbose_preprocessing=False, allow_tqdm=True
    )
    model_folder = './model_weights/Dataset998_PanTSTen/nnUNetTrainer__nnUNetPlans__3d_fullres'
    predictor.initialize_from_trained_model_folder(
        model_folder, use_folds=('all',), checkpoint_name='checkpoint_best.pth',
    )
except Exception as e:
    predictor = None
    print(f"Warning: nnU-Net failed to load. {e}")

try:
    risk_model = joblib.load("model.joblib") 
except FileNotFoundError:
    risk_model = None



# FEATURE EXTRACTION

def extract_features(seg_path, high_risk_threshold_mm=5.0):
    """Calculates metrics directly from a NIfTI mask."""
    seg_img = nib.load(seg_path)
    seg_data = seg_img.get_fdata()
    spacing = seg_img.header.get_zooms()[:3]
    voxel_volume_mm3 = spacing[0] * spacing[1] * spacing[2]

    tumor_mask = (seg_data == 10)
    artery_mask = np.isin(seg_data, [2, 8])
    veins_mask = (seg_data == 9)

    if not np.any(tumor_mask):
        return 50.0, 0.0, 0.0, 0.0 

    ttv_mm3 = np.sum(tumor_mask) * voxel_volume_mm3

    dist_artery = distance_transform_edt(~artery_mask, sampling=spacing) if np.any(artery_mask) else np.full(seg_data.shape, 50.0)
    dist_veins = distance_transform_edt(~veins_mask, sampling=spacing) if np.any(veins_mask) else np.full(seg_data.shape, 50.0)

    tumor_artery_distances = dist_artery[tumor_mask]
    tumor_vein_distances = dist_veins[tumor_mask]

    mac = np.min(tumor_artery_distances)
    aev_mm3 = np.sum(tumor_artery_distances <= high_risk_threshold_mm) * voxel_volume_mm3
    vev_mm3 = np.sum(tumor_vein_distances <= high_risk_threshold_mm) * voxel_volume_mm3

    return mac, vev_mm3, aev_mm3, ttv_mm3

def calculate_risk_and_format(mask_path, no_chemo, pd_surgery):
    """Runs extraction, ML prediction, and formats UI output."""
    # 1. Extract raw metrics
    mac, vev, aev, ttv = extract_features(mask_path)

    # Convert textboxes safely
    try:
        chemo_val = int(no_chemo)
        pd_val = int(pd_surgery)
    except ValueError:
        error_msg = "**Error:** Clinical metadata fields must be numbers (1 or 0)."
        return (gr.update(value=error_msg, visible=True), "", "", "", "", "")

    # 2. Hardcode the Means and Stds from your training data
    # Extracted from your means/stds arrays based on their index positions
    TTV_MEAN, TTV_STD = 8980.23827, 8281.19811
    MAC_MEAN, MAC_STD = 8.53567157, 8.92353600
    AEV_MEAN, AEV_STD = 155.784393, 367.152035
    VEV_MEAN, VEV_STD = 576.025386, 559.944824

    # 3. Apply Z-score Normalization
    scaled_ttv = (ttv - TTV_MEAN) / TTV_STD
    scaled_mac = (mac - MAC_MEAN) / MAC_STD
    scaled_aev = (aev - AEV_MEAN) / AEV_STD
    scaled_vev = (vev - VEV_MEAN) / VEV_STD

    # 4. Predict Risk
    if risk_model is not None:
        # STRICT ORDER MATCHING your column_order:
        # ['Chemo_0', 'Surgery_1', 'Total_Tumor_Volume_mm3_arterial', 'Min_Distance_mm_arterial', 'High_Risk_Volume_mm3_5mm_arterial', 'High_Risk_Volume_mm3_5mm_venous']
        features = np.array([[
            chemo_val,   # Chemo_0
            pd_val,      # Surgery_1
            scaled_ttv,  # Total_Tumor_Volume_mm3_arterial
            scaled_mac,  # Min_Distance_mm_arterial
            scaled_aev,  # High_Risk_Volume_mm3_5mm_arterial
            scaled_vev   # High_Risk_Volume_mm3_5mm_venous
        ]])
        
        prob_r1 = risk_model.predict_proba(features)[0][1] 
    else:
        prob_r1 = 0.72 # Fallback dummy

    # 5. Format Outputs
    risk_percentage = prob_r1
    risk_color = "#C55A11" if risk_percentage > 0.5 else "#B89036"
    risk_display = f"<h2 style='color: {risk_color}; margin-top: 0;'>Predicted probability of R+ resection: {risk_percentage:.2f}</h2>"

    # Notice we still display the RAW un-normalized values to the user in the UI, 
    # but the model calculates risk using the SCALED values.
    return (
        gr.update(value="", visible=False), 
        f"{mac:.1f}", 
        f"{vev:.1f}", 
        f"{aev:.1f}", 
        f"{ttv:.1f}", 
        risk_display
    )


# PIPELINE HANDLERS

def process_raw_ct(nifti_file, no_chemo, pd_surgery):
    """Panel 1: Runs heavy nnU-Net segmentation first."""
    if nifti_file is None:
        return gr.update(value="**Error:** Please upload a Raw CT NIfTI file.", visible=True), *[gr.update(value="")]*4, gr.update(value="")
    
    try:
        with tempfile.TemporaryDirectory() as temp_in_dir, tempfile.TemporaryDirectory() as temp_out_dir:
            original_path = nifti_file.name
            base_name = os.path.basename(original_path)

            # Ensure _0000 suffix
            if not base_name.endswith("_0000.nii.gz") and base_name.endswith(".nii.gz"):
                new_base = base_name.replace(".nii.gz", "_0000.nii.gz")
            else:
                new_base = base_name

            input_file_path = os.path.join(temp_in_dir, new_base)
            shutil.copy(original_path, input_file_path)

            if predictor is None:
                raise RuntimeError("nnU-Net is not initialized.")

            # Run Segmentation
            predictor.predict_from_files(
                temp_in_dir, temp_out_dir, save_probabilities=False, overwrite=False,
                num_processes_preprocessing=2, num_processes_segmentation_export=2, 
                folder_with_segs_from_prev_stage=None, num_parts=1, part_id=0
            )

            output_base = new_base.replace("_0000", "")
            output_mask_path = os.path.join(temp_out_dir, output_base)

            if not os.path.exists(output_mask_path):
                raise FileNotFoundError("nnU-Net failed to generate mask.")

            return calculate_risk_and_format(output_mask_path, no_chemo, pd_surgery)

    except Exception as e:
        error_msg = f"**nnU-Net Pipeline Error:** {str(e)}\n\n```\n{traceback.format_exc()}\n```"
        return gr.update(value=error_msg, visible=True), *[gr.update(value="")]*4, gr.update(value="")

def process_presegmented_mask(mask_file, no_chemo, pd_surgery):
    """Panel 2: Skips nnU-Net, processes mask directly."""
    if mask_file is None:
        return gr.update(value="**Error:** Please upload a pre-segmented Mask NIfTI file.", visible=True), *[gr.update(value="")]*4, gr.update(value="")
    
    try:
        # Directly pass the uploaded mask to the extraction logic
        return calculate_risk_and_format(mask_file.name, no_chemo, pd_surgery)
    except Exception as e:
        error_msg = f"**Fast Pipeline Error:** {str(e)}\n\n```\n{traceback.format_exc()}\n```"
        return gr.update(value=error_msg, visible=True), *[gr.update(value="")]*4, gr.update(value="")


# GRADIO UI LAYOUT

with gr.Blocks() as demo:
    gr.Markdown("# Preoperative PDAC R+ Resection Risk Predictor")
    
    with gr.Row():
        # LEFT COLUMN: Inputs (Tabs + Clinical Data)
        with gr.Column(scale=1):
            
            with gr.Tabs():
                # TAB 1: Slow / Full Pipeline
                with gr.TabItem("Raw CT (Runs nnU-Net)"):
                    gr.Markdown("*Note: Segmentation may take >1 hour on CPU.*")
                    file_input_ct = gr.File(label="Upload Raw CT Scan (.nii.gz)", file_count="single")
                    analyze_ct_btn = gr.Button("Segment & Analyze", variant="primary")
                
                # TAB 2: Fast Pipeline
                with gr.TabItem("Pre-segmented Mask (Fast)"):
                    gr.Markdown("*Bypass AI segmentation. Upload your own mask.*")
                    file_input_mask = gr.File(label="Upload Segmentation Mask (.nii.gz)", file_count="single")
                    analyze_mask_btn = gr.Button("Analyze Mask", variant="primary")
            
            gr.Markdown("---")
            gr.Markdown("### Clinical Metadata (Required)")
            no_chemo = gr.Checkbox(label="Absence of Neoadjuvant Chemo (1/0)")
            pd_surgery = gr.Checkbox(label="Patient underwent PD (Pancreaticoduodenectomy) (1/0)")
            
        # RIGHT COLUMN: Shared Outputs
        with gr.Column(scale=2):
            error_box = gr.Markdown(visible=False)
            
            gr.Markdown("### Segmentation Metrics")
            with gr.Row():
                mac_out = gr.Textbox(label="Minimum Arterial Clearance (mm)", interactive=False)
                vev_out = gr.Textbox(label="Venous Encasement Volume (mm³)", interactive=False)
                aev_out = gr.Textbox(label="Arterial Encasement Volume (mm³)", interactive=False)
                ttv_out = gr.Textbox(label="Total Tumor Volume (mm³)", interactive=False)
            
            gr.Markdown("### Risk Assessment")
            risk_out = gr.HTML(value="<i>Upload a file and run analysis...</i>")

    # Connect buttons to their respective functions, but update the SAME output boxes
    outputs_list = [error_box, mac_out, vev_out, aev_out, ttv_out, risk_out]
    
    analyze_ct_btn.click(
        fn=process_raw_ct,
        inputs=[file_input_ct, no_chemo, pd_surgery],
        outputs=outputs_list
    )
    
    analyze_mask_btn.click(
        fn=process_presegmented_mask,
        inputs=[file_input_mask, no_chemo, pd_surgery],
        outputs=outputs_list
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)