import cv2
import numpy as np
import torch
import albumentations as A
from albumentations.pytorch import ToTensorV2
from tqdm import tqdm
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed


class LungPreprocessor:
    def __init__(self, image_size=224):
        self.image_size = image_size

    def remove_artifacts_and_segment(self, image):
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image.copy()

        h, w = gray.shape
        border = int(min(h, w) * 0.03)
        cropped = gray[border:h-border, border:w-border]
        blur = cv2.GaussianBlur(cropped, (5, 5), 0)

        _, thresh = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        mask = cv2.bitwise_not(thresh)

        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=2)
        kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (30, 30))
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_close, iterations=2)
        mask = cv2.GaussianBlur(mask, (21, 21), 0)

        mask_float = mask.astype(float) / 255.0
        mask_area = mask_float.mean()
        if mask_area < 0.15 or mask_area > 0.85:
            segmented = cropped.copy()
        else:
            segmented = (cropped * mask_float + cropped * 0.2 * (1 - mask_float)).astype(np.uint8)

        return segmented

    def apply_clahe(self, image):
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        return clahe.apply(image)

    def normalize_intensity(self, image):
        p2, p98 = np.percentile(image, (2, 98))
        image = np.clip(image, p2, p98)
        image = ((image - image.min()) / (max(1e-8, image.max() - image.min())) * 255).astype(np.uint8)
        return image

    def preprocess(self, image_path, segment_lung=False):
        image = cv2.imread(str(image_path), cv2.IMREAD_GRAYSCALE)
        if image is None:
            raise ValueError(f"Failed to load image: {image_path}")
        if image.mean() < 1 or image.mean() > 253:
            raise ValueError(f"Image appears corrupted (mean intensity: {image.mean():.1f})")

        if segment_lung:
            image = self.remove_artifacts_and_segment(image)
        image = self.normalize_intensity(image)
        image = self.apply_clahe(image)

        pixel_range = float(image.max()) - float(image.min())
        if pixel_range < 2.0:
            raise ValueError("Preprocessing produced near-uniform image")

        return image


def get_train_transforms(image_size=224):
    return A.Compose([
        A.Resize(height=image_size, width=image_size),
        A.RandomResizedCrop(size=(image_size, image_size), scale=(0.85, 1.0), p=0.5),
        A.HorizontalFlip(p=0.5),
        A.Rotate(limit=10, p=0.7),
        A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.5),
        A.GaussNoise(var_limit=(10.0, 50.0), p=0.3),
        A.GridDistortion(num_steps=5, distort_limit=0.1, p=0.2),
        A.Normalize(mean=[0.5], std=[0.5], max_pixel_value=255.0),
        ToTensorV2()
    ])


def get_val_transforms(image_size=224):
    return A.Compose([
        A.Resize(height=image_size, width=image_size),
        A.Normalize(mean=[0.5], std=[0.5], max_pixel_value=255.0),
        ToTensorV2()
    ])


class PreprocessedDataset(torch.utils.data.Dataset):
    def __init__(self, image_paths, labels, transforms=None, use_preprocessing=True, cache_pt=True, load_to_ram=False):
        self.image_paths = image_paths
        self.labels = labels
        self.transforms = transforms
        self.use_preprocessing = use_preprocessing
        self.cache_pt = cache_pt
        self.preprocessor = LungPreprocessor() if use_preprocessing else None
        self.ram_cache = None

        if load_to_ram:
            self._preload_to_ram()

    def _get_preprocessed(self, img_path, target_size=224):
        if not self.use_preprocessing:
            image = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE)
            image = cv2.resize(image, (target_size, target_size))
            return image

        cache_path = img_path.with_suffix('.pt.cache')
        if self.cache_pt and cache_path.exists():
            return torch.load(cache_path, weights_only=True).numpy()

        image = self.preprocessor.preprocess(str(img_path), segment_lung=False)

        if self.cache_pt:
            torch.save(torch.from_numpy(image), cache_path)

        return image

    def _build_cache(self, paths):
        processor = LungPreprocessor()
        for p in paths:
            try:
                img = processor.preprocess(str(p), segment_lung=False)
                torch.save(torch.from_numpy(img), p.with_suffix('.pt.cache'))
            except Exception:
                pass

    def _preload_to_ram(self):
        n_total = len(self.image_paths)
        print(f"  Pre-loading {n_total} images into RAM...")

        uncached = [p for p in self.image_paths if not p.with_suffix('.pt.cache').exists()]
        if uncached:
            n_workers = min(8, len(uncached))
            chunk_size = (len(uncached) + n_workers - 1) // n_workers
            chunks = [uncached[i:i+chunk_size] for i in range(0, len(uncached), chunk_size)]
            print(f"  Processing {len(uncached)} uncached images with {n_workers} workers...")
            worker = LungPreprocessor()
            for chunk in tqdm(chunks, desc="Caching"):
                for p in chunk:
                    try:
                        img = worker.preprocess(str(p), segment_lung=False)
                        torch.save(torch.from_numpy(img), p.with_suffix('.pt.cache'))
                    except Exception:
                        pass

        self.ram_cache = []
        target_size = None
        if self.transforms:
            for t in self.transforms:
                if hasattr(t, 'height'):
                    target_size = t.height
                    break
        if target_size is None:
            target_size = self.preprocessor.image_size if self.preprocessor else 224

        for path in tqdm(self.image_paths, desc="Loading RAM"):
            try:
                image = torch.load(path.with_suffix('.pt.cache'), weights_only=True).numpy()
                self.ram_cache.append(image)
            except Exception as e:
                print(f"  Warning: Failed to load {path}: {e}")
                self.ram_cache.append(np.zeros((target_size, target_size), dtype=np.uint8))

        n = len(self.ram_cache)
        size_gb = n * target_size * target_size * 1 / 1e9
        print(f"  Cached {n} images in RAM (~{size_gb:.1f} GB)")

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        label = self.labels[idx]

        if self.ram_cache is not None:
            image = self.ram_cache[idx]
        elif self.use_preprocessing:
            image = self._get_preprocessed(self.image_paths[idx])
        else:
            image = cv2.imread(str(self.image_paths[idx]), cv2.IMREAD_GRAYSCALE)
            target_size = self.transforms[0].height if self.transforms and hasattr(self.transforms[0], 'height') else 224
            image = cv2.resize(image, (target_size, target_size))

        if self.transforms:
            augmented = self.transforms(image=image)
            image = augmented['image']
        else:
            image = torch.from_numpy(image).unsqueeze(0).float() / 255.0

        if image.shape[0] == 3:
            image = image.mean(dim=0, keepdim=True)
        elif image.dim() == 2:
            image = image.unsqueeze(0)

        return image, label