!pip install numpy tifffile scikit-image matplotlib huggingface_hub scikit-learn micro_sam
!conda install -c conda-forge micro_sam

import os
import glob
import time
import numpy as np
import tifffile as tiff
import matplotlib.pyplot as plt
from skimage.filters import gaussian, median
from skimage.morphology import square
from skimage import exposure
from huggingface_hub import snapshot_download
from sklearn.metrics import f1_score

# --- micro_sam imports ---
import micro_sam.training as sam_training
from micro_sam.util import export_custom_sam_model
from micro_sam.automatic_segmentation import (
    get_predictor_and_segmenter, 
    automatic_instance_segmentation
)


# ==========================================
# 1. PREPROCESSING FUNCTIONS
# ==========================================

def preprocess_1(input_data):
    """ W3: Cells (Brightfield) -> Blur, Median, Contrast """
    if isinstance(input_data, str):
        img_raw = tiff.imread(input_data).astype(np.float32)
    else:
        img_raw = input_data.astype(np.float32)
        
    img_min, img_max = np.min(img_raw), np.max(img_raw)
    if img_max > img_min:
        img_8bit = ((img_raw - img_min) / (img_max - img_min) * 255.0).astype(np.uint8) 
    else:
        img_8bit = np.zeros(img_raw.shape, dtype=np.uint8)
        
    img_inverted = 255 - img_8bit 
    img_blur = gaussian(img_inverted, sigma=1, preserve_range=True)
    img_filtered = median(img_blur, footprint=square(5)).astype(np.uint8)
    
    p_low, p_high = np.percentile(img_filtered, (0.75, 99.25))
    if p_low >= p_high: 
        final_img = img_filtered
    else:
        final_img = exposure.rescale_intensity(
            img_filtered, in_range=(p_low, p_high), out_range=(0, 255)
        ).astype(np.float32)
        
    # Output float in range [0, 255] for micro_sam dataloader compatibility
    return final_img.astype(np.float32)

def preprocess_2(input_data):
    """ Percentile Clipping (1st, 99.9th), Normalization, Standardization -> Scaled to [0, 255] """
    if isinstance(input_data, str):
        img = tiff.imread(input_data).astype(np.float32)
    else:
        img = input_data.astype(np.float32)
        
    # Percentile clipping
    p_low, p_high = np.percentile(img, (1.0, 99.9))
    img_clipped = np.clip(img, p_low, p_high)
    
    # Scale to [0, 1]
    denom = (p_high - p_low) if (p_high - p_low) > 0 else 1.0
    img_scaled = (img_clipped - p_low) / denom
    
    # Zero-mean, unit-variance standardization
    mean, std = np.mean(img_scaled), np.std(img_scaled)
    img_standardized = (img_scaled - mean) / (std + 1e-8)
    
    # Scale exactly to [0, 255] to bypass micro_sam ValueError
    s_min, s_max = np.min(img_standardized), np.max(img_standardized)
    if s_max > s_min:
        img_final = ((img_standardized - s_min) / (s_max - s_min) * 255.0)
    else:
        img_final = np.zeros_like(img_standardized)
        
    return img_final.astype(np.float32)

def preprocess_1_then_2(input_data):
    """ Sequential application of Preprocessing 1 followed by Preprocessing 2 """
    out_1 = preprocess_1(input_data)
    out_2 = preprocess_2(out_1)
    return out_2


# ==========================================
# 2. METRICS & PLOTTING
# ==========================================

def calculate_metrics(y_true, y_pred):
    """ Calculate IoU and F1 score for binary masks. """
    y_true_bin = (y_true > 0).astype(np.uint8)
    y_pred_bin = (y_pred > 0).astype(np.uint8)
    
    intersection = np.logical_and(y_true_bin, y_pred_bin).sum()
    union = np.logical_or(y_true_bin, y_pred_bin).sum()
    
    iou = intersection / union if union > 0 else 0.0
    f1 = f1_score(y_true_bin.flatten(), y_pred_bin.flatten())
    
    return iou, f1

def plot_results(title, actual_img, preprocessed_img, gt_mask, pred_mask, output_filename):
    """ Generates 1x3 subplot and displays it instantly for review. """
    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
    fig.suptitle(title, fontsize=16)
    
    # Image 1: Actual Image with GT Mask
    axes[0].imshow(actual_img, cmap='gray')
    axes[0].imshow(gt_mask > 0, cmap='Reds', alpha=0.3)
    axes[0].set_title("Actual Image + GT Mask")
    axes[0].axis('off')
    
    # Image 2: Preprocessed Image (Input to model)
    axes[1].imshow(preprocessed_img, cmap='gray')
    axes[1].set_title("Preprocessed Input")
    axes[1].axis('off')
    
    # Image 3: Actual Image with Predicted Mask
    axes[2].imshow(actual_img, cmap='gray')
    axes[2].imshow(pred_mask > 0, cmap='Blues', alpha=0.3)
    axes[2].set_title("Actual Image + Predicted Mask")
    axes[2].axis('off')
    
    plt.tight_layout()
    plt.savefig(output_filename)
    
    # Show the plot immediately so you can monitor direction
    plt.show(block=False)
    plt.pause(2) # Pauses briefly to render the UI before continuing
    plt.close()


# ==========================================
# 3. PIPELINE & DATA HANDLING
# ==========================================

def download_data():
    print("Downloading dataset 'champ7/CEllDataW3'...")
    local_dir = snapshot_download(repo_id="champ7/CEllDataW3", repo_type="dataset")
    return os.path.join(local_dir, "combined_data")

