import torch
import numpy as np
from torch.utils.data import DataLoader
from tqdm import tqdm
import yaml
import wandb
import cv2
from monai.losses import DiceCELoss
from monai.metrics import DiceMetric
from monai.transforms import Compose, Activations, AsDiscrete
from monai.networks.nets import UNet, SwinUNETR
from dataloading.dataloader2D import NiftiSegmentationDataset
import os

# -------------------------
# Load Config
# -------------------------
with open("/workspace/Segmentation/config2d.yaml", "r") as f:
    config = yaml.safe_load(f)

output_dir = config['output_dir']


device = torch.device(config['device'] if torch.cuda.is_available() else "cpu")

# -------------------------
# Load Model
# -------------------------

model =SwinUNETR(img_size=(256, 256), in_channels=3, out_channels=config['num_classes'], spatial_dims=2).to(device)

model.load_state_dict(torch.load("/workspace/Segmentation/checkpoints/best_model.pth", map_location=device))
model.eval()

# -------------------------
# Prepare Test Data
# -------------------------
test_dataset = NiftiSegmentationDataset(config['csv_file_train'], channel_keys=config['channel_keys'], augment=False)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)

# -------------------------
# Metrics
# -------------------------
dice_metric = DiceMetric(include_background=False, reduction="mean", get_not_nans=False)
post_pred = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])

# -------------------------
# Run Inference
# -------------------------
with torch.no_grad():
    for i, (x, y) in enumerate(tqdm(test_loader, desc="Testing")):
        x, y = x.to(device), y.to(device)

        logits = model(x)
        probs = post_pred(logits)
        dice_metric(y_pred=probs, y=y)

        img_np = x[0].detach().cpu().numpy()
        gt_np = y[0].detach().cpu().numpy().astype(np.uint8)
        pred_np = torch.argmax(probs[0], dim=0).detach().cpu().numpy().astype(np.uint8)

        if img_np.shape[0] == 1:
            img_np = img_np[0]

        gt_np = np.squeeze(gt_np).astype(np.uint8)
        pred_np = np.squeeze(pred_np).astype(np.uint8)

        img_norm = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8)
        img_uint8 = (img_norm * 255).astype(np.uint8)
        img_rgb = cv2.cvtColor(img_uint8, cv2.COLOR_GRAY2BGR)

        contour_img = img_rgb.copy()
        contours_gt, _ = cv2.findContours(gt_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(contour_img, contours_gt, -1, (0, 255, 0), 2)

        contours_pred, _ = cv2.findContours(pred_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(contour_img, contours_pred, -1, (0, 0, 255), 2)

        # Save result to output folder
        filename = os.path.basename(test_dataset.samples[i]["image_path"])
        save_path = os.path.join(output_dir, f"overlay_{i}_{filename.replace('.nii.gz', '.png')}")
        cv2.imwrite(save_path, contour_img)

print(dice_metric.aggregate().item())