import torch
import torch.nn.functional as F
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

def color_anisotropy_spectra(image_path, output_path="spectra.png", patch_size=3):
    """
    Compute a 2D spectra map of a colour image using the middle singular value
    of the local 3x3 colour covariance matrix (Singular Value Triple Ellipsoid).
    
    Args:
        image_path: path to input RGB image
        output_path: where to save the resulting grayscale map
        patch_size: should be 3 to form a 3x3 matrix (9 pixels)
    """
    # Load image and convert to float tensor [0,1], shape (C, H, W)
    img = Image.open(image_path).convert('RGB')
    img_tensor = torch.from_numpy(np.array(img)).float() / 255.0
    img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0)  # (1,3,H,W)
    
    C, H, W = img_tensor.shape[1], img_tensor.shape[2], img_tensor.shape[3]
    
    # Pad to keep the same spatial dimensions (mirror padding avoids border artefacts)
    pad = patch_size // 2
    img_pad = F.pad(img_tensor, (pad, pad, pad, pad), mode='reflect')
    
    # Extract all 3x3 patches as (9,3) matrices using unfold
    patches = img_pad.unfold(2, patch_size, 1).unfold(3, patch_size, 1)
    # patches shape: (1, 3, H, W, patch_size, patch_size)
    # Reshape to (H, W, 9, 3): each patch has 9 pixels, each with 3 colour channels
    patches = patches.permute(0, 2, 3, 1, 4, 5).reshape(1, H, W, 9, 3)
    patches = patches.squeeze(0)  # (H, W, 9, 3)
    
    # Compute the 3x3 covariance matrix for each patch
    # Center the data (subtract mean of the 9 pixels per patch)
    mean = patches.mean(dim=2, keepdim=True)  # (H, W, 1, 3)
    centered = patches - mean                 # (H, W, 9, 3)
    # Cov = (1/9) * (centered^T @ centered) -> (H, W, 3, 3)
    cov = torch.einsum('hwij,hwjk->hwik', centered.transpose(-2, -1), centered) / 9.0
    
    # Compute eigenvalues (or singular values) of each 3x3 symmetric matrix
    # We need the middle singular value = sqrt(middle eigenvalue) for PSD matrices
    eigvals, _ = torch.linalg.eigh(cov)  # eigenvalues sorted ascending: λ1 ≤ λ2 ≤ λ3
    # Clamp eigenvalues to 0.0 to prevent negative values due to numerical precision errors
    eigvals = torch.clamp(eigvals, min=0.0)
    # Middle singular value σ2 = sqrt(λ2)   (since λ ≥ 0)
    sigma2 = torch.sqrt(eigvals[..., 1])   # (H, W)
    
    # Normalise to [0,1] for visualisation (stretch to full range)
    vmin, vmax = sigma2.min(), sigma2.max()
    spectra = (sigma2 - vmin) / (vmax - vmin + 1e-8)
    spectra_np = spectra.cpu().numpy()
    
    # Save as grayscale image
    plt.imsave(output_path, spectra_np, cmap='gray')
    print(f"2D spectra saved to {output_path}")
    
    # Also show a side-by-side comparison
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
    ax1.imshow(np.array(img))
    ax1.set_title("Original RGB Image")
    ax1.axis('off')
    ax2.imshow(spectra_np, cmap='viridis')
    ax2.set_title("Color Anisotropy Spectra (σ₂ map)")
    ax2.axis('off')
    plt.tight_layout()
    plt.show()
    
    return spectra_np

# Example usage (uncomment and provide an image path)
#spectra_map = color_anisotropy_spectra("input.jpg", "output_spectra.png")
spectra_map = color_anisotropy_spectra("input.png", "output_spectra.png")