"""Sharpness subspace calibration: PCA on blur stimuli in FLUX VAE patch space."""

import math
from dataclasses import dataclass
from typing import Dict, Tuple

import torch
import torch.nn.functional as F
import comfy.utils

from .patchify import patchify


@dataclass
class SharpnessData:
    """Calibration data for the sharpness subspace.

    Produced by PCA on FLUX VAE-encoded images at varying blur levels.
    PC1 captures ~93% of sharpness/blur variance.
    """

    basis: torch.Tensor   # [64, K] PCA basis (columns), K typically 1-2
    mean: torch.Tensor    # [64] PCA mean
    pc1_std: float        # Standard deviation along PC1 (for normalizing strength)
    sign: float           # +1 or -1: ensures positive strength = sharper

    def to(self, device, dtype=None):
        """Move all tensors to device/dtype."""
        kw = {"device": device}
        if dtype is not None:
            kw["dtype"] = dtype
        return SharpnessData(
            basis=self.basis.to(**kw),
            mean=self.mean.to(**kw),
            pc1_std=self.pc1_std,
            sign=self.sign,
        )


# Cache for Gaussian kernels to avoid recomputation
_gaussian_kernel_cache: Dict[Tuple[int, float, torch.dtype, torch.device], torch.Tensor] = {}


