# app.py - Hugging Face Spaces Carpet Weaving Map Generator
import gradio as gr
import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageFilter, ImageEnhance, ImageDraw, ImageFont
import io
from sklearn.cluster import KMeans
import pandas as pd
import tempfile
import os

class CarpetWeavingMap:
    def __init__(self):
        self.original_image = None
        self.knot_map = None
        self.color_palette = None
        self.grid_pattern = None
        
    def extract_color_palette(self, image, n_colors=16):
        """Extract dominant colors for carpet knots using K-means clustering"""
        # Convert image to RGB array
        img_array = np.array(image)
        pixels = img_array.reshape(-1, 3)
        
        # Use K-means to find dominant colors
        kmeans = KMeans(n_clusters=n_colors, random_state=42, n_init=10)
        kmeans.fit(pixels)
        
        # Get the colors and sort by frequency
        colors = kmeans.cluster_centers_.astype(int)
        labels = kmeans.labels_
        
        # Count frequency of each color
        unique_labels, counts = np.unique(labels, return_counts=True)
        
        # Sort colors by frequency (most common first)
        sorted_indices = np.argsort(counts)[::-1]
        self.color_palette = colors[sorted_indices]
        
        return self.color_palette
    
    def create_knot_grid(self, image, knots_per_cm=10, carpet_width_cm=100, carpet_height_cm=150):
        """Create a grid pattern showing individual knot colors"""
        # Calculate grid dimensions
        grid_width = carpet_width_cm * knots_per_cm
        grid_height = carpet_height_cm * knots_per_cm
        
        # Resize image to match grid dimensions
        resized_image = image.resize((grid_width, grid_height), Image.Resampling.LANCZOS)
        img_array = np.array(resized_image)
        
        # Extract color palette
        self.extract_color_palette(resized_image)
        
        # Create knot map by mapping each pixel to nearest palette color
        self.knot_map = np.zeros((grid_height, grid_width, 3), dtype=np.uint8)
        
        for y in range(grid_height):
            for x in range(grid_width):
                pixel = img_array[y, x]
                # Find closest color in palette
                distances = np.sqrt(np.sum((self.color_palette - pixel)**2, axis=1))
                closest_color_idx = np.argmin(distances)
                self.knot_map[y, x] = self.color_palette[closest_color_idx]
        
        self.knots_per_cm = knots_per_cm
        self.carpet_width_cm = carpet_width_cm
        self.carpet_height_cm = carpet_height_cm
        
        return self.knot_map
    
    def draw_grid_overlay(self, knot_map, grid_spacing=10):
        """Draw grid lines over the knot map"""
        grid_image = Image.fromarray(knot_map)
        draw = ImageDraw.Draw(grid_image)
        
        height, width = knot_map.shape[:2]
        
        # Draw vertical lines (every grid_spacing knots)
        for x in range(0, width, grid_spacing):
            draw.line([(x, 0), (x, height)], fill='black', width=1)
        
        # Draw horizontal lines (every grid_spacing knots)
        for y in range(0, height, grid_spacing):
            draw.line([(0, y), (width, y)], fill='black', width=1)
        
        return np.array(grid_image)
    
    def create_technical_specifications(self):
        """Create technical specifications dictionary"""
        specs = {
            'Carpet Dimensions': f"{self.carpet_width_cm} x {self.carpet_height_cm} cm",
            'Knot Density': f"{self.knots_per_cm} knots per cm",
            'Total Knots Width': f"{self.carpet_width_cm * self.knots_per_cm} knots",
            'Total Knots Height': f"{self.carpet_height_cm * self.knots_per_cm} knots",
            'Total Knots': f"{self.carpet_width_cm * self.carpet_height_cm * self.knots_per_cm**2:,} knots",
            'Color Palette Size': f"{len(self.color_palette)} colors",
            'Grid Sections (10x10)': f"{(self.carpet_width_cm * self.knots_per_cm)//10} x {(self.carpet_height_cm * self.knots_per_cm)//10}"
        }
        return specs
    
    def create_visualization(self, show_grid=True, grid_spacing=10):
        """Create the complete carpet weaving map visualization"""
        if self.knot_map is None:
            return None, None, None
        
        # Create figure with subplots
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(20, 15))
        
        # Main knot map
        if show_grid:
            display_map = self.draw_grid_overlay(self.knot_map, grid_spacing)
        else:
            display_map = self.knot_map
            
        ax1.imshow(display_map)
        ax1.set_title(f'Carpet Weaving Map ({self.knots_per_cm} knots/cm)', fontsize=16, fontweight='bold')
        ax1.set_xlabel('Knots (Width)')
        ax1.set_ylabel('Knots (Height)')
        
        # Add scale markings
        width_knots = self.carpet_width_cm * self.knots_per_cm
        height_knots = self.carpet_height_cm * self.knots_per_cm
        
        # X-axis ticks every 10 cm
        x_ticks = np.arange(0, width_knots + 1, 10 * self.knots_per_cm)
        x_labels = [f"{int(x/(self.knots_per_cm))}cm" for x in x_ticks]
        ax1.set_xticks(x_ticks)
        ax1.set_xticklabels(x_labels)
        
        # Y-axis ticks every 10 cm
        y_ticks = np.arange(0, height_knots + 1, 10 * self.knots_per_cm)
        y_labels = [f"{int(y/(self.knots_per_cm))}cm" for y in y_ticks]
        ax1.set_yticks(y_ticks)
        ax1.set_yticklabels(y_labels)
        
        # Original image
        ax2.imshow(self.original_image)
        ax2.set_title('Original Image', fontsize=14, fontweight='bold')
        ax2.axis('off')
        
        # Color palette
        self.display_color_palette_in_subplot(ax3)
        
        # Technical specifications
        self.display_specifications_in_subplot(ax4)
        
        plt.tight_layout()
        
        # Save the figure
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png', dir=tempfile.gettempdir())
        plt.savefig(temp_file.name, dpi=300, bbox_inches='tight')
        plt.close()
        
        return temp_file.name
    
    def display_color_palette_in_subplot(self, ax):
        """Display color palette in a subplot"""
        if self.color_palette is None:
            return
            
        n_colors = len(self.color_palette)
        cols = 4
        rows = (n_colors + cols - 1) // cols
        
        for i, color in enumerate(self.color_palette):
            row = i // cols
            col = i % cols
            
            rect = plt.Rectangle((col, rows - row - 1), 0.8, 0.8, 
                               facecolor=color/255.0, edgecolor='black', linewidth=1)
            ax.add_patch(rect)
            
            # Add color number
            ax.text(col + 0.4, rows - row - 0.5, f"C{i+1:02d}", 
                   ha='center', va='center', fontsize=10, fontweight='bold')
        
        ax.set_xlim(0, cols)
        ax.set_ylim(0, rows)
        ax.set_aspect('equal')
        ax.axis('off')
        ax.set_title('Color Palette', fontsize=12, fontweight='bold')
    
    def display_specifications_in_subplot(self, ax):
        """Display technical specifications in a subplot"""
        specs = self.create_technical_specifications()
        
        ax.axis('off')
        ax.set_title('Technical Specifications', fontsize=12, fontweight='bold')
        
        y_pos = 0.9
        for key, value in specs.items():
            ax.text(0.1, y_pos, f"{key}:", fontsize=10, fontweight='bold')
            ax.text(0.6, y_pos, value, fontsize=10)
            y_pos -= 0.12
        
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
    
    def export_pattern_files(self):
        """Export pattern data as downloadable files"""
        if self.knot_map is None:
            return None, None
        
        # Create color mapping
        color_map = {}
        for i, color in enumerate(self.color_palette):
            color_map[f"C{i+1:02d}"] = f"RGB({color[0]},{color[1]},{color[2]})"
        
        # Create pattern grid with color codes
        height, width = self.knot_map.shape[:2]
        pattern_grid = np.zeros((height, width), dtype='U4')
        
        for y in range(height):
            for x in range(width):
                pixel = self.knot_map[y, x]
                # Find matching color in palette
                for i, color in enumerate(self.color_palette):
                    if np.array_equal(pixel, color):
                        pattern_grid[y, x] = f"C{i+1:02d}"
                        break
        
        # Save pattern grid as CSV
        csv_file = tempfile.NamedTemporaryFile(delete=False, suffix='.csv', mode='w', dir=tempfile.gettempdir())
        np.savetxt(csv_file.name, pattern_grid, delimiter=',', fmt='%s')
        csv_file.close()
        
        # Save color mapping as text file
        txt_file = tempfile.NamedTemporaryFile(delete=False, suffix='.txt', mode='w', dir=tempfile.gettempdir())
        txt_file.write("Carpet Pattern Color Mapping\n")
        txt_file.write("=" * 30 + "\n\n")
        for code, rgb in color_map.items():
            txt_file.write(f"{code}: {rgb}\n")
        txt_file.write(f"\nSpecifications:\n")
        specs = self.create_technical_specifications()
        for key, value in specs.items():
            txt_file.write(f"{key}: {value}\n")
        txt_file.close()
        
        return csv_file.name, txt_file.name

