import streamlit as st
import numpy as np
import cv2
from PIL import Image
import matplotlib.pyplot as plt
from matplotlib import cm
from skimage import data
from skimage.transform import hough_line, hough_line_peaks
from skimage.feature import canny
from skimage.draw import line as draw_line
import io

# Page configuration
st.set_page_config(
    page_title="Hough Transform & RANSAC Line Detection",
    page_icon="📐",
    layout="wide"
)

# Title and introduction
st.title("📐 Hough Transform & RANSAC Line Detection")
st.markdown("""
This interactive tool helps you understand how **Hough Transform** and **RANSAC** detect lines in images.
Experiment with parameters to build intuition about how these algorithms work!
""")

# Sidebar for image selection
st.sidebar.header("1️⃣ Image Selection")

# Example images
example_images = {
    "Camera": data.camera(),
    "Coins": data.coins(),
    "Checkerboard": data.checkerboard(),
    "Brick": data.brick(),
    "Text": data.text(),
}

image_source = st.sidebar.radio("Choose image source:", ["Example Images", "Upload Image"])

if image_source == "Example Images":
    selected_example = st.sidebar.selectbox("Select example:", list(example_images.keys()))
    image = example_images[selected_example]
else:
    uploaded_file = st.sidebar.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
    if uploaded_file is not None:
        image = np.array(Image.open(uploaded_file).convert('L'))
    else:
        st.sidebar.info("Please upload an image or select an example.")
        image = example_images["Camera"]

# Edge detection parameters
st.sidebar.header("2️⃣ Edge Detection")
st.sidebar.markdown("*Canny edge detector extracts edge points*")

sigma = st.sidebar.slider("Gaussian sigma (blur)", 1.0, 5.0, 2.0, 0.5)
low_threshold = st.sidebar.slider("Low threshold", 0, 100, 20, 5)
high_threshold = st.sidebar.slider("High threshold", 0, 255, 50, 5)

# Apply Canny edge detection
edges = canny(image, sigma=sigma, low_threshold=low_threshold, high_threshold=high_threshold)

# Main content area with tabs
tab1, tab2, tab3 = st.tabs(["🔍 Overview", "📊 Hough Transform", "🎯 RANSAC"])

with tab1:
    st.header("Image Processing Pipeline")
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("Original Image")
        fig, ax = plt.subplots(figsize=(6, 6))
        ax.imshow(image, cmap='gray')
        ax.axis('off')
        st.pyplot(fig)
        plt.close()
    
    with col2:
        st.subheader("Edge Detection (Canny)")
        fig, ax = plt.subplots(figsize=(6, 6))
        ax.imshow(edges, cmap='gray')
        ax.axis('off')
        st.pyplot(fig)
        plt.close()
    
    st.info("""
    **Edge Detection** identifies points where intensity changes rapidly. 
    These edge points are candidates for line detection algorithms.
    """)

