"""
Hugging Face / local Gradio app for exploring Collatz structures.

Row 1:
  - Inverse tree controls
  - Minimal subtree controls
  - Statistics for the currently displayed graph

Row 2:
  - Image display area (Zoom & Scroll or Fit to Width)
"""

from __future__ import annotations
import io
import matplotlib.pyplot as plt

from typing import Any
from pathlib import Path
import base64

import gradio as gr

from src.utils import (
    build_and_render_collatz_tree,
    build_and_render_minimal_subtree,
    safe_int,
)
from src.collatz.metrics import compute_basic_graph_stats, format_stats_markdown


# ============================================================
# Helpers
# ============================================================

def image_file_to_html(
    path: str,
    mode: str = "Zoom & Scroll",
    box_height: int = 650,
) -> str:
    """
    Convert an image file into an HTML block.

    Modes:
      - "Zoom & Scroll": full resolution inside fixed-height scroll-box
      - "Fit to Width" : scaled to column width, whole graph visible
    """
    img_path = Path(path)
    if not img_path.is_file():
        return "<p style='color:red;'>Error: image file not found.</p>"

    data = img_path.read_bytes()
    encoded = base64.b64encode(data).decode("ascii")

    if mode == "Fit to Width":
        # Show whole graph scaled to container width
        html = f"""
        <div style="
            border:1px solid #ddd;
            border-radius:6px;
            padding:4px;
            background-color:#fafafa;
        ">
          <img src="data:image/png;base64,{encoded}"
               style="display:block; max-width:100%; height:auto; margin:0 auto;" />
        </div>
        """
    else:
        # Zoom & scroll (full resolution)
        html = f"""
        <div style="
            display:flex;
            justify-content:center;
            width:100%;
        ">
            <div style="
                height:{box_height}px;
                overflow:auto;
                border:1px solid #ddd;
                border-radius:6px;
                padding:4px;
                background-color:#fafafa;
                width:fit-content;
                max-width:100%;
            ">
              <img src="data:image/png;base64,{encoded}"
                   style="display:block; max-width:none; max-height:none;" />
            </div>
        </div>
        """

    return html

def parity_histogram_html(stats: dict) -> str:
    """
    Create a small odd vs even histogram as an embedded PNG <img> tag.
    """
    num_odd = stats.get("num_odd", 0)
    num_even = stats.get("num_even", 0)

    # If no nodes, nothing to plot
    if num_odd == 0 and num_even == 0:
        return "<p>_No nodes to plot._</p>"

    labels = ["Odd", "Even"]
    values = [num_odd, num_even]

    fig, ax = plt.subplots(figsize=(3.5, 2.5))
    ax.bar(labels, values)
    ax.set_ylabel("Count")
    ax.set_title("Odd vs Even Nodes")
    fig.tight_layout()

    buf = io.BytesIO()
    fig.savefig(buf, format="png")
    plt.close(fig)
    encoded = base64.b64encode(buf.getvalue()).decode("ascii")

    return f'<img src="data:image/png;base64,{encoded}" style="max-width:100%; height:auto;" />'

# ============================================================
# Callbacks
# ============================================================

def inverse_tree_callback(
    backbone_length: Any,
    branch_length: Any,
    max_depth: Any,
    view_mode: str,
):
    """
    Generate the inverse structural tree and return (image_html, stats_md).
    """

    b_len = safe_int(backbone_length, default=8)
    r_len = safe_int(branch_length, default=4)
    depth = safe_int(max_depth, default=2)

    # clamp for demo
    b_len = max(4, min(b_len, 10))
    r_len = max(1, min(r_len, 7))
    depth = max(0, min(depth, 4))

    image_path, df_edges = build_and_render_collatz_tree(
        backbone_length=b_len,
        branch_length=r_len,
        max_depth=depth,
        return_edges=True,
    )

    html_block = image_file_to_html(image_path, view_mode, 650)

    stats = compute_basic_graph_stats(df_edges)
    stats_md = format_stats_markdown(stats)

    hist_html = parity_histogram_html(stats)

    return html_block, stats_md, hist_html


def minimal_subtree_callback(
    N: Any,
    view_mode: str,
):
    """
    Generate the minimal subtree up to N and return (image_html, stats_md).
    """

    N = safe_int(N, default=7)
    # Cap N for demo to prevent huge graphs
    N = max(1, min(N, 2000))

    image_path, df_edges = build_and_render_minimal_subtree(
        N,
        return_edges=True,
        filename=f"minimal_subtree",
    )

    html_block = image_file_to_html(image_path, view_mode, 650)

    stats = compute_basic_graph_stats(df_edges)
    stats_md = format_stats_markdown(stats)
    hist_html = parity_histogram_html(stats)

    return html_block, stats_md, hist_html


