import os
import glob
import cv2
import numpy as np
import tifffile
import pandas as pd
import matplotlib.subplots as plt_subplots
import matplotlib.pyplot as plt
from scipy.ndimage import median_filter
from skimage.filters import gaussian
from skimage import exposure
from huggingface_hub import snapshot_download
from csbdeep.utils import normalize
from stardist.models import Config2D, StarDist2D
from stardist import fill_label_holes, calculate_extents, random_label_cmap
from stardist.matching import matching

# --- 1. Data Loading & Preparation ---
def load_and_split_dataset(repo_id="champ7/celldatamag", test_size=10):
    print(f"Downloading dataset from Hugging Face: {repo_id}...")
    local_dir = snapshot_download(repo_id=repo_id, repo_type="dataset")
    
    # Target specifically _w2.TIF based on the dataset structure
    search_path = os.path.join(local_dir, "**", "*_w2.TIF")
    image_paths = sorted(glob.glob(search_path, recursive=True))
    
    X, Y = [], []
    for img_path in image_paths:
        # Map _w2.TIF to _w3_seg.npy based on the provided screenshot
        mask_path = img_path.replace("_w2.TIF", "_w3_seg.npy")
        
        if os.path.exists(mask_path):
            img = tifffile.imread(img_path)
            mask = np.load(mask_path, allow_pickle=True)
            
            if mask.ndim == 0 and isinstance(mask.item(), dict):
                mask = mask.item().get('masks', mask.item())
            
            img = np.squeeze(img)
            mask = np.squeeze(mask)
                
            mask = fill_label_holes(mask.astype(np.int32))
            X.append(img)
            Y.append(mask)

    print(f"Loaded {len(X)} valid 2D image-mask pairs.")
    
    if len(X) <= test_size:
        raise ValueError("Not enough images to create a split. Reduce test_size or check the dataset.")

    X_test, Y_test = X[:test_size], Y[:test_size]
    X_train, Y_train = X[test_size:], Y[test_size:]
    
    return X_train, Y_train, X_test, Y_test

# --- 2. Preprocessing Definitions ---
def process_w2(image_obj):
    """ W2: Organelles (Red Fluorescence) -> CLAHE & Top-Hat preprocessing """
    if image_obj.ndim == 3: 
        image_obj = image_obj[0]
    p1, p99 = np.percentile(image_obj, (1, 99.9))
    robust_img = exposure.rescale_intensity(image_obj, in_range=(p1, p99), out_range=(0, 255)).astype(np.uint8)
    clahe = cv2.createCLAHE(clipLimit=5.0, tileGridSize=(8, 8))
    clahe_img = clahe.apply(robust_img)
    return clahe_img

def preprocess_2(img):
    """ Standard StarDist Percentile Normalization """
    return normalize(img, 1, 99.8, axis=(0,1))

def preprocess_none(img):
    """ Basic Min-Max scaling to prevent float crashes on raw inputs """
    return (img - np.min(img)) / (np.ptp(img) + 1e-8)

def apply_preprocessing(img, mode="none", to_rgb=False):
    """ Router for preprocessing modes """
    if mode == "prep1":
        out = process_w2(img)
    elif mode == "prep2":
        out = preprocess_2(img)
    elif mode == "prep1_then_prep2":
        out = process_w2(img)
        out = preprocess_2(out)
    else:
        out = preprocess_none(img)
        
    out = out.astype(np.float32)

    if to_rgb and out.ndim == 2:
        out = np.stack((out,) * 3, axis=-1)
        
    return out

# --- 3. Visualization Function ---
def plot_test_results(X_test, Y_test, Y_pred, X_input, experiment_name, mode_label, num_images=2):
    """Plots 3 columns: Original+GT, Model Input, Model Input+Pred"""
    print(f"\nGenerating plots for: {experiment_name}")
    lbl_cmap = random_label_cmap()
    
    for i in range(min(num_images, len(X_test))):
        fig, axes = plt.subplots(1, 3, figsize=(18, 6))
        
        # Column 1: Original Image + GT
        axes[0].imshow(X_test[i], cmap='gray')
        axes[0].imshow(Y_test[i], cmap=lbl_cmap, alpha=0.5)
        axes[0].set_title(f"Image {i+1}: Original + Ground Truth")
        axes[0].axis('off')
        
        # Column 2: The actual input fed to the model
        disp_img = X_input[i]
        if disp_img.ndim == 3: 
            disp_img = disp_img[..., 0] 
            
        axes[1].imshow(disp_img, cmap='gray')
        axes[1].set_title(f"Image {i+1}: Model Input ({mode_label})")
        axes[1].axis('off')
        
        # Column 3: Model Input + Prediction
        axes[2].imshow(disp_img, cmap='gray')
        axes[2].imshow(Y_pred[i], cmap=lbl_cmap, alpha=0.5)
        axes[2].set_title(f"Image {i+1}: Model Input + Prediction")
        axes[2].axis('off')
        
        plt.tight_layout()
        plt.show()