with tab2:
    st.header("Hough Transform for Line Detection")
    
    st.markdown("""
    ### How it works:
    1. **Parameter Space**: Lines are represented as ρ = x·cos(θ) + y·sin(θ)
    2. **Voting**: Each edge point votes for all possible lines passing through it
    3. **Accumulator Array (Sinogram)**: Peaks indicate detected lines
    """)
    
    # Hough Transform parameters
    col1, col2 = st.columns([1, 2])
    
    with col1:
        st.subheader("Parameters")
        
        theta_res = st.slider(
            "Theta resolution (degrees)", 
            0.1, 5.0, 1.0, 0.1,
            help="Angular resolution in Hough space"
        )
        
        num_peaks = st.slider(
            "Number of lines to detect", 
            1, 20, 5, 1,
            help="Top N peaks in accumulator array"
        )
        
        threshold_percentile = st.slider(
            "Threshold (percentile)", 
            50, 99, 85, 1,
            help="Minimum votes needed (as percentile of max)"
        )
        
        min_distance = st.slider(
            "Min peak distance", 
            5, 50, 20, 5,
            help="Minimum distance between peaks in Hough space"
        )
        
        min_angle = st.slider(
            "Min angle distance (degrees)", 
            1, 45, 15, 1,
            help="Minimum angle separation between lines"
        )
    
    # Perform Hough Transform
    tested_angles = np.deg2rad(np.arange(0, 180, theta_res))
    h, theta, d = hough_line(edges, theta=tested_angles)
    
    # Find peaks
    hough_threshold = np.percentile(h, threshold_percentile)
    h_peaks, angles, dists = hough_line_peaks(
        h, theta, d, 
        min_distance=min_distance, 
        min_angle=min_angle,
        threshold=hough_threshold,
        num_peaks=num_peaks
    )
    
    with col2:
        st.subheader("Visualizations")
        
        # Create visualization with detected lines
        fig, axes = plt.subplots(1, 2, figsize=(12, 5))
        
        # Detected lines on edges
        ax = axes[0]
        ax.imshow(edges, cmap='gray')
        
        for _, angle, dist in zip(h_peaks, angles, dists):
            y0, y1 = 0, edges.shape[0]
            if np.abs(np.sin(angle)) > 1e-10:
                x0 = (dist - y0 * np.sin(angle)) / np.cos(angle)
                x1 = (dist - y1 * np.sin(angle)) / np.cos(angle)
                ax.plot([x0, x1], [y0, y1], 'r-', linewidth=2, alpha=0.7)
        
        ax.set_title(f'Detected Lines (n={len(angles)})')
        ax.axis('off')
        
        # Hough accumulator (sinogram)
        ax = axes[1]
        im = ax.imshow(
            np.log(1 + h),
            extent=[np.rad2deg(theta[0]), np.rad2deg(theta[-1]), d[-1], d[0]],
            cmap='hot',
            aspect='auto'
        )
        ax.scatter(np.rad2deg(angles), dists, s=100, c='cyan', marker='x', linewidths=3)
        ax.set_xlabel('Theta (degrees)')
        ax.set_ylabel('Rho (pixels)')
        ax.set_title('Hough Space (Sinogram)')
        plt.colorbar(im, ax=ax, label='Log(Votes)')
        
        plt.tight_layout()
        st.pyplot(fig)
        plt.close()
    
    # Educational explanations
    st.markdown("---")
    st.subheader("📚 Understanding the Sinogram")
    
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.markdown("""
        **Theta (θ)**
        - Angle of line normal (0-180°)
        - Horizontal resolution = your theta resolution setting
        """)
    
    with col2:
        st.markdown("""
        **Rho (ρ)**
        - Distance from origin to line
        - Vertical axis in sinogram
        - Range: [-diagonal, +diagonal]
        """)
    
    with col3:
        st.markdown("""
        **Bright Spots (Peaks)**
        - High vote counts
        - Each peak = one detected line
        - Cyan X marks = selected peaks
        """)
    
    st.info(f"""
    **Current Results**: Detected **{len(angles)}** lines from **{np.sum(edges)}** edge points.
    The sinogram shows **{h.shape[0]} × {h.shape[1]}** bins (ρ × θ).
    """)

