import streamlit as st
import cv2
import numpy as np
import networkx as nx
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
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
import json

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

# Page configuration - MUST be first Streamlit command
st.set_page_config(
    page_title="Kolam Design Analyzer",
    page_icon="🎨",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- Session state initialization ---
def initialize_session_state():
    """Initialize all session state variables"""
    defaults = {
        'uploaded_image': None,
        'analysis_complete': False,
        'analysis_results': {},
        'processing': False,
        'image_uploaded': False,
        'analysis_hash': None,
        'cached_figures': {},
        'params_changed': False,
        'file_hash': None
    }
    
    for key, value in defaults.items():
        if key not in st.session_state:
            st.session_state[key] = value

# Initialize session state
initialize_session_state()

# Custom CSS for professional styling and anti-flicker
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;
    }
    
    .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;
    }
    
    /* Anti-flicker CSS */
    .main .block-container {
        padding-top: 1rem;
        padding-bottom: 1rem;
        max-width: 100%;
    }
    
    .stTabs [data-baseweb="tab-list"] {
        gap: 2px;
    }
    
    .stTabs [data-baseweb="tab"] {
        height: 50px;
    }
    
    .element-container {
        width: 100% !important;
    }
    
    /* Prevent layout shifts */
    .stPlotlyChart, .stPyplot {
        width: 100%;
        min-height: 400px;
    }
    
    /* Stabilize metrics */
    [data-testid="metric-container"] {
        min-height: 80px;
    }
</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)

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:
            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:
            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:
            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:
            return {}

# Initialize analyzer - cached to prevent recreation
@st.cache_resource
def get_analyzer():
    return KolamAnalyzer()

analyzer = get_analyzer()

# Check library availability
PANDAS_AVAILABLE = True
PLOTLY_AVAILABLE = True
CRYPTO_AVAILABLE = True

try:
    import pandas as pd
except ImportError:
    PANDAS_AVAILABLE = False

try:
    import plotly.graph_objects as go
    import plotly.express as px
except ImportError:
    PLOTLY_AVAILABLE = False

try:
    from cryptography.fernet import Fernet
except ImportError:
    CRYPTO_AVAILABLE = False

# Helper function to create stable matplotlib figures
@st.cache_data(hash_funcs={np.ndarray: lambda x: x.tobytes()})
def create_processing_figure(original_img, thresh_img, edges_img):
    """Create cached matplotlib figure for image processing"""
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    
    axes[0].imshow(original_img, cmap='gray')
    axes[0].set_title('Original Grayscale', fontsize=12, fontweight='bold')
    axes[0].axis('off')
    
    axes[1].imshow(thresh_img, cmap='gray')
    axes[1].set_title('Binary Threshold', fontsize=12, fontweight='bold')
    axes[1].axis('off')
    
    axes[2].imshow(edges_img, cmap='gray')
    axes[2].set_title('Edge Detection', fontsize=12, fontweight='bold')
    axes[2].axis('off')
    
    plt.tight_layout()
    return fig

@st.cache_data(hash_funcs={np.ndarray: lambda x: x.tobytes()})
def create_nodes_figure(original_img, nodes):
    """Create cached matplotlib figure for detected nodes"""
    img_with_nodes = original_img.copy()
    for x, y in nodes:
        cv2.circle(img_with_nodes, (int(x), int(y)), 3, (255), -1)
    
    fig, ax = plt.subplots(1, 1, figsize=(8, 8))
    ax.imshow(img_with_nodes, cmap='gray')
    ax.set_title(f'Detected Nodes: {len(nodes)}', fontsize=14, fontweight='bold')
    ax.axis('off')
    return fig

@st.cache_data(hash_funcs={nx.Graph: lambda g: str(sorted(g.edges()))})
def create_interactive_graph(G):
    """Create cached interactive graph visualization"""
    if G.number_of_nodes() == 0:
        return None
        
    pos = nx.get_node_attributes(G, 'pos')
    
    if not pos:
        pos = nx.spring_layout(G, seed=42)  # Fixed seed for consistency
    
    # 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',
        name='Edges'
    )
    
    # 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,
        name='Nodes',
        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',
            height=500  # Fixed height to prevent layout shifts
        )
    )
    
    return fig

