import gradio as gr
import numpy as np
import plotly.graph_objects as go
import sympy
import cv2
import time
import os
import sys
from collections import Counter
from PIL import Image

# Ensure logos package is importable
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
if parent_dir not in sys.path:
    sys.path.append(parent_dir)

try:
    from logos.dsp_bridge import DSPBridge
    from logos.logos_core import PRIME_MODULO
except ImportError as e:
    print(f"Error importing LOGOS components: {e}")
    DSPBridge = None
    PRIME_MODULO = 9973

# ==========================================
# PRIME VISUALIZATION (From User Snippet)
# ==========================================

def get_gpf(n):
    """Returns the Greatest Prime Factor."""
    if n <= 1: return 1
    i = 2
    while i * i <= n:
        if n % i:
            i += 1
        else:
            n //= i
    return n

def visualize_potentiality_flow():
    """
    Tab 1: Directed Graph (Sankey) showing Digit Constraints.
    """
    labels = ["Integer Stream"] + [f"Ends in {i}" for i in range(10)] + ["Composite Sink", "Prime Potential (P_n)"]
    sources, targets, values, colors = [], [], [], []

    # Layer 1: Stream -> Digits
    for i in range(10):
        sources.append(0); targets.append(i+1); values.append(10); colors.append("#444")

    # Layer 2: Digits -> Destination
    prime_lanes = [1, 3, 7, 9]
    for i in range(10):
        sources.append(i+1)
        if i in prime_lanes:
            targets.append(12) # Prime Potential
            colors.append("#00ffea") # Cyan
        else:
            targets.append(11) # Sink
            colors.append("#ff0055") # Red
        values.append(10)

    fig = go.Figure(data=[go.Sankey(
        node=dict(pad=15, thickness=20, line=dict(color="black", width=0.5), label=labels, color=["white"]+["#333"]*10+["#ff0055", "#00ffea"]),
        link=dict(source=sources, target=targets, value=values, color=colors)
    )])
    fig.update_layout(title="Prime Potentiality Flow (Mod 10 Constraints)", template="plotly_dark", height=600)
    return fig

def visualize_prime_network(max_integer, show_links):
    """
    Tab 2: Radial Topology.
    """
    fig = go.Figure()
    positions, gpf_map, prime_counts = {}, {}, Counter()

    for n in range(1, max_integer + 1):
        angle = np.pi/2 - (2 * np.pi * (n % 10)) / 10 # Clockwise from Top
        radius = n 
        positions[n] = (radius * np.cos(angle), radius * np.sin(angle))

        if n > 1 and not sympy.isprime(n):
            gpf = get_gpf(n)
            gpf_map[n] = gpf
            prime_counts[gpf] += 1

    if show_links:
        edge_x, edge_y = [], []
        for n, base in gpf_map.items():
            if base in positions:
                x0, y0 = positions[n]
                x1, y1 = positions[base]
                edge_x.extend([x0, x1, None])
                edge_y.extend([y0, y1, None])
        fig.add_trace(go.Scatter(x=edge_x, y=edge_y, mode='lines', line=dict(color='rgba(100,100,100,0.15)', width=0.5), hoverinfo='none', name='GPF Gravity'))

    # Draw Nodes
    px, py, ps, pt = [], [], [], []
    cx, cy, ct = [], [], []

    for n in range(1, max_integer + 1):
        x, y = positions[n]
        if sympy.isprime(n) or n == 1:
            px.append(x); py.append(y)
            ps.append(5 + (np.log(prime_counts[n]+1)*6))
            pt.append(f"PRIME: {n}<br>Gravity: {prime_counts[n]}")
        else:
            cx.append(x); cy.append(y)
            ct.append(f"Composite: {n}")

    fig.add_trace(go.Scatter(x=cx, y=cy, mode='markers', marker=dict(size=3, color='#ff0055', opacity=0.5), text=ct, hoverinfo='text', name='Composites'))
    fig.add_trace(go.Scatter(x=px, y=py, mode='markers', marker=dict(size=ps, color='#00ffea', line=dict(width=1, color='white')), text=pt, hoverinfo='text', name='Primes'))

    # Spokes
    for i in range(10):
        angle = np.pi/2 - (2 * np.pi * i) / 10
        fig.add_trace(go.Scatter(x=[0, max_integer*1.1*np.cos(angle)], y=[0, max_integer*1.1*np.sin(angle)], mode='lines', line=dict(color='#222', width=1, dash='dot'), showlegend=False))

    fig.update_layout(title=f"Radial Prime-Indexed Topology", template="plotly_dark", height=800, width=800, xaxis=dict(visible=False), yaxis=dict(visible=False))
    return fig