def process_image(image, knots_per_cm, carpet_width_cm, carpet_height_cm, n_colors, show_grid):
    """Main processing function for Gradio interface"""
    if image is None:
        return None, None, None, "Please upload an image first."
    
    try:
        # Create mapper instance
        mapper = CarpetWeavingMap()
        mapper.original_image = Image.fromarray(image).convert('RGB')
        
        # Create knot grid
        mapper.create_knot_grid(
            mapper.original_image, 
            knots_per_cm=knots_per_cm,
            carpet_width_cm=carpet_width_cm, 
            carpet_height_cm=carpet_height_cm
        )
        
        # Extract color palette
        mapper.extract_color_palette(mapper.original_image, n_colors=n_colors)
        
        # Create visualization
        visualization_path = mapper.create_visualization(show_grid=show_grid)
        
        # Export pattern files
        csv_path, txt_path = mapper.export_pattern_files()
        
        # Create specifications text
        specs = mapper.create_technical_specifications()
        specs_text = "CARPET WEAVING SPECIFICATIONS:\n" + "="*40 + "\n"
        for key, value in specs.items():
            specs_text += f"{key}: {value}\n"
        
        return visualization_path, csv_path, txt_path, specs_text
        
    except Exception as e:
        return None, None, None, f"Error processing image: {str(e)}"

