"""
Interactive Image Sampling and Quantization Demo
An educational tool for graduate-level image analysis courses
"""

import io
import numpy as np
import cv2 as cv
import streamlit as st
from PIL import Image
from huggingface_hub import hf_hub_download


@st.cache_resource
def load_sample_image():
    """Load a sample image for the demo. Falls back to generated image if download fails."""
    try:
        # Try to download from HuggingFace
        image_path = hf_hub_download(
            repo_id="amithjkamath/exampleimages",
            filename="sample-image.jpg",
            repo_type="dataset",
        )
        img = cv.imread(image_path)
        img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
    except Exception:
        # Generate a sample image with varied content
        img = generate_sample_image()
    
    return img


def generate_sample_image(size=512):
    """Generate a sample image with interesting features for demonstration."""
    img = np.zeros((size, size, 3), dtype=np.uint8)
    
    # Create a gradient background
    for i in range(size):
        img[i, :, 0] = int(255 * i / size)  # Red gradient
        img[:, i, 1] = int(255 * i / size)  # Green gradient
    
    # Add some geometric shapes
    cv.circle(img, (size//4, size//4), size//8, (255, 255, 255), -1)
    cv.rectangle(img, (size//2, size//2), (3*size//4, 3*size//4), (255, 0, 0), -1)
    cv.circle(img, (3*size//4, size//4), size//12, (0, 255, 255), -1)
    
    # Add some text
    cv.putText(img, "Sample", (size//4, 3*size//4), 
               cv.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 3)
    
    return img


def downsample_image(img, sampling_rate):
    """
    Downsample image by reducing spatial resolution.
    
    Args:
        img: Input image (RGB)
        sampling_rate: Factor by which to reduce resolution (1 = original, 2 = half, etc.)
    
    Returns:
        Downsampled image, upsampled back to original size for comparison
    """
    if sampling_rate == 1:
        return img
    
    h, w = img.shape[:2]
    new_h, new_w = h // sampling_rate, w // sampling_rate
    
    # Downsample using area interpolation (better quality)
    downsampled = cv.resize(img, (new_w, new_h), interpolation=cv.INTER_AREA)
    
    # Upsample back to original size using nearest neighbor (shows pixelation)
    upsampled = cv.resize(downsampled, (w, h), interpolation=cv.INTER_NEAREST)
    
    return upsampled


def quantize_image(img, bits_per_pixel):
    """
    Quantize image by reducing the number of bits per pixel.
    
    Args:
        img: Input image (RGB)
        bits_per_pixel: Number of bits per pixel (1-8)
    
    Returns:
        Quantized image
    """
    if bits_per_pixel == 8:
        return img
    
    # Calculate number of levels
    num_levels = 2 ** bits_per_pixel
    
    # Quantize by dividing into levels
    quantized = np.floor(img / (256.0 / num_levels)) * (256.0 / num_levels)
    quantized = np.clip(quantized, 0, 255).astype(np.uint8)
    
    return quantized


def apply_sampling_and_quantization(img, sampling_rate, bits_per_pixel):
    """Apply both sampling and quantization to an image."""
    # First downsample
    sampled = downsample_image(img, sampling_rate)
    
    # Then quantize
    result = quantize_image(sampled, bits_per_pixel)
    
    return result


def calculate_file_size(img_shape, sampling_rate, bits_per_pixel, compression_type="none"):
    """
    Calculate estimated file size.
    
    Args:
        img_shape: Shape of the original image (h, w, c)
        sampling_rate: Downsampling factor
        bits_per_pixel: Bits per pixel per channel
        compression_type: "none", "png", or "jpeg"
    
    Returns:
        File size in bytes
    """
    h, w, c = img_shape
    
    # Calculate actual number of pixels after sampling
    num_pixels = (h // sampling_rate) * (w // sampling_rate)
    
    # Calculate raw size in bytes
    raw_size = num_pixels * c * bits_per_pixel / 8
    
    # Apply compression estimate
    if compression_type == "png":
        # PNG typically achieves 60-80% of raw size for typical images
        size = raw_size * 0.7
    elif compression_type == "jpeg":
        # JPEG can achieve much better compression (20-40% of raw)
        size = raw_size * 0.3
    else:
        size = raw_size
    
    return int(size)


def format_file_size(size_bytes):
    """Format file size in human-readable format."""
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024 * 1024:
        return f"{size_bytes / 1024:.2f} KB"
    else:
        return f"{size_bytes / (1024 * 1024):.2f} MB"


def compress_image_jpeg(img, quality=50):
    """Compress image using JPEG and return the result."""
    # Convert to BGR for OpenCV
    img_bgr = cv.cvtColor(img, cv.COLOR_RGB2BGR)
    
    # Encode as JPEG
    encode_param = [int(cv.IMWRITE_JPEG_QUALITY), quality]
    _, buffer = cv.imencode('.jpg', img_bgr, encode_param)
    
    # Decode back
    img_decoded = cv.imdecode(buffer, cv.IMREAD_COLOR)
    img_rgb = cv.cvtColor(img_decoded, cv.COLOR_BGR2RGB)
    
    return img_rgb, len(buffer)


def compress_image_png(img, compression_level=6):
    """Compress image using PNG and return the result."""
    # Convert to BGR for OpenCV
    img_bgr = cv.cvtColor(img, cv.COLOR_RGB2BGR)
    
    # Encode as PNG
    encode_param = [int(cv.IMWRITE_PNG_COMPRESSION), compression_level]
    _, buffer = cv.imencode('.png', img_bgr, encode_param)
    
    # Decode back
    img_decoded = cv.imdecode(buffer, cv.IMREAD_COLOR)
    img_rgb = cv.cvtColor(img_decoded, cv.COLOR_BGR2RGB)
    
    return img_rgb, len(buffer)


def main_loop():
    """Main application loop."""
    st.set_page_config(layout="wide", page_title="Sampling & Quantization Demo")
    
    st.title("Interactive Image Sampling and Quantization Demo")
    st.markdown("""
    Welcome! This interactive demo teaches fundamental concepts in digital image processing.
    Explore how **sampling** (spatial resolution) and **quantization** (bit depth) affect 
    image quality and storage requirements.
    """)
    
    # Load sample image
    sample_img = load_sample_image()
    
    # Option to upload custom image
    st.sidebar.header("Image Input")
    uploaded_file = st.sidebar.file_uploader("Upload your own image (optional)", 
                                             type=['png', 'jpg', 'jpeg'])
    
    if uploaded_file is not None:
        # Use uploaded image
        file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
        img = cv.imdecode(file_bytes, cv.IMREAD_COLOR)
        img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
        
        # Resize if too large
        max_size = 512
        h, w = img.shape[:2]
        if max(h, w) > max_size:
            scale = max_size / max(h, w)
            new_w, new_h = int(w * scale), int(h * scale)
            img = cv.resize(img, (new_w, new_h), interpolation=cv.INTER_AREA)
    else:
        img = sample_img
    
    # Main controls
    st.sidebar.header("Controls")
    
    st.sidebar.subheader("Spatial Sampling")
    sampling_rate = st.sidebar.slider(
        "Sampling Grid Size (pixels)",
        min_value=1,
        max_value=16,
        value=1,
        step=1,
        help="Higher values = more pixelated image (fewer pixels stored)"
    )
    
    st.sidebar.subheader("Quantization")
    bits_per_pixel = st.sidebar.slider(
        "Bits per Pixel per Channel",
        min_value=1,
        max_value=8,
        value=8,
        step=1,
        help="Lower values = fewer colors/gray levels (less storage per pixel)"
    )
    
    # Calculate number of possible values
    num_levels = 2 ** bits_per_pixel
    st.sidebar.info(f"**{num_levels}** intensity levels per channel\n\n"
                   f"**{num_levels**3:,}** total colors possible")
    
    # Process image
    processed_img = apply_sampling_and_quantization(img, sampling_rate, bits_per_pixel)
    
    # Display images
    st.markdown("---")
    st.markdown("## Visual Comparison")
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("### Original Image")
        st.image(img, use_column_width=True)
        st.caption(f"Size: {img.shape[1]}x{img.shape[0]} pixels, 8 bits/channel")
    
    with col2:
        st.markdown("### Processed Image")
        st.image(processed_img, use_column_width=True)
        st.caption(f"Size: {img.shape[1]//sampling_rate}x{img.shape[0]//sampling_rate} pixels, "
                  f"{bits_per_pixel} bits/channel")
    
    # File size analysis
    st.markdown("---")
    st.markdown("## Storage Analysis")
    
    col1, col2, col3 = st.columns(3)
    
    original_size = calculate_file_size(img.shape, 1, 8, "none")
    processed_size = calculate_file_size(img.shape, sampling_rate, bits_per_pixel, "none")
    reduction = (1 - processed_size / original_size) * 100
    
    with col1:
        st.metric("Original (Uncompressed)", format_file_size(original_size))
    
    with col2:
        st.metric("Processed (Uncompressed)", format_file_size(processed_size))
    
    with col3:
        st.metric("Size Reduction", f"{reduction:.1f}%")
    
    # Detailed breakdown
    with st.expander("Size Calculation Details"):
        st.markdown(f"""
        **Original Image:**
        - Dimensions: {img.shape[1]} x {img.shape[0]} pixels
        - Channels: 3 (RGB)
        - Bits per pixel: 8 x 3 = 24 bits
        - Total bits: {img.shape[1]} x {img.shape[0]} x 24 = {img.shape[1] * img.shape[0] * 24:,} bits
        - **Uncompressed size: {format_file_size(original_size)}**
        
        **Processed Image:**
        - Dimensions: {img.shape[1]//sampling_rate} x {img.shape[0]//sampling_rate} pixels
        - Channels: 3 (RGB)
        - Bits per pixel: {bits_per_pixel} x 3 = {bits_per_pixel * 3} bits
        - Total bits: {img.shape[1]//sampling_rate} x {img.shape[0]//sampling_rate} x {bits_per_pixel * 3} = {(img.shape[1]//sampling_rate) * (img.shape[0]//sampling_rate) * bits_per_pixel * 3:,} bits
        - **Uncompressed size: {format_file_size(processed_size)}**
        """)
    
    # Compression comparison
    st.markdown("---")
    st.markdown("## Compression Methods Comparison")
    
    st.markdown("""
    Now see how different compression algorithms affect the processed image.
    **PNG** uses lossless compression, while **JPEG** uses lossy compression.
    """)
    
    # JPEG compression
    col1, col2 = st.columns([1, 3])
    
    with col1:
        st.subheader("JPEG Settings")
        jpeg_quality = st.slider(
            "JPEG Quality",
            min_value=1,
            max_value=100,
            value=50,
            help="Lower quality = more compression = smaller file = more artifacts"
        )
    
    # Compress images
    jpeg_img, jpeg_size = compress_image_jpeg(processed_img, jpeg_quality)
    png_img, png_size = compress_image_png(processed_img, compression_level=6)
    
    with col2:
        st.markdown("### Compression Results")
        
        col_png, col_jpeg = st.columns(2)
        
        with col_png:
            st.markdown("**PNG (Lossless)**")
            st.image(png_img, use_column_width=True)
            st.metric("PNG File Size", format_file_size(png_size))
            st.caption("Exact reconstruction, no quality loss")
        
        with col_jpeg:
            st.markdown("**JPEG (Lossy)**")
            st.image(jpeg_img, use_column_width=True)
            st.metric("JPEG File Size", format_file_size(jpeg_size))
            compression_ratio = (1 - jpeg_size / png_size) * 100
            st.caption(f"{compression_ratio:.1f}% smaller than PNG")
    
    # Show blocking artifacts
    if jpeg_quality < 30:
        st.warning("**Low JPEG quality detected!** Look closely at the image to see blocking artifacts.")
    
    # Educational section
    st.markdown("---")
    st.markdown("## Educational Insights")
    
    tab1, tab2, tab3 = st.tabs(["Sampling", "Quantization", "Compression"])
    
    with tab1:
        st.markdown("""
        ### Spatial Sampling
        
        **What is it?**
        - Sampling determines the spatial resolution of an image
        - A sampling rate of N means we keep every Nth pixel in each direction
        - This reduces the total number of pixels by a factor of N²
        
        **Key Concepts:**
        - **Nyquist-Shannon Sampling Theorem**: To avoid aliasing, sampling rate must be at least 
          twice the highest frequency in the image
        - **Pixelation**: When sampling rate is too low, fine details are lost and edges become blocky
        - **Storage Impact**: Directly proportional to pixel count
        
        **Try it:**
        - Increase the sampling grid size slider above
        - Notice how the image becomes more pixelated
        - Watch the file size decrease as fewer pixels are stored
        """)
        
        if sampling_rate > 1:
            st.info(f"Current sampling reduces pixel count by {sampling_rate**2}x "
                   f"({img.shape[0]*img.shape[1]:,} to {(img.shape[0]//sampling_rate)*(img.shape[1]//sampling_rate):,} pixels)")
    
    with tab2:
        st.markdown("""
        ### Quantization (Bit Depth)
        
        **What is it?**
        - Quantization determines how many distinct values each pixel can have
        - With N bits per pixel per channel, we can represent 2^N different intensity levels
        - This affects color depth and tonal range
        
        **Key Concepts:**
        - **8 bits** = 256 levels per channel = 16.7 million colors (standard)
        - **4 bits** = 16 levels per channel = 4,096 colors
        - **1 bit** = 2 levels per channel = 8 colors (effectively binary)
        - **Posterization**: Visible bands in gradients when quantization is too coarse
        
        **Storage Impact:**
        - Each pixel requires: bits_per_channel x number_of_channels bits
        - For RGB: 3 x bits_per_pixel bits per pixel
        
        **Try it:**
        - Decrease the bits per pixel slider above
        - Notice color banding in smooth gradients
        - Watch file size decrease as fewer bits are used per pixel
        """)
        
        if bits_per_pixel < 8:
            st.info(f"Current quantization uses {bits_per_pixel * 3} bits per pixel "
                   f"(vs 24 bits normally), saving {(1 - bits_per_pixel/8)*100:.0f}% per pixel")
    
    with tab3:
        st.markdown("""
        ### Image Compression
        
        **PNG (Portable Network Graphics) - Lossless**
        - Uses DEFLATE compression algorithm (similar to ZIP)
        - Exploits spatial redundancy in images
        - Perfect reconstruction - no quality loss
        - Better for graphics, text, screenshots
        - Typically 50-80% of raw size
        
        - **JPEG (Joint Photographic Experts Group) - Lossy**
        - Uses Discrete Cosine Transform (DCT) on 8x8 blocks
        - Quantizes frequency components (loses information)
        - Much better compression ratios (10-40% of raw size)
        - Better for photographs with gradual color changes
        - **Blocking Artifacts**: Visible 8x8 blocks at low quality
        
        **Quality vs. Size Trade-off:**
        - High JPEG quality (90-100): Minimal artifacts, larger files
        - Medium quality (50-70): Good balance for photos
        - Low quality (1-30): Heavy artifacts, smallest files
        
        **Try it:**
        - Adjust the JPEG quality slider above
        - At low quality (<30), look for 8x8 blocking patterns
        - Compare file sizes: JPEG can be 5-10x smaller than PNG
        """)
        
        st.info(f"Current JPEG is {(png_size / jpeg_size):.1f}x smaller than PNG, "
               f"and {(processed_size / jpeg_size):.1f}x smaller than uncompressed")
    
    # Footer
    st.markdown("---")
    st.markdown("""
    <small>
    Educational Demo for Image Analysis Courses | 
    Built with Streamlit | 
    <a href="https://github.com/ubern-image-analysis/sampling-quantization" target="_blank">View Source</a>
    </small>
    """, unsafe_allow_html=True)


if __name__ == "__main__":
    main_loop()