import streamlit as st
import cv2
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import pandas as pd
import io
import base64
from PIL import Image
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.express as px

# --- Session state initialization ---
def initialize_session_state():
    """Initialize all session state variables"""
    if 'uploaded_image' not in st.session_state:
        st.session_state['uploaded_image'] = None
    if 'analysis_complete' not in st.session_state:
        st.session_state['analysis_complete'] = False
    if 'analysis_results' not in st.session_state:
        st.session_state['analysis_results'] = {}
    if 'processing' not in st.session_state:
        st.session_state['processing'] = False

# Fix for Hugging Face Spaces permissions
import os
import tempfile
os.environ['STREAMLIT_BROWSER_GATHER_USAGE_STATS'] = 'false'
os.environ['MPLCONFIGDIR'] = tempfile.gettempdir()

# Initialize session state
initialize_session_state()

# Page configuration
st.set_page_config(
    page_title="Kolam Design Analyzer",
    page_icon="🎨",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for professional styling
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(90deg, #FF6B35 0%, #F7931E 100%);
        padding: 2rem;
        border-radius: 10px;
        margin-bottom: 2rem;
        color: white;
        text-align: center;
    }
    
    .metric-card {
        background: white;
        padding: 1rem;
        border-radius: 8px;
        border-left: 4px solid #FF6B35;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        margin: 0.5rem 0;
    }
    
    .analysis-section {
        background: #f8f9fa;
        padding: 1.5rem;
        border-radius: 10px;
        margin: 1rem 0;
    }
    
    .upload-section {
        border: 2px dashed #FF6B35;
        padding: 2rem;
        border-radius: 10px;
        text-align: center;
        margin: 1rem 0;
        background: #fff9f7;
    }
    
    .stButton > button {
        background: linear-gradient(90deg, #FF6B35 0%, #F7931E 100%);
        color: white;
        border: none;
        border-radius: 5px;
        padding: 0.5rem 2rem;
        font-weight: bold;
        transition: all 0.3s;
    }
    
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 4px 8px rgba(0,0,0,0.2);
    }
    
    /* Prevent page jumping */
    .main .block-container {
        padding-top: 1rem;
        padding-bottom: 1rem;
    }
</style>
""", unsafe_allow_html=True)

# Title and header
st.markdown("""
<div class="main-header">
    <h1>🎨 Kolam Design Analyzer</h1>
    <h3>Smart India Hackathon 2025 - AI-Powered Traditional Art Analysis</h3>
    <p>Discover the mathematical principles and geometric patterns behind traditional Kolam designs</p>
</div>
""", unsafe_allow_html=True)

# Sidebar with consistent parameters
with st.sidebar:
    st.markdown("### 🔧 Analysis Parameters")
    
    # Use session state for parameters to prevent re-runs
    if 'params' not in st.session_state:
        st.session_state['params'] = {
            'image_size': 256,
            'threshold_value': 127,
            'canny_low': 30,
            'canny_high': 100,
            'max_corners': 100,
            'min_line_length': 5
        }
    
    image_size = st.slider("Image Processing Size", 128, 512, 
                          st.session_state['params']['image_size'], step=64)
    threshold_value = st.slider("Binary Threshold", 50, 200, 
                               st.session_state['params']['threshold_value'])
    canny_low = st.slider("Canny Low Threshold", 10, 100, 
                         st.session_state['params']['canny_low'])
    canny_high = st.slider("Canny High Threshold", 50, 200, 
                          st.session_state['params']['canny_high'])
    max_corners = st.slider("Maximum Corners", 50, 200, 
                           st.session_state['params']['max_corners'])
    min_line_length = st.slider("Minimum Line Length", 3, 20, 
                               st.session_state['params']['min_line_length'])
    
    # Update parameters in session state
    st.session_state['params'].update({
        'image_size': image_size,
        'threshold_value': threshold_value,
        'canny_low': canny_low,
        'canny_high': canny_high,
        'max_corners': max_corners,
        'min_line_length': min_line_length
    })
    
    st.markdown("---")
    st.markdown("### 📊 About This Tool")
    st.info("This application uses computer vision and graph theory to analyze traditional Kolam designs, extracting geometric patterns and design principles.")
    
    # Reset button
    if st.button("🔄 Reset Analysis"):
        st.session_state['analysis_complete'] = False
        st.session_state['uploaded_image'] = None
        st.session_state['analysis_results'] = {}
        st.session_state['processing'] = False
        st.rerun()

class KolamAnalyzer:
    def __init__(self):
        self.cipher = None
        self.encryption_key = None
        
    def generate_encryption_key(self):
        """Generate encryption key for graph data"""
        try:
            from cryptography.fernet import Fernet
            self.encryption_key = Fernet.generate_key()
            self.cipher = Fernet(self.encryption_key)
            return self.encryption_key.decode()
        except ImportError:
            return "Encryption not available"
        
    def preprocess_image(self, image, size, threshold_val, canny_low, canny_high):
        """Preprocess uploaded image"""
        try:
            # Convert PIL image to OpenCV format
            img_array = np.array(image)
            if len(img_array.shape) == 3:
                img_gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
            else:
                img_gray = img_array
                
            # Resize image
            img_resized = cv2.resize(img_gray, (size, size))
            
            # Apply binary threshold
            _, thresh = cv2.threshold(img_resized, threshold_val, 255, cv2.THRESH_BINARY_INV)
            
            # Edge detection
            edges = cv2.Canny(thresh, canny_low, canny_high)
            
            return img_resized, thresh, edges
        except Exception as e:
            st.error(f"Error in image preprocessing: {str(e)}")
            return None, None, None
    
    def detect_nodes(self, edges, max_corners):
        """Detect corner points as graph nodes"""
        try:
            corners = cv2.goodFeaturesToTrack(
                edges, 
                maxCorners=max_corners, 
                qualityLevel=0.01, 
                minDistance=5
            )
            if corners is None:
                return []
            return [tuple(pt.ravel()) for pt in corners.astype(int)]
        except Exception as e:
            st.error(f"Error in node detection: {str(e)}")
            return []
    
    def detect_edges(self, edges, nodes, min_line_length):
        """Detect lines and create graph edges"""
        try:
            lines = cv2.HoughLinesP(
                edges, 
                1, 
                np.pi/180, 
                threshold=30, 
                minLineLength=min_line_length, 
                maxLineGap=10
            )
            
            graph_edges = []
            if lines is not None and len(nodes) > 0:
                for x1, y1, x2, y2 in lines[:,0]:
                    n1 = min(range(len(nodes)), 
                           key=lambda i: np.linalg.norm(np.array(nodes[i]) - np.array([x1,y1])))
                    n2 = min(range(len(nodes)), 
                           key=lambda i: np.linalg.norm(np.array(nodes[i]) - np.array([x2,y2])))
                    if n1 != n2 and (n1, n2) not in graph_edges and (n2, n1) not in graph_edges:
                        graph_edges.append((n1, n2))
            
            # Fallback: connect nearby nodes if no lines detected
            if len(graph_edges) == 0:
                graph_edges = self.connect_nearby_nodes(nodes, max_distance=30)
                
            return graph_edges
        except Exception as e:
            st.error(f"Error in edge detection: {str(e)}")
            return []
    
    def connect_nearby_nodes(self, nodes, max_distance=30):
        """Connect nearby nodes as fallback"""
        edges = []
        for i, (x1, y1) in enumerate(nodes):
            for j, (x2, y2) in enumerate(nodes):
                if i < j:
                    distance = np.linalg.norm(np.array([x1, y1]) - np.array([x2, y2]))
                    if distance <= max_distance:
                        edges.append((i, j))
        return edges
    
    def build_graph(self, nodes, edges):
        """Build NetworkX graph from nodes and edges"""
        try:
            G = nx.Graph()
            for idx, node in enumerate(nodes):
                G.add_node(idx, pos=node)
            for n1, n2 in edges:
                G.add_edge(n1, n2)
            return G
        except Exception as e:
            st.error(f"Error in graph building: {str(e)}")
            return nx.Graph()
    
    def extract_graph_features(self, G):
        """Extract mathematical features from the graph"""
        try:
            num_nodes = G.number_of_nodes()
            num_edges = G.number_of_edges()
            degrees = [d for _, d in G.degree()]
            avg_degree = np.mean(degrees) if degrees else 0
            max_degree = max(degrees) if degrees else 0
            min_degree = min(degrees) if degrees else 0
            
            # Calculate cycles
            try:
                num_cycles = sum(1 for c in nx.cycle_basis(G))
            except:
                num_cycles = 0
                
            # Calculate connectivity
            is_connected = nx.is_connected(G) if num_nodes > 0 else False
            num_components = nx.number_connected_components(G)
            
            # Calculate centrality measures
            try:
                betweenness = nx.betweenness_centrality(G)
                avg_betweenness = np.mean(list(betweenness.values())) if betweenness else 0
                
                closeness = nx.closeness_centrality(G)
                avg_closeness = np.mean(list(closeness.values())) if closeness else 0
            except:
                avg_betweenness = 0
                avg_closeness = 0
            
            return {
                "num_nodes": num_nodes,
                "num_edges": num_edges,
                "avg_degree": round(avg_degree, 2),
                "max_degree": max_degree,
                "min_degree": min_degree,
                "num_cycles": num_cycles,
                "is_connected": is_connected,
                "num_components": num_components,
                "avg_betweenness": round(avg_betweenness, 4),
                "avg_closeness": round(avg_closeness, 4),
                "density": round(nx.density(G), 4) if num_nodes > 1 else 0
            }
        except Exception as e:
            st.error(f"Error in feature extraction: {str(e)}")
            return {}
    
    def encrypt_graph(self, G):
        """Encrypt graph data for security"""
        try:
            if not self.cipher:
                self.generate_encryption_key()
            
            adj_matrix = nx.to_numpy_array(G)
            adj_bytes = adj_matrix.tobytes()
            encrypted = self.cipher.encrypt(adj_bytes)
            return encrypted
        except Exception as e:
            return None
    
    def create_interactive_graph(self, G):
        """Create interactive graph visualization using Plotly"""
        try:
            pos = nx.get_node_attributes(G, 'pos')
            
            if not pos:
                # If no positions, use spring layout
                pos = nx.spring_layout(G)
            
            # Extract edges
            edge_x = []
            edge_y = []
            for edge in G.edges():
                x0, y0 = pos[edge[0]]
                x1, y1 = pos[edge[1]]
                edge_x.extend([x0, x1, None])
                edge_y.extend([y0, y1, None])
            
            # Create edge trace
            edge_trace = go.Scatter(
                x=edge_x, y=edge_y,
                line=dict(width=2, color='#FF6B35'),
                hoverinfo='none',
                mode='lines'
            )
            
            # Extract nodes
            node_x = []
            node_y = []
            node_text = []
            node_degree = []
            
            for node in G.nodes():
                x, y = pos[node]
                node_x.append(x)
                node_y.append(y)
                degree = G.degree(node)
                node_degree.append(degree)
                node_text.append(f'Node {node}<br>Degree: {degree}')
            
            # Create node trace
            node_trace = go.Scatter(
                x=node_x, y=node_y,
                mode='markers',
                hoverinfo='text',
                text=node_text,
                marker=dict(
                    size=[max(10, d*3) for d in node_degree],
                    color=node_degree,
                    colorscale='Viridis',
                    colorbar=dict(
                        thickness=15,
                        len=0.5,
                        x=1.02,
                        title="Node Degree"
                    ),
                    line=dict(width=2, color='white')
                )
            )
            
            # Create figure
            fig = go.Figure(data=[edge_trace, node_trace],
                           layout=go.Layout(
                                title='Interactive Kolam Graph Structure',
                                titlefont_size=16,
                                showlegend=False,
                                hovermode='closest',
                                margin=dict(b=20,l=5,r=5,t=40),
                                annotations=[ dict(
                                    text="Node size represents degree centrality",
                                    showarrow=False,
                                    xref="paper", yref="paper",
                                    x=0.005, y=-0.002,
                                    xanchor="left", yanchor="bottom",
                                    font=dict(size=12)
                                )],
                                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                                plot_bgcolor='white'
                           ))
            
            return fig
        except Exception as e:
            st.error(f"Error creating interactive graph: {str(e)}")
            return None

# Initialize analyzer
@st.cache_resource
def get_analyzer():
    return KolamAnalyzer()

analyzer = get_analyzer()

# Main content area
col1, col2 = st.columns([1, 2])

# Check library availability
try:
    import pandas as pd
    PANDAS_AVAILABLE = True
except ImportError as e:
    st.warning("⚠️ Pandas not available due to NumPy compatibility. Using basic data structures.")
    PANDAS_AVAILABLE = False

try:
    import plotly.graph_objects as go
    import plotly.express as px
    PLOTLY_AVAILABLE = True
except ImportError as e:
    st.warning("⚠️ Plotly not available due to NumPy compatibility. Using matplotlib for visualizations.")
    PLOTLY_AVAILABLE = False

try:
    from cryptography.fernet import Fernet
    CRYPTO_AVAILABLE = True
except ImportError as e:
    st.warning("⚠️ Cryptography not available. Encryption features disabled.")
    CRYPTO_AVAILABLE = False

with col1:
    st.markdown("### 📤 Upload Kolam Image")
    
    uploaded_file = st.file_uploader(
        "Choose a Kolam image...",
        type=["png", "jpg", "jpeg"],
        help="Upload a clear image of a Kolam design for analysis",
        key="file_uploader"
    )

    # Handle file upload
    if uploaded_file is not None:
        # Only process if it's a new file
        if (st.session_state['uploaded_image'] is None or 
            uploaded_file.name != getattr(st.session_state.get('uploaded_file'), 'name', None)):
            st.session_state['uploaded_image'] = Image.open(uploaded_file)
            st.session_state['uploaded_file'] = uploaded_file
            st.session_state['analysis_complete'] = False
            st.session_state['analysis_results'] = {}

    # Display uploaded image
    if st.session_state['uploaded_image'] is not None:
        st.image(st.session_state['uploaded_image'], caption="Uploaded Kolam", use_column_width=True)
        
        # Analysis button
        if st.button("🔍 Analyze Kolam Design", key="analyze_btn", disabled=st.session_state.get('processing', False)):
            st.session_state['processing'] = True
            
            with st.spinner("Analyzing Kolam design..."):
                try:
                    # Process image with current parameters
                    params = st.session_state['params']
                    original, thresh, edges = analyzer.preprocess_image(
                        st.session_state['uploaded_image'], 
                        params['image_size'], 
                        params['threshold_value'], 
                        params['canny_low'], 
                        params['canny_high']
                    )
                    
                    if original is not None:
                        # Detect nodes and edges
                        nodes = analyzer.detect_nodes(edges, params['max_corners'])
                        graph_edges = analyzer.detect_edges(edges, nodes, params['min_line_length'])
                        
                        # Build graph
                        G = analyzer.build_graph(nodes, graph_edges)
                        
                        # Extract features
                        features = analyzer.extract_graph_features(G)
                        
                        # Store results in session state
                        st.session_state['analysis_results'] = {
                            'original_img': original,
                            'thresh_img': thresh,
                            'edges_img': edges,
                            'nodes': nodes,
                            'graph': G,
                            'features': features
                        }
                        
                        # Generate encryption key
                        encryption_key = analyzer.generate_encryption_key()
                        st.session_state['analysis_results']['encryption_key'] = encryption_key
                        
                        st.session_state['analysis_complete'] = True
                        st.success("✅ Analysis completed successfully!")
                    else:
                        st.error("Failed to process the image. Please try with different parameters.")
                        
                except Exception as e:
                    st.error(f"Analysis failed: {str(e)}")
                finally:
                    st.session_state['processing'] = False

with col2:
    st.markdown("### 📊 Analysis Results")
    
    if st.session_state['analysis_complete'] and st.session_state['analysis_results']:
        results = st.session_state['analysis_results']
        
        # Create tabs for different visualizations
        tab1, tab2, tab3, tab4 = st.tabs(["🖼️ Image Processing", "📈 Graph Analysis", "📊 Features", "🔐 Security"])
        
        with tab1:
            st.markdown("#### Image Processing Pipeline")
            
            try:
                # Create subplot for processed images
                fig, axes = plt.subplots(1, 3, figsize=(15, 5))
                
                axes[0].imshow(results['original_img'], cmap='gray')
                axes[0].set_title('Original Grayscale', fontsize=12, fontweight='bold')
                axes[0].axis('off')
                
                axes[1].imshow(results['thresh_img'], cmap='gray')
                axes[1].set_title('Binary Threshold', fontsize=12, fontweight='bold')
                axes[1].axis('off')
                
                axes[2].imshow(results['edges_img'], cmap='gray')
                axes[2].set_title('Edge Detection', fontsize=12, fontweight='bold')
                axes[2].axis('off')
                
                plt.tight_layout()
                st.pyplot(fig)
                plt.close()
                
                # Show detected nodes
                st.markdown("#### Detected Corner Points")
                img_with_nodes = results['original_img'].copy()
                for x, y in results['nodes']:
                    cv2.circle(img_with_nodes, (int(x), int(y)), 3, (255), -1)
                
                fig_nodes, ax_nodes = plt.subplots(1, 1, figsize=(8, 8))
                ax_nodes.imshow(img_with_nodes, cmap='gray')
                ax_nodes.set_title(f'Detected Nodes: {len(results["nodes"])}', 
                                  fontsize=14, fontweight='bold')
                ax_nodes.axis('off')
                st.pyplot(fig_nodes)
                plt.close()
            except Exception as e:
                st.error(f"Error displaying image processing results: {str(e)}")
        
        with tab2:
            st.markdown("#### Interactive Graph Visualization")
            
            try:
                # Create interactive graph
                if results['graph'].number_of_nodes() > 0:
                    fig_interactive = analyzer.create_interactive_graph(results['graph'])
                    if fig_interactive:
                        st.plotly_chart(fig_interactive, use_container_width=True)
                else:
                    st.warning("No graph structure detected in the image.")
                
                # Graph statistics
                col_a, col_b = st.columns(2)
                with col_a:
                    st.metric("Total Nodes", results['features'].get('num_nodes', 0))
                    st.metric("Total Edges", results['features'].get('num_edges', 0))
                    st.metric("Graph Density", results['features'].get('density', 0))
                
                with col_b:
                    st.metric("Average Degree", results['features'].get('avg_degree', 0))
                    st.metric("Number of Cycles", results['features'].get('num_cycles', 0))
                    st.metric("Connected Components", results['features'].get('num_components', 0))
            except Exception as e:
                st.error(f"Error displaying graph analysis: {str(e)}")
        
        with tab3:
            st.markdown("#### Mathematical Properties")
            
            try:
                # Create metrics dataframe
                if PANDAS_AVAILABLE:
                    features_df = pd.DataFrame([
                        {"Property": "Nodes", "Value": results['features'].get('num_nodes', 0)},
                        {"Property": "Edges", "Value": results['features'].get('num_edges', 0)},
                        {"Property": "Average Degree", "Value": results['features'].get('avg_degree', 0)},
                        {"Property": "Maximum Degree", "Value": results['features'].get('max_degree', 0)},
                        {"Property": "Minimum Degree", "Value": results['features'].get('min_degree', 0)},
                        {"Property": "Cycles", "Value": results['features'].get('num_cycles', 0)},
                        {"Property": "Graph Density", "Value": results['features'].get('density', 0)},
                        {"Property": "Average Betweenness", "Value": results['features'].get('avg_betweenness', 0)},
                        {"Property": "Average Closeness", "Value": results['features'].get('avg_closeness', 0)},
                        {"Property": "Connected", "Value": "Yes" if results['features'].get('is_connected', False) else "No"},
                        {"Property": "Components", "Value": results['features'].get('num_components', 0)}
                    ])
                    
                    st.dataframe(features_df, use_container_width=True)
                else:
                    # Display as simple table without pandas
                    for key, value in results['features'].items():
                        st.write(f"**{key.replace('_', ' ').title()}**: {value}")
                
                # Visualize degree distribution
                if results['graph'].number_of_nodes() > 0 and PLOTLY_AVAILABLE:
                    degrees = [d for _, d in results['graph'].degree()]
                    fig_hist = px.histogram(
                        x=degrees, 
                        title="Degree Distribution",
                        labels={'x': 'Node Degree', 'y': 'Frequency'},
                        color_discrete_sequence=['#FF6B35']
                    )
                    fig_hist.update_layout(
                        plot_bgcolor='white',
                        paper_bgcolor='white'
                    )
                    st.plotly_chart(fig_hist, use_container_width=True)
            except Exception as e:
                st.error(f"Error displaying features: {str(e)}")
        
        with tab4:
            st.markdown("#### Security & Data Protection")
            
            try:
                if CRYPTO_AVAILABLE:
                    # Encrypt graph
                    encrypted_data = analyzer.encrypt_graph(results['graph'])
                    
                    if encrypted_data:
                        col_x, col_y = st.columns(2)
                        with col_x:
                            st.success("🔐 Graph data encrypted successfully!")
                            st.info(f"Encrypted data size: {len(encrypted_data)} bytes")
                            
                        with col_y:
                            if results.get('encryption_key'):
                                st.code(f"Encryption Key:\n{results['encryption_key']}", language="text")
                    else:
                        st.error("Failed to encrypt graph data")
                else:
                    st.warning("🔒 Encryption not available due to package compatibility issues.")
                    st.info("Graph data will be stored in plain text format.")
                    
                # Download options
                st.markdown("#### 📥 Download Results")
                
                col_dl1, col_dl2 = st.columns(2)
                with col_dl1:
                    # Prepare features for download
                    if PANDAS_AVAILABLE:
                        features_json = pd.DataFrame([results['features']]).to_json(orient='records')
                    else:
                        import json
                        features_json = json.dumps([results['features']], indent=2)
                        
                    st.download_button(
                        "📊 Download Features (JSON)",
                        data=features_json,
                        file_name="kolam_features.json",
                        mime="application/json"
                    )
                
                with col_dl2:
                    # Prepare adjacency matrix
                    adj_matrix = nx.to_numpy_array(results['graph'])
                    adj_buffer = io.BytesIO()
                    np.save(adj_buffer, adj_matrix)
                    st.download_button(
                        "🔢 Download Adjacency Matrix",
                        data=adj_buffer.getvalue(),
                        file_name="kolam_adjacency.npy",
                        mime="application/octet-stream"
                    )
            except Exception as e:
                st.error(f"Error in security section: {str(e)}")
    else:
        st.info("👆 Please upload a Kolam image and click 'Analyze' to see results")

# Footer
st.markdown("---")
st.markdown("""
<div style="text-align: center; color: #666; padding: 1rem;">
    <p><strong>Kolam Design Analyzer</strong> | Smart India Hackathon 2025</p>
    <p>Preserving traditional art through modern technology 🎨✨</p>
</div>
""", unsafe_allow_html=True)