#!/usr/bin/env python3
"""
ZO-1 Network Analysis Tool - Main App
Main Gradio interface for Hugging Face Spaces deployment
"""

import gradio as gr
import traceback
import logging
from logging.handlers import RotatingFileHandler
import cv2
import numpy as np
from zo1_core import *

# Configure logging (console + rotating file)
logger = logging.getLogger("zo1")
logger.setLevel(logging.DEBUG)
if not logger.handlers:
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)
    fh = RotatingFileHandler("debug.log", maxBytes=2_000_000, backupCount=2)
    fh.setLevel(logging.DEBUG)
    fmt = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s - %(message)s")
    ch.setFormatter(fmt)
    fh.setFormatter(fmt)
    logger.addHandler(ch)
    logger.addHandler(fh)
# color scheme for POB lab
POB_LAB_PRIMARY = "#24AE8F"
POB_LAB_PRIMARY_HOVER = "#01BC8D"
POB_LAB_TEXT = "#C0C6CC"          # darker grey, not near-white
POB_LAB_TEXT_SUBTLE = "#9AA2A9"   # mid grey
POB_LAB_TEXT_MUTED = "#7A828A"    # darker grey
POB_LAB_BG = "#0B0F14"
POB_LAB_PANEL = "#10161D"
POB_LAB_INPUT = "#141B23"
POB_LAB_BORDER = "#22303D"
POB_LAB_DIVIDER = "#1A252F"
POB_LAB_LINK = "#63E6BE"
POB_LAB_SHADOW = "0 6px 24px rgba(0,0,0,0.35)"