def visualize_gpf_counts(sequence_length):
    """
    Tab 3: GPF Density (The Orange Graph).
    """
    gpf_counts = Counter()
    for n in range(4, sequence_length):
        if not sympy.isprime(n): gpf_counts[get_gpf(n)] += 1

    sorted_gpfs = sorted(gpf_counts.keys())
    counts = [gpf_counts[p] for p in sorted_gpfs]

    fig = go.Figure(data=go.Bar(x=sorted_gpfs, y=counts, marker_color='#ff7f00', name="Composite Count"))
    fig.update_layout(title="Composite Density by GPF Base", xaxis_title="Prime Base", yaxis_title="Count", template="plotly_dark", xaxis=dict(type='category'))
    return fig

# ==========================================
# LOGOS MACHINE SHOP
# ==========================================

def process_dsp(image, grid_size=8, workers=16):
    """
    Process image through LOGOS SPCW DSP Bridge.
    """
    if image is None:
        return None, "No Image Provided"
    
    if DSPBridge is None:
        return None, "LOGOS Components Not Loaded"
        
    # Save temp file
    temp_path = "temp_ingest.png"
    if isinstance(image, np.ndarray):
        img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        cv2.imwrite(temp_path, img_bgr)
    elif isinstance(image, str):
        temp_path = image
    else:
        # PIL Image
        try:
            image.save(temp_path)
        except Exception:
            # Maybe numpy conversion needed
            np_img = np.array(image)
            cv2.imwrite(temp_path, cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR))

    # Init Bridge
    bridge = DSPBridge(num_workers=int(workers), viewport_size=(1024, 768))
    
    try:
        # Transmit (Encoding -> Decoding)
        stats = bridge.transmit(temp_path, show_window=False)
        
        # Get result
        recon = bridge.get_canvas()
        if recon is not None:
            # Convert to PIL
            output_img = Image.fromarray(recon)
        else:
            output_img = None
            
        # Format stats
        status = f"""
        ### LOGOS SPCW TRANSMISSION REPORT
        - **State**: CRYSTALLINE
        - **Mod**: {PRIME_MODULO}
        - **Waves**: {stats.total_tiles} ({stats.parallel_waves} parallel)
        - **Atoms**: {stats.atoms_sent}
        - **Throughput**: {stats.throughput_mbps:.2f} MB/s
        - **Latency**: {stats.elapsed_ms:.1f} ms
        - **SSIM**: {stats.ssim:.6f}
        """
        
        return output_img, status
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        return None, f"ERROR: {str(e)}"

# ==========================================
# APP LAYOUT
# ==========================================

with gr.Blocks(theme=gr.themes.Monochrome(), title="LOGOS SPCW Protocol") as demo:
    gr.Markdown("# LOGOS: Structured Prime Composite Waveform (SPCW)")
    gr.Markdown("_\"The Machine Shop\" - Research & Development Lab_")
    
    with gr.Tabs():
        with gr.Tab("Prime Network Visualizer"):
            gr.Markdown("## Layer 1: Mathematical Blueprints")
            with gr.Row():
                with gr.Column(scale=1):
                    max_int = gr.Slider(50, 5000, value=1000, label="Max Integer", step=50)
                    show_links = gr.Checkbox(True, label="Show Gravity Links (GPF)")
                    btn_viz = gr.Button("Generate Network", variant="primary")
                with gr.Column(scale=2):
                    plot_radial = gr.Plot(label="Radial Topology")
            
            with gr.Row():
                with gr.Column():
                    plot_flow = gr.Plot(label="Potentiality Flow")
                with gr.Column():
                    plot_counts = gr.Plot(label="Composite Density")
            
            btn_viz.click(visualize_prime_network, [max_int, show_links], plot_radial)
            btn_viz.click(visualize_potentiality_flow, None, plot_flow)
            btn_viz.click(visualize_gpf_counts, [max_int], plot_counts)

        with gr.Tab("The Machine Shop (DSP Bridge)"):
            gr.Markdown("## Layer 3: SPCW Transmission Engine")
            with gr.Row():
                with gr.Column():
                    input_img = gr.Image(type="numpy", label="Input Reality (The Cake)")
                    workers = gr.Slider(1, 64, value=16, label="Parallel Workers", step=1)
                    btn_proc = gr.Button("Bake & Eat (Transmit)", variant="primary")
                with gr.Column():
                    output_img = gr.Image(type="pil", label="Reconstructed Hologram")
                    output_stats = gr.Markdown("Waiting for transmission...")
            
            btn_proc.click(process_dsp, [input_img, workers], [output_img, output_stats])

if __name__ == "__main__":
    demo.launch()