# --- 4. Training Function ---
def train_model(X_train, Y_train, model_name, prep_mode="none"):
    print(f"\n--- Training Custom Model ONCE: {model_name} (Training Data Preprocessing: {prep_mode}) ---")
    
    X_trn_proc = [apply_preprocessing(x, mode=prep_mode, to_rgb=False) for x in X_train]
    
    rng = np.random.RandomState(42)
    ind = rng.permutation(len(X_trn_proc))
    n_val = max(1, int(round(0.15 * len(ind))))
    ind_train, ind_val = ind[:-n_val], ind[-n_val:]
    
    X_trn = [X_trn_proc[i] for i in ind_train]
    Y_trn = [Y_train[i] for i in ind_train]
    X_val = [X_trn_proc[i] for i in ind_val]
    Y_val = [Y_train[i] for i in ind_val]
    
    conf = Config2D(
        n_rays=32,
        grid=(2,2),
        n_channel_in=1,
        train_patch_size=(256,256),
        train_batch_size=4,
    )
    model = StarDist2D(conf, name=model_name, basedir='models')
    
    median_size = calculate_extents(list(Y_train), np.median)
    fov = np.array(model._axes_tile_overlap('YX'))
    if any(median_size > fov):
        print("WARNING: Median object size is larger than the field of view.")
    
    model.train(X_trn, Y_trn, validation_data=(X_val, Y_val), epochs=30, steps_per_epoch=50)
    model.optimize_thresholds(X_val, Y_val)
    
    return model

# --- 5. Evaluation Loop ---
def run_evaluation(model, X_test, Y_test, experiment_name, prep_mode, is_pretrained_he=False):
    print(f"\nEvaluating: {experiment_name}")
    metrics = []
    
    Y_pred_list = []
    X_input_list = []
    
    for i, (img, gt) in enumerate(zip(X_test, Y_test)):
        img_input = apply_preprocessing(img, mode=prep_mode, to_rgb=is_pretrained_he)
        X_input_list.append(img_input)
        
        pred_mask, _ = model.predict_instances(img_input, prob_thresh=0.5, nms_thresh=0.3)
        Y_pred_list.append(pred_mask)
        
        # IoU Threshold 0.0 for matching calculation
        res = matching(gt, pred_mask, thresh=0.0)
        
        metrics.append({
            "Image": i+1, 
            "F1": res.f1, 
            "IoU": res.mean_matched_score 
        })
    
    df = pd.DataFrame(metrics)
    mean_f1 = df['F1'].mean()
    mean_iou = df['IoU'].mean()
    
    print(f"Results for {experiment_name}:")
    print(f"Mean F1-Score: {mean_f1:.4f}")
    print(f"Mean IoU:      {mean_iou:.4f}\n")
    
    plot_test_results(X_test, Y_test, Y_pred_list, X_input_list, experiment_name, prep_mode, num_images=2)
    
    return mean_f1, mean_iou

# --- Main Execution Block ---
if __name__ == "__main__":
    np.random.seed(42) 
    
    # 1. Load Data with updated repo
    X_train, Y_train, X_test, Y_test = load_and_split_dataset(repo_id="champ7/celldatamag", test_size=10)
    
    all_results = []
    
    # 2. Preload the versatile H&E model
    print("\nLoading pre-trained '2D_versatile_he' model...")
    pretrained_model = StarDist2D.from_pretrained('2D_versatile_he')

    # 3. Train the custom model EXACTLY ONCE (using 'none' as the baseline training preprocessing)
    custom_trained_model = train_model(X_train, Y_train, model_name="single_custom_model", prep_mode="none")

    # 4. Define all 8 experiments
    experiments = [
        # (Experiment Name, Preprocessing Mode, Model to Evaluate, is_pretrained_flag)
        ("Exp 1: Pretrained, No Preprocessing", "none", pretrained_model, True),
        ("Exp 2: Pretrained, Preprocessing 1", "prep1", pretrained_model, True),
        ("Exp 3: Pretrained, Preprocessing 2", "prep2", pretrained_model, True),
        ("Exp 4: Pretrained, Prep 1 -> Prep 2", "prep1_then_prep2", pretrained_model, True),
        
        ("Exp 5: Trained, No Preprocessing", "none", custom_trained_model, False),
        ("Exp 6: Trained, Preprocessing 1", "prep1", custom_trained_model, False),
        ("Exp 7: Trained, Preprocessing 2", "prep2", custom_trained_model, False),
        ("Exp 8: Trained, Prep 1 -> Prep 2", "prep1_then_prep2", custom_trained_model, False),
    ]

    # 5. Run through all inference/evaluation states
    for exp_name, prep_mode, model_to_eval, is_pretrained in experiments:
        f1, iou = run_evaluation(model_to_eval, X_test, Y_test, exp_name, prep_mode, is_pretrained_he=is_pretrained)
        all_results.append({
            "Experiment": exp_name,
            "Mean F1-Score": round(f1, 4),
            "Mean IoU": round(iou, 4)
        })

    # --- 6. Final Summary Table ---
    print("\n" + "="*70)
    print("FINAL METRICS SUMMARY (IoU Threshold = 0.0)".center(70))
    print("="*70)
    
    results_df = pd.DataFrame(all_results)
    print(results_df.to_string(index=False, justify='left'))
    print("="*70)