# ============================================================
# Build UI
# ============================================================

def build_demo() -> gr.Blocks:

    with gr.Blocks(title="Collatz Explorer") as demo:

        gr.Markdown(
            """ 
<h1 style="text-align:center; margin-bottom:20px;">
    🔷 <span style="font-weight:700;">Collatz Structural Explorer</span> 🔷
</h1>

<div style="text-align:justify;">

<div style="margin-left:20px; margin-bottom:15px;">
The <em>Collatz Structural Explorer</em> accompanies the research article 
<a href="https://www.tandfonline.com/doi/full/10.1080/27684830.2025.2542052" target="_blank" style="color:#1a73e8; text-decoration:none; font-weight:600;">
Unfolding the Collatz Tree: An Indirect Structural Proof of the Collatz Conjecture
</a>, published in the <em>Journal of Experimental Mathematics</em> (Taylor and Francis). 
This interactive demonstration is intended to visually illustrate key structural ideas from the paper using a dynamic inverse-tree perspective.
</div>

<div style="margin-left:20px;">
It highlights how the inverse Collatz map, structural branch rules, and the minimal subtree containing all natural numbers up to a chosen bound N collectively reconstruct the forward Collatz dynamics in an organized and interpretable way. 
Through real-time visualization and graph statistics, readers can explore the hierarchical structure of the Collatz process and gain an intuitive understanding of the theoretical insights developed in the publication.
</div>

</div>
<div style="height:50px;"></div>
            """
            )

        # ============================
        # Row 1: controls + stats
        # ============================
        with gr.Row():
            # Inverse tree controls
            with gr.Column(scale=1, min_width=260):
                gr.Markdown("### Inverse Collatz Tree")

                backbone_input = gr.Slider(
                    4, 10, value=8, step=1,
                    label="Backbone length (powers of 2)",
                )
                branch_input = gr.Slider(
                    1, 7, value=4, step=1,
                    label="Branch length",
                )
                depth_input = gr.Slider(
                    0, 4, value=2, step=1,
                    label="Branch recursion depth",
                )

                view_mode_inverse = gr.Radio(
                    ["Zoom & Scroll", "Fit to Width"],
                    value="Zoom & Scroll",
                    label="View mode for inverse tree",
                )

                gen_inverse = gr.Button("Generate Inverse Tree")

            # Minimal subtree controls
            with gr.Column(scale=1, min_width=260):
                gr.Markdown("### Minimal Subtree up to N")

                N_input = gr.Number(
                    value=7, precision=0,
                    label="Upper bound N (includes all 1..N)",
                    info="Demo max = 2000",
                )

                view_mode_minimal = gr.Radio(
                    ["Zoom & Scroll", "Fit to Width"],
                    value="Zoom & Scroll",
                    label="View mode for minimal subtree",
                )

                gen_minimal = gr.Button("Generate Minimal Subtree")

            # Stats panel
            # Stats + histogram (side by side)
            with gr.Column(scale=1):
                gr.Markdown("### Current Graph Statistics")

                with gr.Row():
                    # Column for text statistics
                    with gr.Column(scale=2, min_width=140):
                        stats_output = gr.Markdown(
                            value="_No graph generated yet._"
                        )

                    # Column for histogram (right side)
                    with gr.Column(scale=2, min_width=140):
                        hist_output = gr.HTML(
                            value="",
                            label="Odd vs Even Histogram",
                        )

        # ============================
        # Row 2: image display area
        # ============================
        with gr.Row():
            with gr.Column():
                image_output = gr.HTML(
                    label="Current Collatz Graph",
                )
                gr.Markdown(
                    """
                    **Display tips:**  
                    - In **Zoom & Scroll** mode, use the scrollbars to explore large graphs.  
                    - In **Fit to Width** mode, the graph is scaled to the available width.  
                    - You can right-click the image to open it in a new tab or save it.
                    """
                )

        # Wire buttons: both update the same image + stats
        gen_inverse.click(
            fn=inverse_tree_callback,
            inputs=[backbone_input, branch_input, depth_input, view_mode_inverse],
            outputs=[image_output, stats_output, hist_output],
        )

        gen_minimal.click(
            fn=minimal_subtree_callback,
            inputs=[N_input, view_mode_minimal],
            outputs=[image_output, stats_output, hist_output],
        )

    return demo


demo = build_demo()

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