# Gradio Interface
def create_gradio_interface():
    """Create the Gradio interface"""
    
    title = "🧶 Carpet Weaving Map Generator"
    description = """
    Convert any image into a technical carpet weaving pattern with precise knot mapping and color specifications.
    
    **Features:**
    - Generate grid-based knot patterns
    - Customizable knot density (4-20 knots/cm)
    - Color palette extraction and mapping
    - Exportable CSV pattern files
    - Technical specifications for weavers
    
    **Instructions:**
    1. Upload your image
    2. Adjust carpet dimensions and knot density
    3. Set number of colors for the palette
    4. Download the pattern files for weaving
    """
    
    with gr.Blocks(title=title, theme=gr.themes.Soft()) as interface:
        gr.Markdown(f"# {title}")
        gr.Markdown(description)
        
        with gr.Row():
            with gr.Column(scale=1):
                # Input controls
                image_input = gr.Image(
                    label="Upload Image",
                    type="numpy"
                )
                
                with gr.Row():
                    knots_per_cm = gr.Slider(
                        minimum=4, maximum=20, value=10, step=1,
                        label="Knots per CM (density)"
                    )
                    n_colors = gr.Slider(
                        minimum=8, maximum=32, value=16, step=1,
                        label="Number of Colors"
                    )
                
                with gr.Row():
                    carpet_width = gr.Number(
                        value=80, label="Carpet Width (cm)", precision=0
                    )
                    carpet_height = gr.Number(
                        value=120, label="Carpet Height (cm)", precision=0
                    )
                
                show_grid = gr.Checkbox(
                    value=True, label="Show Grid Lines"
                )
                
                process_btn = gr.Button("Generate Carpet Pattern", variant="primary")
            
            with gr.Column(scale=2):
                # Output displays
                output_image = gr.Image(
                    label="Carpet Weaving Map",
                    type="filepath"
                )
                
                specifications = gr.Textbox(
                    label="Technical Specifications",
                    lines=10,
                    max_lines=15
                )
        
        with gr.Row():
            # Download files
            csv_download = gr.File(
                label="Download Pattern Grid (CSV)",
                visible=True
            )
            txt_download = gr.File(
                label="Download Color Mapping (TXT)",
                visible=True
            )
        
        # Example configurations
        gr.Markdown("""
        ### 💡 Recommended Settings:
        
        **Fine Persian Style**: 15 knots/cm, 60×90 cm, 20 colors  
        **Standard Quality**: 10 knots/cm, 80×120 cm, 16 colors  
        **Rustic/Tribal**: 6 knots/cm, 100×150 cm, 12 colors
        """)
        
        # Processing function
        process_btn.click(
            fn=process_image,
            inputs=[image_input, knots_per_cm, carpet_width, carpet_height, n_colors, show_grid],
            outputs=[output_image, csv_download, txt_download, specifications]
        )
        
        # Footer information
        gr.Markdown("""
        ### Knot Density Guide:
        - **Fine Persian Carpets**: 15-20 knots/cm (very detailed)
        - **Standard Quality**: 8-12 knots/cm (good balance)
        - **Rustic/Tribal**: 4-6 knots/cm (traditional style)
        
        ### File Outputs:
        - **CSV File**: Grid pattern with color codes for each knot position
        - **TXT File**: Color mapping with RGB values and specifications
        - **PNG Image**: Visual weaving map with grid overlay
        """)
    
    return interface

# Create and launch the interface
if __name__ == "__main__":
    interface = create_gradio_interface()
    interface.launch()

# Requirements for requirements.txt:
"""
gradio>=3.40.0
opencv-python>=4.8.0
numpy>=1.24.0
matplotlib>=3.7.0
Pillow>=10.0.0
scikit-learn>=1.3.0
pandas>=2.0.0
"""