with gr.Blocks(
    title="ZO-1 Network Analysis Tool",
    css=f"""
    html, body {{
      background: var(--poblab-bg) !important;
    }}

    html, body {{ background: var(--poblab-bg) !important; }}

    .gradio-container {{
      --poblab-primary: {POB_LAB_PRIMARY};
      --poblab-primary-hover: {POB_LAB_PRIMARY_HOVER};
      --poblab-text: {POB_LAB_TEXT};
      --poblab-text-subtle: {POB_LAB_TEXT_SUBTLE};
      --poblab-text-muted: {POB_LAB_TEXT_MUTED};
      --poblab-bg: {POB_LAB_BG};
      --poblab-panel: {POB_LAB_PANEL};
      --poblab-input: {POB_LAB_INPUT};
      --poblab-border: {POB_LAB_BORDER};
      --poblab-divider: {POB_LAB_DIVIDER};
      --poblab-link: {POB_LAB_LINK};
      color: var(--poblab-text);
      background: var(--poblab-bg);
      --body-background-fill: var(--poblab-bg);
      --block-background-fill: var(--poblab-panel);
      --panel-background-fill: var(--poblab-panel);
      --input-background-fill: var(--poblab-input);
      --body-text-color: var(--poblab-text);
      --text-color-subtle: var(--poblab-text-subtle);
      --border-color-primary: var(--poblab-border);
    }}

    /* Panels/blocks/forms/tabs to dark surfaces */
    .gr-box, .gr-panel, .gr-group, .gr-block {{ background: var(--poblab-panel) !important; border-color: var(--poblab-border) !important; }}
    .gr-form, .gr-row, .gr-column, .form, .gradio-container .form {{ background: var(--poblab-panel) !important; }}
    .gr-tabs, .gr-tabs .tab-nav, .gr-tabitem, .tabitem, .tabs {{ background: var(--poblab-bg) !important; border-color: var(--poblab-border) !important; }}

    /* Buttons */
    .gr-button.primary {{
      background: var(--poblab-primary) !important;
      color: #C0C6CC !important;   /* dark grey text, not white */
      border-color: transparent !important;
    }}
    .gr-button.primary:hover {{
      background: var(--poblab-primary-hover) !important;
    }}

    /* Markdown emphasis */
    .gr-markdown strong {{ color: #7FA89C !important; }}

    /* Inputs */
    .gr-input, .gr-textbox textarea, .gr-dropdown, .gr-slider, .gr-image, .gr-file, .gradio-file, input, select, textarea {{ background: var(--poblab-input) !important; color: var(--poblab-text) !important; border-color: var(--poblab-border) !important; }}
    .gr-file, .gradio-file {{ background: var(--poblab-input) !important; border: 1px solid var(--poblab-border) !important; }}
    .gr-file .file-wrap, .gradio-file .file-wrap {{ background: var(--poblab-input) !important; border: 0 !important; padding: 8px 10px !important; }}
    .gr-file input[type="file"], .gradio-file input[type="file"] {{ color: var(--poblab-text) !important; }}
    .gr-file .file-preview, .gradio-file .file-preview {{ background-color: var(--poblab-panel) !important; border: 1px solid var(--poblab-border) !important; }}

    /* Labels */
    label, .block-label, .gradio-container label, [data-testid="block-label"] {{
      color: var(--poblab-text) !important;
      font-weight: 600;
    }}

    /* Subhead style used via gr.HTML */
    .poblab-subhead {{
      background: transparent !important;
      color: var(--poblab-text) !important;
      font-weight: 700;
      font-size: 0.95rem;
      text-transform: uppercase;
      letter-spacing: 0.03em;
      margin-bottom: 6px;
    }}
    /* Restyle block labels as flat subheadings */
    [data-testid="block-label"], 
    .block-label, 
    .gradio-container label {{
      background: transparent !important;   /* remove the lighter strip */
      color: #C0C6CC !important;            /* mid-grey text */
      font-weight: 700;                     /* heavier */
      font-size: 0.95rem;                   /* slightly larger than body text */
      margin-bottom: 4px;                   /* a little spacing before component */
      text-transform: uppercase;            /* optional: KuCoin-like style */
      letter-spacing: 0.03em;               /* subtle spacing */
    }}

    /* Style section titles inside Markdown blocks */
    .gr-group [data-testid="markdown"] h3,
    .gr-group [data-testid="markdown"] p {{
      font-weight: 700 !important;           /* make it bold */
      font-size: 0.95rem !important;         /* slightly larger */
      text-transform: uppercase !important;  /* ALL CAPS */
      letter-spacing: 0.03em !important;     /* spaced out letters */
      color: var(--poblab-text) !important;  /* use our dark grey text */
      background: var(--poblab-panel) !important; /* dark panel bg */
      border: 1px solid var(--poblab-border) !important;
      display: inline-block !important;
      padding: 4px 8px !important;
      border-radius: 6px !important;
      margin: 0 0 6px 0 !important;
    }}
    """

) as demo:
    gr.Markdown("""
    # 🔬 ZO-1 Network Analysis & Quantification
    
    Advanced segmentation and RIS analysis using cutting-edge AI magic ✨
    
    This tool analyzes ZO-1 junction networks using either:
    - **🔵 RIS (Radial Integrity Score)**: Concentric circles approach (recommended)
    - **📊 TiJOR**: Expanding rectangles method (legacy)
    """)
    
    with gr.Tabs():
        with gr.TabItem("📸 Image Upload & Segmentation"):
            with gr.Row():
                with gr.Column():
                    def preprocess_image(image):
                        """Preprocess uploaded image for display (robust to 16-bit/float TIFF)."""
                        if image is None:
                            return None
                        
                        # Ensure dtype is supported by OpenCV before any color conversion
                        if getattr(image, 'dtype', None) is not None and image.dtype not in (np.uint8, np.uint16, np.float32):
                            img_min = float(image.min())
                            img_max = float(image.max())
                            if img_max > img_min:
                                image = ((image - img_min) / (img_max - img_min) * 255).astype(np.uint8)
                            else:
                                image = np.zeros_like(image, dtype=np.uint8)
                        
                        # Handle different image formats and data types
                        if len(image.shape) == 3:
                            # Convert RGBA to RGB
                            if image.shape[2] == 4:
                                image = image[:, :, :3]
                            display_image = image
                        else:
                            # Convert grayscale to RGB for display (after dtype safety above)
                            display_image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
                        
                        # Ensure 8-bit for display (TIFF files might be 16-bit or float)
                        if display_image.dtype != np.uint8:
                            dmin = float(display_image.min())
                            dmax = float(display_image.max())
                            if dmax > dmin:
                                display_image = ((display_image - dmin) / (dmax - dmin) * 255).astype(np.uint8)
                            else:
                                display_image = np.zeros_like(display_image, dtype=np.uint8)
                        
                        # Scale to 25% for preview
                        height, width = display_image.shape[:2]
                        new_height, new_width = int(height * 0.25), int(width * 0.25)
                        display_image = cv2.resize(display_image, (new_width, new_height), interpolation=cv2.INTER_AREA)
                        
                        return display_image
                    
                    image_input = gr.File(label="Upload ZO-1 Image", file_count="single", type="filepath", file_types=["image"]) 
                    image_preview = gr.Image(label="Image Preview", interactive=False, height=180)
                    cell_diameter = gr.Slider(10, 250, value=30, step=5, label="Cell Diameter Estimate (pixels)")
                    scale_factor = gr.Slider(0.1, 1.0, value=1.0, step=0.1, label="Scale Factor (1.0 = full size; lower if slow)")
                    enable_validation = gr.Checkbox(label="Enable AI Contour Validation", value=False)
                    validation_method = gr.Dropdown(
                        ["K-means clustering", "Gaussian Mixture Model (GMM)", "Otsu thresholding"],
                        value="K-means clustering",
                        label="Validation Method"
                    )
                    process_btn = gr.Button("🚀 Run Segmentation", variant="primary")
                    gr.Markdown("Tip: If segmentation isn't satisfactory, adjust the Cell Diameter and rerun. When satisfied, switch to the Analysis tab.")
                
                with gr.Column():
                    segmentation_output = gr.Textbox(label="Segmentation Status", lines=3)
                    segmentation_viz = gr.Image(label="Segmentation Results")
            
            def safe_process_image(*args):
                try:
                    # args[0] is now a filepath from gr.File
                    image_path = args[0]
                    # Preview: load minimal and downscale for display
                    try:
                        if isinstance(image_path, str):
                            img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
                            if img is not None:
                                disp = img if len(img.shape)==3 else cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
                                h, w = disp.shape[:2]
                                scale = 180.0 / max(h, w)
                                disp = cv2.resize(disp, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_AREA)
                            else:
                                disp = None
                        else:
                            disp = None
                    except Exception:
                        disp = None
                    # Process image (zo1_core handles filepath or numpy)
                    result = process_image(*args)
                    return result
                except Exception as e:
                    error_msg = f"Error: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
                    return error_msg, None, None
            
            # Connect image upload to preview
            # Preview generation wired to file input
            # We reuse safe_process_image preview logic by calling preprocess separately
            # For simplicity, just let segmentation click show first preview (keeps UI compact)
            
            process_btn.click(
                safe_process_image,
                inputs=[image_input, cell_diameter, scale_factor, enable_validation, validation_method],
                outputs=[segmentation_output, segmentation_viz, gr.State()]
            )
        
        with gr.TabItem("🔬 Analysis"):
            with gr.Row():
                with gr.Column():
                    analysis_geometry = gr.Dropdown(
                        ["Circles (RIS - recommended)", "Rectangles (TiJOR)"],
                        value="Circles (RIS - recommended)",
                        label="Analysis Method"
                    )
                    
                    with gr.Group():
                        gr.HTML('<div class="poblab-subhead">RIS Parameters</div>')
                        packing_factor = gr.Slider(1.2, 2.0, value=1.5, step=0.1, label="Packing Factor (κ)")
                        min_radius_percent = gr.Slider(5, 25, value=10, step=5, label="Min Radius (% of image)")
                        max_radius_percent = gr.Slider(30, 100, value=80, step=10, label="Max Radius (% of image)")
                        num_circles = gr.Slider(5, 30, value=15, step=1, label="Number of Circles")
                        min_separation = gr.Slider(1, 20, value=5, step=1, label="Min Separation Between Intersections (px)")
                    
                    with gr.Group():
                        gr.HTML('<div class="poblab-subhead">TiJOR Parameters</div>')
                        initial_size = gr.Slider(1, 100, value=10, step=1, label="Initial Rectangle Size (%)")
                        max_size = gr.Slider(1, 100, value=90, step=1, label="Max Rectangle Size (%)")
                        num_steps = gr.Slider(5, 20, value=10, step=1, label="Number of Steps")
                        min_distance = gr.Slider(1, 20, value=5, step=1, label="Min Cross-section Distance (px)")
                    
                    with gr.Group():
                        gr.HTML('<div class="poblab-subhead">Display Options</div>')
                        show_contours = gr.Checkbox(label="Show cell contours", value=False)
                        show_rectangles = gr.Checkbox(label="Show analysis geometry", value=True)
                        show_cross_sections = gr.Checkbox(label="Show cross-sections", value=True)
                    
                    analyze_btn = gr.Button("🔬 Run Analysis", variant="primary")
                
                with gr.Column():
                    analysis_output = gr.Textbox(label="Analysis Results", lines=8)
                    analysis_viz = gr.Image(label="Analysis Visualization")
            
            def safe_run_analysis(*args):
                try:
                    logger.debug(f"run_analysis args: geometry={args[0]}, initial_size={args[1]}, max_size={args[2]}, num_steps={args[3]}, min_distance={args[4]}, packing_factor={args[5]}, show_contours={args[10]}, show_rectangles={args[11]}, show_cross_sections={args[12]}")
                    result = run_analysis(*args)
                    logger.debug("run_analysis completed")
                    return result
                except Exception as e:
                    error_msg = f"Error: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
                    print(error_msg)
                    logger.error(error_msg)
                    return error_msg, None, None
            
            analyze_btn.click(
                safe_run_analysis,
                inputs=[analysis_geometry, initial_size, max_size, num_steps, min_distance, packing_factor, min_radius_percent, max_radius_percent, num_circles, min_separation, show_contours, show_rectangles, show_cross_sections],
                outputs=[analysis_output, analysis_viz, gr.State()]
            )
        
        with gr.TabItem("💾 Export Results"):
            with gr.Row():
                with gr.Column():
                    export_format = gr.Dropdown(["CSV", "Text Report"], value="CSV", label="Export Format")
                    export_btn = gr.Button("💾 Export Results", variant="primary")
                
                with gr.Column():
                    export_output = gr.Textbox(label="Export Data", lines=10)
                    download_btn = gr.File(label="Download File")
            
            def safe_export(format_choice):
                try:
                    text, path = export_results(format_choice)
                    return text, path
                except Exception as e:
                    err = f"Export failed: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
                    print(err)
                    return err, None

            export_btn.click(
                safe_export,
                inputs=[export_format],
                outputs=[export_output, download_btn]
            )
    
    gr.Markdown("""
    ---
    **🔬 ZO-1 Network Analysis Tool | Powered by Cutting-Edge AI | Enhanced with RIS & TiJOR Quantification ✨**
    
    For support: pierre.bagnaninchi@ed.ac.uk
    """)

if __name__ == "__main__":
    import os
    logger.info("Launching Gradio app...")
    is_space = bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID"))
    if is_space:
        demo.launch(show_error=True, ssr_mode=False)
    else:
        demo.launch(share=True, server_name="127.0.0.1", server_port=7862, debug=True, show_error=True)