with tab3:
    st.header("RANSAC Line Fitting")
    
    st.markdown("""
    ### How it works:
    1. **Random Sampling**: Pick 2 random edge points
    2. **Model Fitting**: Fit a line through these points
    3. **Consensus**: Count inliers (points close to the line)
    4. **Iteration**: Repeat and keep the best model
    """)
    
    # Get edge points
    edge_points = np.column_stack(np.where(edges))
    
    if len(edge_points) < 2:
        st.warning("Not enough edge points detected. Adjust edge detection parameters.")
    else:
        col1, col2 = st.columns([1, 2])
        
        with col1:
            st.subheader("Parameters")
            
            ransac_iterations = st.slider(
                "Number of iterations",
                100, 5000, 1000, 100,
                help="More iterations = higher chance of finding best fit"
            )
            
            ransac_threshold = st.slider(
                "Distance threshold (pixels)",
                1.0, 10.0, 3.0, 0.5,
                help="Max distance for a point to be an inlier"
            )
            
            min_inliers = st.slider(
                "Minimum inliers",
                10, 200, 50, 10,
                help="Minimum points needed for valid line"
            )
            
            num_lines_ransac = st.slider(
                "Number of lines (RANSAC)",
                1, 10, 3, 1,
                help="How many lines to detect sequentially"
            )
        
        # RANSAC implementation
        def ransac_line(points, iterations, threshold, min_inliers):
            """Fit a line using RANSAC"""
            best_inliers = []
            best_model = None
            
            for _ in range(iterations):
                # Random sample
                idx = np.random.choice(len(points), 2, replace=False)
                p1, p2 = points[idx]
                
                # Fit line: ax + by + c = 0
                if p1[1] == p2[1]:  # Vertical line
                    continue
                    
                # Calculate line parameters
                dx = p2[1] - p1[1]
                dy = p2[0] - p1[0]
                
                if dx == 0 and dy == 0:
                    continue
                
                # Normal form
                norm = np.sqrt(dx**2 + dy**2)
                a = -dy / norm
                b = dx / norm
                c = -(a * p1[1] + b * p1[0])
                
                # Calculate distances
                distances = np.abs(a * points[:, 1] + b * points[:, 0] + c)
                inliers = distances < threshold
                
                if np.sum(inliers) > len(best_inliers):
                    best_inliers = inliers
                    best_model = (a, b, c)
            
            if len(best_inliers) >= min_inliers:
                return best_model, best_inliers
            return None, []
        
        # Detect multiple lines
        remaining_points = edge_points.copy()
        detected_lines = []
        all_inliers = []
        
        for i in range(num_lines_ransac):
            if len(remaining_points) < min_inliers:
                break
                
            model, inliers_mask = ransac_line(
                remaining_points, 
                ransac_iterations, 
                ransac_threshold, 
                min_inliers
            )
            
            if model is not None:
                detected_lines.append(model)
                inlier_points = remaining_points[inliers_mask]
                all_inliers.append(inlier_points)
                
                # Remove inliers for next iteration
                remaining_points = remaining_points[~inliers_mask]
        
        with col2:
            st.subheader("Visualizations")
            
            fig, axes = plt.subplots(1, 2, figsize=(12, 5))
            
            # RANSAC detected lines
            ax = axes[0]
            ax.imshow(edges, cmap='gray')
            
            colors = plt.cm.rainbow(np.linspace(0, 1, len(detected_lines)))
            
            for (a, b, c), color in zip(detected_lines, colors):
                y0, y1 = 0, edges.shape[0]
                if abs(a) > 1e-10:
                    x0 = -(b * y0 + c) / a
                    x1 = -(b * y1 + c) / a
                else:
                    x0 = -c / b
                    x1 = -c / b
                ax.plot([x0, x1], [y0, y1], color=color, linewidth=2, alpha=0.8)
            
            ax.set_title(f'RANSAC Lines (n={len(detected_lines)})')
            ax.axis('off')
            
            # Show inliers/outliers
            ax = axes[1]
            ax.imshow(image, cmap='gray', alpha=0.3)
            
            # Plot all edge points as outliers (gray)
            if len(remaining_points) > 0:
                ax.scatter(remaining_points[:, 1], remaining_points[:, 0], 
                          c='gray', s=1, alpha=0.5, label='Outliers')
            
            # Plot inliers for each line
            for i, (inliers, color) in enumerate(zip(all_inliers, colors)):
                if len(inliers) > 0:
                    ax.scatter(inliers[:, 1], inliers[:, 0], 
                              c=[color], s=2, alpha=0.8, label=f'Line {i+1} inliers')
            
            ax.set_title('Inliers vs Outliers')
            ax.axis('off')
            ax.legend(loc='upper right', fontsize=8)
            
            plt.tight_layout()
            st.pyplot(fig)
            plt.close()
        
        # Statistics
        st.markdown("---")
        st.subheader("📊 RANSAC Statistics")
        
        col1, col2, col3 = st.columns(3)
        
        with col1:
            total_inliers = sum(len(inliers) for inliers in all_inliers)
            st.metric("Total Edge Points", len(edge_points))
            st.metric("Points Used (Inliers)", total_inliers)
        
        with col2:
            st.metric("Lines Detected", len(detected_lines))
            if len(all_inliers) > 0:
                avg_inliers = np.mean([len(inliers) for inliers in all_inliers])
                st.metric("Avg Inliers per Line", f"{avg_inliers:.0f}")
        
        with col3:
            if len(edge_points) > 0:
                inlier_percentage = (total_inliers / len(edge_points)) * 100
                st.metric("Inlier Percentage", f"{inlier_percentage:.1f}%")
        
        st.info("""
        **Key Insight**: RANSAC is robust to outliers. Even if many edge points don't belong to 
        lines (e.g., noise, curves), RANSAC can still find the dominant linear structures.
        """)

# Comparison section
st.markdown("---")
st.header("🔄 Hough Transform vs RANSAC")

col1, col2 = st.columns(2)

with col1:
    st.subheader("Hough Transform")
    st.markdown("""
    **Strengths:**
    - Detects all lines simultaneously
    - Works well for multiple parallel lines
    - Global optimization approach
    - Good for complete line detection
    
    **Parameters:**
    - Angular resolution (θ)
    - Distance resolution (ρ)
    - Vote threshold
    """)

with col2:
    st.subheader("RANSAC")
    st.markdown("""
    **Strengths:**
    - Robust to outliers
    - Works with noisy data
    - Can fit partial lines
    - Good when lines are interrupted
    
    **Parameters:**
    - Number of iterations
    - Distance threshold
    - Minimum inliers
    """)

# Footer
st.markdown("---")
st.markdown("""
### 🎓 Educational Notes

**For Students**: 
- Try different edge detection parameters and observe how they affect line detection
- Compare how Hough Transform's sinogram changes with parameter adjustments
- Experiment with RANSAC parameters to see the trade-off between iterations and accuracy
- Notice how both methods handle noise and incomplete lines differently

**Tips for Exploration**:
1. Start with default parameters to see basic functionality
2. Increase theta resolution in Hough Transform for more precise angles
3. Increase RANSAC iterations for better line fitting
4. Try different images to see how algorithms perform on various structures
""")