def get_file_pairs(data_dir):
    tif_files = sorted(glob.glob(os.path.join(data_dir, "*.TIF")))
    pairs = []
    for tif in tif_files:
        mask_path = tif.replace(".TIF", "_seg.npy")
        if os.path.exists(mask_path):
            pairs.append((tif, mask_path))
    return pairs

def prepare_training_data(pairs, prep_func, prep_name):
    """ Prepares data for the subset pairs. """
    prep_dir = f"./preprocessed_train_data_{prep_name}"
    os.makedirs(prep_dir, exist_ok=True)
    
    for img_path, mask_path in pairs:
        base_name = os.path.basename(img_path)
        img_prep = prep_func(img_path)
        tiff.imwrite(os.path.join(prep_dir, base_name), img_prep)
        
        mask = np.load(mask_path)
        tiff.imwrite(os.path.join(prep_dir, base_name.replace(".TIF", "_seg.tif")), mask.astype(np.int32))
        
    return prep_dir

def train_model(train_dir, model_name):
    print(f"\n--- Training {model_name} (Single Execution) ---")
    
    train_loader = sam_training.default_sam_loader(
        raw_paths=train_dir,
        raw_key="*.TIF",
        label_paths=train_dir,
        label_key="*_seg.tif",
        patch_shape=(1, 512, 512),
        batch_size=2,
        with_segmentation_decoder=True,
        is_train=True,
        num_workers=2,
        shuffle=True
    )
    
    sam_training.train_sam(
        name=model_name,
        model_type="vit_b",
        train_loader=train_loader,
        val_loader=train_loader,
        n_epochs=10, 
        n_objects_per_batch=5,
        with_segmentation_decoder=True,
        device="cuda"
    )
    
    checkpoint_path = os.path.join("checkpoints", model_name, "best.pt")
    export_path = f"./{model_name}_exported.pth"
    export_custom_sam_model(checkpoint_path=checkpoint_path, model_type="vit_b", save_path=export_path)
    return export_path


def run_inference_and_evaluate(pairs, prep_func, prep_name, model_type, model_checkpoint, is_pretrained=True):
    model_status = "Pretrained" if is_pretrained else "Trained"
    print(f"\nEvaluating {prep_name} with {model_status} Model...")
    
    predictor, segmenter = get_predictor_and_segmenter(
        model_type=model_type,
        checkpoint=model_checkpoint if not is_pretrained else None,
    )
    
    total_iou, total_f1 = 0, 0
    
    for idx, (img_path, mask_path) in enumerate(pairs):
        raw_img = tiff.imread(img_path)
        gt_mask = np.load(mask_path)
        
        input_img = prep_func(raw_img)
        pred_mask = automatic_instance_segmentation(predictor, segmenter, input_img)
        
        iou, f1 = calculate_metrics(gt_mask, pred_mask)
        total_iou += iou
        total_f1 += f1
        
        # Plot 2 images per combination as requested
        if idx < 2:
            plot_name = f"{prep_name}_{model_status}_sample_{idx+1}.png"
            title = f"{prep_name} | {model_status} | IoU: {iou:.3f} | F1: {f1:.3f}"
            plot_results(title, raw_img, input_img, gt_mask, pred_mask, plot_name)
            
    avg_iou = total_iou / len(pairs)
    avg_f1 = total_f1 / len(pairs)
    print(f"Finished {model_status} | {prep_name} -> F1: {avg_f1:.4f}, IoU: {avg_iou:.4f}")
    return avg_iou, avg_f1


# ==========================================
# 4. MAIN EXECUTION
# ==========================================

if __name__ == "__main__":
    data_dir = download_data()
    
    # CONSTRAINT MET: Only take 10 images from the dataset
    file_pairs = get_file_pairs(data_dir)[:10]
    print(f"Using a subset of {len(file_pairs)} image-mask pairs for faster processing.")

    combinations = [
        ("Preprocessing 1", preprocess_1, "prep1"),
        ("Preprocessing 2", preprocess_2, "prep2"),
        ("Preprocessing 1 + 2", preprocess_1_then_2, "prep12")
    ]
    
    results = []

    # 1. Evaluate Pretrained Model (3 Iterations)
    for prep_title, prep_func, prep_tag in combinations:
        iou, f1 = run_inference_and_evaluate(
            file_pairs, prep_func, prep_title, 
            model_type="vit_b", model_checkpoint=None, is_pretrained=True
        )
        results.append({"Pipeline": f"{prep_title} -> Pretrained Model", "F1": f1, "IoU@0.0": iou})

    # 2. Train Model (CONSTRAINT MET: Train Once)
    # We prepare training data using 'Preprocessing 1 + 2' 
    train_data_path = prepare_training_data(file_pairs, preprocess_1_then_2, "prep12")
    trained_model_path = train_model(train_data_path, "sam_finetuned_single")

    # 3. Evaluate Trained Model (3 Iterations)
    for prep_title, prep_func, prep_tag in combinations:
        iou, f1 = run_inference_and_evaluate(
            file_pairs, prep_func, prep_title, 
            model_type="vit_b", model_checkpoint=trained_model_path, is_pretrained=False
        )
        results.append({"Pipeline": f"{prep_title} -> Trained Model", "F1": f1, "IoU@0.0": iou})

    # 4. Output Final Table
    print("\n### Final Evaluation Metrics (10 Image Subset)")
    print("| Pipeline Combination | F1 Score | IoU@0.0 |")
    print("| :--- | :--- | :--- |")
    for res in results:
        print(f"| {res['Pipeline']} | {res['F1']:.4f} | {res['IoU@0.0']:.4f} |")