def _gaussian_kernel_1d(kernel_size: int, sigma: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    """Create a 1D Gaussian kernel with caching."""
    cache_key = (kernel_size, sigma, dtype, device)
    if cache_key in _gaussian_kernel_cache:
        return _gaussian_kernel_cache[cache_key]

    x = torch.arange(kernel_size, dtype=dtype, device=device) - kernel_size // 2
    gauss = torch.exp(-x.pow(2) / (2 * sigma ** 2))
    gauss = gauss / gauss.sum()

    _gaussian_kernel_cache[cache_key] = gauss
    return gauss


def _apply_gaussian_blur(images: torch.Tensor, blur_sigma: float) -> torch.Tensor:
    """Apply Gaussian blur to a batch of images [B, C, H, W].

    Uses separable convolution for kernel_size > 15 (O(2k) vs O(k²)).

    Returns blurred images on same device/dtype as input.
    blur_sigma=0 returns input unchanged.
    """
    if blur_sigma < 1e-6:
        return images

    # Kernel size: 6*sigma rounded up to odd
    kernel_size = int(math.ceil(blur_sigma * 6)) | 1
    kernel_size = max(kernel_size, 3)

    B, C, H, W = images.shape
    device, dtype = images.device, images.dtype

    if kernel_size <= 15:
        # Direct 2D convolution for small kernels
        kernel_1d = _gaussian_kernel_1d(kernel_size, blur_sigma, device, dtype)
        kernel_2d = kernel_1d[:, None] @ kernel_1d[None, :]  # outer product
        kernel = kernel_2d.view(1, 1, kernel_size, kernel_size).expand(C, 1, -1, -1)
        pad = kernel_size // 2
        images_grouped = images.view(B * C, 1, H, W)
        blurred = F.conv2d(images_grouped, kernel, padding=pad, groups=1)
        return blurred.view(B, C, H, W)
    else:
        # Separable convolution: apply 1D Gaussian in X then Y direction
        # This reduces O(k²) to O(2k) operations
        gauss_1d = _gaussian_kernel_1d(kernel_size, blur_sigma, device, dtype)
        pad = kernel_size // 2

        # Horizontal pass: convolve along W dimension
        kernel_h = gauss_1d.view(1, 1, 1, kernel_size).expand(C, 1, 1, kernel_size)
        blurred = F.conv2d(images, kernel_h, padding=(0, pad), groups=C)

        # Vertical pass: convolve along H dimension
        kernel_v = gauss_1d.view(1, 1, kernel_size, 1).expand(C, 1, kernel_size, 1)
        blurred = F.conv2d(blurred, kernel_v, padding=(pad, 0), groups=C)

        return blurred


def calibrate_sharpness(vae, num_samples: int = 64, image_size: int = 512,
                        blur_levels: Tuple[float, ...] = (0, 1, 2, 4, 8, 16),
                        batch_size: int = 8) -> SharpnessData:
    """Compute sharpness subspace data (PCA basis, mean, sign) from FLUX VAE.

    1. Generate num_samples random noise images (spatial detail needed for blur)
    2. For each blur level, apply Gaussian blur to the SAME images
    3. VAE encode → patchify → average patches → [64] vector
    4. PCA on all vectors → extract PC1 (+ PC2)
    5. Determine sign: positive strength = sharper
    6. Compute pc1_std from spread of PC1 scores

    Returns: SharpnessData
    """
    n_levels = len(blur_levels)
    total_images = num_samples * n_levels

    print(f"\n[LCS Sharpness Calibration] Starting: {num_samples} images × {n_levels} blur levels = {total_images} samples")
    print(f"[LCS Sharpness Calibration] Blur sigmas: {list(blur_levels)}")

    # Deterministic seed for reproducibility
    rng = torch.Generator().manual_seed(42)

    # Step 1: Generate all base images upfront [num_samples, 3, H, W]
    # Use noise (not solid color) — spatial detail is needed for blur to have effect.
    # Per-batch generation to limit memory.
    print(f"[LCS Sharpness Calibration] Generating {num_samples} noise images...")
    base_images = torch.rand(num_samples, 3, image_size, image_size, generator=rng)

    # Step 2+3: For each blur level, apply blur to ALL base images, then encode
    vectors = []
    blur_labels = []  # track blur sigma per vector for sign determination
    pbar = comfy.utils.ProgressBar(total_images)

    for blur_sigma in blur_levels:
        print(f"[LCS Sharpness Calibration] Processing blur σ={blur_sigma}...")

        for batch_start in range(0, num_samples, batch_size):
            batch_end = min(batch_start + batch_size, num_samples)
            batch = base_images[batch_start:batch_end]  # [B, 3, H, W]
            actual_batch = batch.shape[0]

            # Apply blur (no-op for sigma=0)
            blurred = _apply_gaussian_blur(batch, blur_sigma)

            # Convert BCHW → BHWC for ComfyUI VAE
            imgs_bhwc = blurred.permute(0, 2, 3, 1).contiguous().cpu()

            # VAE encode → [B, 16, H/8, W/8]
            latent = vae.encode(imgs_bhwc)

            # Patchify → [B, L, 64], average across patches → [B, 64]
            patches, _, _ = patchify(latent)
            avg = patches.mean(dim=1).cpu()

            vectors.extend(avg.unbind(0))
            blur_labels.extend([blur_sigma] * actual_batch)

            pbar.update(actual_batch)

    # Stack all vectors: [N, 64]
    X = torch.stack(vectors, dim=0).float()
    blur_labels_t = torch.tensor(blur_labels, dtype=torch.float32)
    print(f"[LCS Sharpness Calibration] Collected {X.shape[0]} vectors of dimension {X.shape[1]}")

    # Step 4: PCA
    print("[LCS Sharpness Calibration] Computing PCA...")
    mean = X.mean(dim=0)  # [64]
    X_centered = X - mean
    U, S, Vh = torch.linalg.svd(X_centered, full_matrices=False)
    # Top 2 components
    basis = Vh[:2].T  # [64, 2]

    # Variance explained
    total_var = (S ** 2).sum()
    explained = (S[:2] ** 2) / total_var
    print(f"[LCS Sharpness Calibration] PC1: {explained[0]:.1%}, PC2: {explained[1]:.1%} ({(explained[0]+explained[1]):.1%} total)")

    # Step 5: Determine sign convention
    # Project all vectors onto PC1
    pc1_scores = X_centered @ basis[:, 0]  # [N]

    # Correlate PC1 score with blur sigma
    # If positive correlation (more blur = higher score), flip sign
    correlation = torch.corrcoef(torch.stack([pc1_scores, blur_labels_t]))[0, 1]
    sign = -1.0 if correlation > 0 else 1.0
    print(f"[LCS Sharpness Calibration] PC1-blur correlation: {correlation:.3f} → sign = {sign:+.0f}")

    # Step 6: Compute pc1_std
    pc1_std = float(pc1_scores.std())
    print(f"[LCS Sharpness Calibration] PC1 std: {pc1_std:.4f}")
    print(f"[LCS Sharpness Calibration] Complete! Basis shape: {basis.shape}")

    return SharpnessData(
        basis=basis,
        mean=mean,
        pc1_std=pc1_std,
        sign=sign,
    )