import os
import glob
import random
import numpy as np
import pandas as pd
import tifffile as tiff
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 cellpose_omni import models, metrics, core

# ==========================================
# 1. Custom Preprocessing (Prep 1)
# ==========================================
def preprocess_w3(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)
    img_8bit = ((img_raw - img_min) / (img_max - img_min) * 255).astype(np.uint8) if img_max > img_min else 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_filter(img_blur, size=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.uint8)
        
    # Standardize to 0.0 - 1.0 floats
    return (final_img / 255.0).astype(np.float32)

# ==========================================
# 2. Data Preparation
# ==========================================
def prepare_data(repo_id="champ7/CEllDataW3"):
    print(f"Downloading dataset {repo_id} from Hugging Face...")
    dataset_path = snapshot_download(repo_id=repo_id, repo_type="dataset")
    data_dir = os.path.join(dataset_path, "combined_data")
    
    image_files = sorted(glob.glob(os.path.join(data_dir, "*.TIF")))
    seg_files = sorted(glob.glob(os.path.join(data_dir, "*_seg.npy")))
    
    images, masks = [], []
    
    for img_path in image_files:
        base_name = img_path.replace(".TIF", "")
        seg_path = base_name + "_seg.npy"
        
        if seg_path in seg_files:
            img = tiff.imread(img_path).astype(np.float32)
            images.append(img)
            
            seg_data = np.load(seg_path, allow_pickle=True)
            if seg_data.shape == ():
                seg_dict = seg_data.item()
                mask_raw = seg_dict['masks']
            else:
                mask_raw = seg_data
                
            # Convert to standard 32-bit integer labels
            masks.append(mask_raw.astype(np.int32))
            
    X_train, y_train = images, masks
    
    combined = list(zip(images, masks))
    random.seed(42)
    sample_10 = random.sample(combined, min(10, len(combined)))
    X_test, y_test = zip(*sample_10)
    
    return list(X_train), list(X_test), list(y_train), list(y_test)

# ==========================================
# 3. Metrics & Utilities
# ==========================================
def compute_metrics(masks_true, masks_pred, threshold=0.0):
    ap, tp, fp, fn = metrics.average_precision(masks_true, masks_pred, threshold=[threshold])
    
    tp_sum = np.sum(tp)
    fp_sum = np.sum(fp)
    fn_sum = np.sum(fn)
    
    f1_score = 0.0 if (2 * tp_sum + fp_sum + fn_sum) == 0 else (2 * tp_sum) / (2 * tp_sum + fp_sum + fn_sum)
    mean_ap = np.mean(ap)
    
    return f1_score, mean_ap

def to_numpy(data):
    if hasattr(data, 'cpu'):
        return data.cpu().detach().numpy()
    elif isinstance(data, list):
        return [to_numpy(d) for d in data]
    elif isinstance(data, tuple):
        return tuple(to_numpy(d) for d in data)
    elif isinstance(data, dict):
        return {k: to_numpy(v) for k, v in data.items()}
    elif isinstance(data, np.ndarray) and data.dtype == object:
        return np.array([to_numpy(d) for d in data])
    return data

# ==========================================
# 4. Visualization Helper (4 Panels)
# ==========================================
def visualize_results(raw_img, fed_img, y_true, y_pred, config_name, model_type, img_idx):
    fig, axes = plt.subplots(1, 4, figsize=(24, 6))
    fig.suptitle(f"{config_name} | {model_type} | Image {img_idx}", fontsize=16, fontweight='bold')

    axes[0].imshow(raw_img, cmap='gray')
    axes[0].set_title("Original Image", fontsize=14)
    axes[0].axis('off')

    axes[1].imshow(raw_img, cmap='gray')
    mask_true_display = np.ma.masked_where(y_true == 0, y_true)
    axes[1].imshow(mask_true_display, cmap='nipy_spectral', alpha=0.5, interpolation='none')
    axes[1].set_title("Original Image + GT Mask", fontsize=14)
    axes[1].axis('off')

    axes[2].imshow(fed_img, cmap='gray')
    axes[2].set_title("Preprocessed Image (Fed to Model)", fontsize=14)
    axes[2].axis('off')

    axes[3].imshow(raw_img, cmap='gray')
    mask_pred_display = np.ma.masked_where(y_pred == 0, y_pred)
    axes[3].imshow(mask_pred_display, cmap='nipy_spectral', alpha=0.5, interpolation='none')
    axes[3].set_title("Original Image + Predicted Mask", fontsize=14)
    axes[3].axis('off')

    plt.tight_layout()
    plt.show()

