import cv2
import numpy as np
import torch

def apply_phase1_hardening(patch, landmarks, apply_dropout=True, apply_denoise=True, apply_clahe=True, apply_noise=True, apply_jitter=True):
    """
    Applies Domain Adaptation Phase 1 (Data Hardening) to a single sample.
    patch: (4, H, W) numpy array, uint8
    landmarks: (N, 2) or (956,) numpy array
    """
    processed_patch = patch.copy()
    processed_landmarks = landmarks.copy()
    
    # 1. Resolution Dropout (16x16 -> 8x8)
    if apply_dropout and np.random.rand() > 0.5:
        h, w = processed_patch.shape[1], processed_patch.shape[2]
        for i in range(4):
            # Simulated distance: downsample and upsample
            small = cv2.resize(processed_patch[i], (8, 8), interpolation=cv2.INTER_CUBIC)
            processed_patch[i] = cv2.resize(small, (w, h), interpolation=cv2.INTER_CUBIC)

    # 2. Edge-Preserving Denoising (Bilateral Filter)
    if apply_denoise:
        for i in range(4):
            processed_patch[i] = cv2.bilateralFilter(processed_patch[i], 5, 20, 20)

    # 3. Tuned CLAHE
    if apply_clahe:
        clahe = cv2.createCLAHE(clipLimit=1.1, tileGridSize=(4, 4))
        for i in range(4):
            processed_patch[i] = clahe.apply(processed_patch[i])

    # 4. Random Gaussian Noise
    if apply_noise and np.random.rand() > 0.5:
        noise = np.random.normal(0, 3, processed_patch.shape).astype(np.float32)
        processed_patch = np.clip(processed_patch.astype(np.float32) + noise, 0, 255).astype(np.uint8)

    # 5. Landmark Jitter
    if apply_jitter and np.random.rand() > 0.5:
        noise_lm = np.random.normal(0, 0.003, processed_landmarks.shape).astype(np.float32)
        processed_landmarks = processed_landmarks + noise_lm

    return processed_patch, processed_landmarks

def preprocess_for_model(patch, landmarks):
    """
    Final conversion to torch tensors and normalization.
    """
    patch_tensor = torch.from_numpy(patch).float() / 255.0
    landmarks_tensor = torch.from_numpy(landmarks).float().view(-1)
    return patch_tensor, landmarks_tensor