# Sidebar with parameters
with st.sidebar:
    st.markdown("### 🔧 Analysis Parameters")
    
    # Initialize default parameters
    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
        }
    
    # Get current parameters
    current_params = st.session_state['params'].copy()
    
    # Parameter sliders
    image_size = st.slider("Image Processing Size", 128, 512, current_params['image_size'], step=64)
    threshold_value = st.slider("Binary Threshold", 50, 200, current_params['threshold_value'])
    canny_low = st.slider("Canny Low Threshold", 10, 100, current_params['canny_low'])
    canny_high = st.slider("Canny High Threshold", 50, 200, current_params['canny_high'])
    max_corners = st.slider("Maximum Corners", 50, 200, current_params['max_corners'])
    min_line_length = st.slider("Minimum Line Length", 3, 20, current_params['min_line_length'])
    
    # Update parameters and check if changed
    new_params = {
        '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
    }
    
    if new_params != st.session_state['params']:
        st.session_state['params'] = new_params
        st.session_state['params_changed'] = True
    
    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"):
        for key in ['analysis_complete', 'analysis_results', 'uploaded_image', 
                   'processing', 'analysis_hash', 'cached_figures', 'file_hash']:
            if key in st.session_state:
                del st.session_state[key]
        st.cache_data.clear()
        st.rerun()

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

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"
    )

    # Handle file upload with hash checking
    if uploaded_file is not None:
        file_hash = hash(uploaded_file.read())
        uploaded_file.seek(0)  # Reset file pointer
        
        if st.session_state['file_hash'] != file_hash:
            st.session_state['uploaded_image'] = Image.open(uploaded_file)
            st.session_state['file_hash'] = file_hash
            st.session_state['analysis_complete'] = False
            st.session_state['analysis_results'] = {}
            st.cache_data.clear()  # Clear cache for new image

    # 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
        analyze_disabled = (st.session_state.get('processing', False) or 
                          (st.session_state.get('analysis_complete', False) and 
                           not st.session_state.get('params_changed', False)))
        
        if st.button("🔍 Analyze Kolam Design", disabled=analyze_disabled):
            st.session_state['processing'] = True
            st.session_state['params_changed'] = False
            
            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.get('analysis_complete', False) and st.session_state.get('analysis_results'):
        results = st.session_state['analysis_results']
        
        # Create stable tabs
        tab1, tab2, tab3, tab4 = st.tabs(["🖼️ Image Processing", "📈 Graph Analysis", "📊 Features", "🔐 Security"])
        
        with tab1:
            st.markdown("#### Image Processing Pipeline")
            
            # Use cached figure creation
            fig = create_processing_figure(
                results['original_img'], 
                results['thresh_img'], 
                results['edges_img']
            )
            st.pyplot(fig, clear_figure=True)
            
            # Show detected nodes with cached figure
            st.markdown("#### Detected Corner Points")
            fig_nodes = create_nodes_figure(results['original_img'], results['nodes'])
            st.pyplot(fig_nodes, clear_figure=True)
        
        with tab2:
            st.markdown("#### Interactive Graph Visualization")
            
            # Create interactive graph with caching
            fig_interactive = create_interactive_graph(results['graph'])
            if fig_interactive:
                st.plotly_chart(fig_interactive, use_container_width=True, key="main_graph")
            else:
                st.warning("No graph structure detected in the image.")
            
            # Stable 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))
        
        with tab3:
            st.markdown("#### Mathematical Properties")
            
            # Create stable dataframe
            if PANDAS_AVAILABLE:
                features_data = [
                    {"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)}
                ]
                
                features_df = pd.DataFrame(features_data)
                st.dataframe(features_df, use_container_width=True, hide_index=True)
            else:
                # Display as table without pandas
                for key, value in results['features'].items():
                    st.write(f"**{key.replace('_', ' ').title()}**: {value}")
            
            # Degree distribution with fixed height
            if results['graph'].number_of_nodes() > 0 and PLOTLY_AVAILABLE:
                degrees = [d for _, d in results['graph'].degree()]
                if degrees:
                    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',
                        height=400  # Fixed height
                    )
                    st.plotly_chart(fig_hist, use_container_width=True, key="degree_hist")
        
        with tab4:
            st.markdown("#### Security & Data Protection")
            
            if CRYPTO_AVAILABLE:
                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'][:50]}...", language="text")
                else:
                    st.error("Failed to encrypt graph data")
            else:
                st.warning("🔒 Encryption not available.")
                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:
                # Features 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:
                # 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"
                )
    else:
        st.info("👆 Please upload a Kolam image and click 'Analyze' to see results")
        
        # Show placeholder content to maintain layout stability
        tab1, tab2, tab3, tab4 = st.tabs(["🖼️ Image Processing", "📈 Graph Analysis", "📊 Features", "🔐 Security"])
        
        with tab1:
            st.write("Upload an image and run analysis to see image processing results.")
        
        with tab2:
            st.write("Upload an image and run analysis to see graph visualization.")
        
        with tab3:
            st.write("Upload an image and run analysis to see mathematical features.")
        
        with tab4:
            st.write("Upload an image and run analysis to see security options.")

# 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)