# ==========================================
# 5. Pipeline Execution
# ==========================================
def run_pipeline():
    # Use GPU for fast inference/evaluation
    use_gpu_eval = core.use_gpu()
    channels = [0, 0]
    base_model_name = 'cyto2'
    
    X_train_raw, X_test_raw, y_train, y_test = prepare_data()
    
    X_train_p1 = [preprocess_w3(img) for img in X_train_raw]
    X_test_p1 = [preprocess_w3(img) for img in X_test_raw]

    configs = [
        {
            "name": "Prep 1", 
            "train_data": X_train_p1, "test_data": X_test_p1,
            "train_kwargs": {"normalize": False, "rescale": None},
            "eval_kwargs": {"normalize": False, "flow_threshold": 0, "mask_threshold": -1, "rescale": None}
        },
        {
            "name": "Prep 2", 
            "train_data": X_train_raw, "test_data": X_test_raw,
            "train_kwargs": {"normalize": True, "rescale": None},
            "eval_kwargs": {"normalize": True, "flow_threshold": 0, "mask_threshold": -1, "rescale": None}
        },
        {
            "name": "Prep 1 + Prep 2", 
            "train_data": X_train_p1, "test_data": X_test_p1,
            "train_kwargs": {"normalize": True, "rescale": None},
            "eval_kwargs": {"normalize": True, "flow_threshold": 0, "mask_threshold": -1, "rescale": None}
        }
    ]

    results = []

    for cfg in configs:
        print(f"\n--- Running Configuration: {cfg['name']} ---")
        
        # ------------------------------------------
        # A. Pretrained Model Evaluation (Runs on GPU if available)
        # ------------------------------------------
        pretrained_model = models.CellposeModel(
            gpu=use_gpu_eval, 
            model_type=base_model_name, 
            omni=True, 
            nchan=2, 
            nclasses=2
        )
        preds_pre, _, _ = pretrained_model.eval(cfg["test_data"], channels=channels, **cfg["eval_kwargs"])
        
        preds_pre = to_numpy(preds_pre)
        f1_pre, iou_pre = compute_metrics(y_test, preds_pre, threshold=0.0)
        
        results.append({
            "Configuration": f"{cfg['name']} + Pretrained + Prediction",
            "F1 Score": f"{f1_pre:.4f}",
            "Mean IoU (AP)": f"{iou_pre:.4f}"
        })
        
        for viz_idx in range(min(2, len(cfg["test_data"]))):
            visualize_results(
                raw_img=X_test_raw[viz_idx], fed_img=cfg["test_data"][viz_idx], 
                y_true=y_test[viz_idx], y_pred=preds_pre[viz_idx],
                config_name=cfg["name"], model_type="Pretrained Model", img_idx=viz_idx + 1
            )
        
        # ------------------------------------------
        # B. Fine-tuning Model (gpu=False BYPASSES CUDA LOSS ASSERTIONS COMPLETELY)
        # ------------------------------------------
        print(f"Fine-tuning {base_model_name} on CPU to guarantee zero CUDA assertion crashes...")
        finetune_model = models.CellposeModel(
            #gpu=False,  # <--- THIS REMOVES THE ASSERTION PROBLEM ENTIRELY
            model_type=base_model_name, 
            omni=True,
            nchan=2,
            nclasses=2
        )
        
        model_path = finetune_model.train(
            cfg["train_data"], y_train, 
            train_links=[None] * len(y_train), 
            channels=channels, 
            save_path=f"./omnipose_{cfg['name'].replace(' ', '_')}", 
            n_epochs=150, 
            learning_rate=0.1,
            batch_size=2,
            **cfg["train_kwargs"]
        )
        
        # ------------------------------------------
        # C. Evaluate Fine-Tuned Model (Evaluates back on GPU)
        # ------------------------------------------
        custom_model = models.CellposeModel(
            gpu=use_gpu_eval, 
            pretrained_model=model_path, 
            model_type=base_model_name, 
            omni=True,
            nchan=2,
            nclasses=2
        )
        preds_fine, _, _ = custom_model.eval(cfg["test_data"], channels=channels, **cfg["eval_kwargs"])
        
        preds_fine = to_numpy(preds_fine)
        f1_fine, iou_fine = compute_metrics(y_test, preds_fine, threshold=0.0)
        
        results.append({
            "Configuration": f"{cfg['name']} + Training + Prediction",
            "F1 Score": f"{f1_fine:.4f}",
            "Mean IoU (AP)": f"{iou_fine:.4f}"
        })
        
        for viz_idx in range(min(2, len(cfg["test_data"]))):
            visualize_results(
                raw_img=X_test_raw[viz_idx], fed_img=cfg["test_data"][viz_idx], 
                y_true=y_test[viz_idx], y_pred=preds_fine[viz_idx],
                config_name=cfg["name"], model_type="Fine-tuned Model", img_idx=viz_idx + 1
            )

    # ==========================================
    # 6. Final Output Compilation
    # ==========================================
    print("\n================ FINAL METRICS TABLE ================")
    results_df = pd.DataFrame(results)
    print(results_df.to_markdown(index=False))

if __name__ == "__main__":
    run_pipeline()