src/app.py

import streamlit as st
import numpy as np
from PIL import Image
import io
import sys
import os

# Bug 3 & 4 fixes - Top level imports, recursion limit
sys.setrecursionlimit(10000)

sys.path.insert(0, os.path.dirname(__file__))
from quadtree_engine import (
    read_ppm_bytes, process_image, arr_to_pil, compress_image, decompress_image, pad_to_square_pow2
)

st.set_page_config(
    page_title="QuadTree Image Engine",
    page_icon="🌲",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ── CSS ──────────────────────────────────────────────────────────────────────
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=DM+Mono:wght@400;500;700&family=IBM+Plex+Sans:wght@400;500;600&display=swap');

html, body, [data-testid="stAppViewContainer"], .stApp {
    background: #0d0d0d !important;
    color: #f0f0f0 !important;
    font-family: 'IBM Plex Sans', sans-serif !important;
}

header[data-testid="stHeader"] {
    display: none !important;
}


h1, h2, h3, h4, h5, h6, [data-testid="stMarkdownContainer"] h1, [data-testid="stMarkdownContainer"] h2, [data-testid="stMarkdownContainer"] h3 {
    font-family: 'DM Mono', monospace !important;
}

[data-testid="stSidebar"] {
    background: #141414 !important;
    border-right: 1px solid #2a2a2a !important;
}

.stButton > button {
    background-color: #00ff87 !important;
    color: #0d0d0d !important;
    border: none !important;
    border-radius: 4px !important;
    font-family: 'DM Mono', monospace !important;
    font-weight: 700 !important;
    width: 100%;
    margin-top: 1rem;
}

.stButton > button:hover {
    background-color: #00cc6a !important;
    color: #000000 !important;
}

.stDownloadButton > button {
    background-color: transparent !important;
    color: #00ff87 !important;
    border: 1px solid #00ff87 !important;
    border-radius: 4px !important;
    font-family: 'DM Mono', monospace !important;
    font-weight: 700 !important;
}

.stDownloadButton > button:hover {
    background-color: #00ff87 !important;
    color: #0d0d0d !important;
}

.accent-text { color: #00ff87; }
.warning-text { color: #ffb700; }
.error-text { color: #ff4444; }

.step-row {
    display: flex;
    align-items: center;
    margin-bottom: 0.5rem;
    background-color: #141414;
    border: 1px solid #2a2a2a;
    padding: 0.75rem;
}
.step-num {
    color: #00ff87;
    font-family: 'DM Mono', monospace;
    font-weight: 700;
    margin-right: 1rem;
    min-width: 60px;
}

.code-block-custom {
    background-color: #141414;
    border: 1px solid #2a2a2a;
    padding: 1rem;
    font-family: 'DM Mono', monospace;
    color: #f0f0f0;
}
</style>
""", unsafe_allow_html=True)

# ── Helpers ───────────────────────────────────────────────────────────────────

PPM_DIR = os.path.join(os.path.dirname(__file__), "c_project_files")

def list_ppms():
    result = []
    if not os.path.isdir(PPM_DIR):
        return result
    for f in sorted(os.listdir(PPM_DIR)):
        if not f.endswith(".ppm"):
            continue
        fpath = os.path.join(PPM_DIR, f)
        try:
            with open(fpath, "rb") as fh:
                header = fh.read(3)
            if header.startswith(b'P6'):
                result.append(f)
        except Exception:
            pass
    return result

def load_array(uploaded=None, ppm_name=None):
    if uploaded is not None:
        uploaded.seek(0)
        raw = uploaded.read()
    elif ppm_name:
        path = os.path.join(PPM_DIR, ppm_name)
        with open(path, "rb") as f:
            raw = f.read()
    else:
        return None, None

    try:
        arr = np.array(Image.open(io.BytesIO(raw)).convert("RGB"), dtype=np.uint8)
        return arr, raw
    except Exception:
        pass

    try:
        arr = read_ppm_bytes(raw)
        return arr, raw
    except Exception as e:
        fname = uploaded.name if uploaded else ppm_name
        raise ValueError(f"Cannot read '{fname}' as an image.\nDetail: {e}")

def count_nodes_and_leaves(node, depth=0):
    if node is None:
        return 0, 0, 0 # total, leaves, max_d
    if node.is_leaf():
        return 1, 1, depth
    counts = [count_nodes_and_leaves(c, depth+1) for c in [node.topLeft, node.topRight, node.bottomLeft, node.bottomRight]]
    total = 1 + sum(c[0] for c in counts)
    leaves = sum(c[1] for c in counts)
    max_d = max(c[2] for c in counts)
    return total, leaves, max_d

# ── Sidebar ──────────────────────────────────────────────────────────────────

with st.sidebar:
    st.markdown("### ⬆ INPUT")
    
    input_options = ["Upload file"]
    has_ppms = os.path.isdir(PPM_DIR) and len(list_ppms()) > 0
    if has_ppms:
        input_options.append("Use repo PPM")
        
    input_mode = st.radio("Source", input_options, label_visibility="collapsed")
    
    arr1 = None
    file_id = None
    
    if input_mode == "Upload file":
        up1 = st.file_uploader("Primary image", type=["ppm","png","jpg","jpeg"], key="u1")
        if up1:
            file_id = f"{up1.name}_{up1.size}"
            try:
                arr1, _ = load_array(uploaded=up1)
            except ValueError as e:
                st.error(str(e))
    else:
        sel = st.selectbox("Select PPM", list_ppms())
        if sel:
            file_id = f"repo_{sel}"
            try:
                arr1, _ = load_array(ppm_name=sel)
            except ValueError as e:
                st.error(str(e))
                
    # Bug 5: Image Size Guard
    if arr1 is not None:
        MAX_DIM = 2048
        h, w = arr1.shape[:2]
        if h > MAX_DIM or w > MAX_DIM:
            st.markdown(f"<div style='color: #ffb700; font-size: 0.85rem; margin-bottom: 1rem; border: 1px solid #ffb700; padding: 0.5rem;'>⚠️ Image is {w}×{h}. Images over {MAX_DIM}px may be slow or crash. Consider resizing first.</div>", unsafe_allow_html=True)

    # Bug 6: Stale Result Flash (immediately wipe on new file)
    if file_id != st.session_state.get("last_file_id"):
        st.session_state.pop("result", None)
        st.session_state.pop("tree", None)
        st.session_state.pop("op_done", None)
        st.session_state["last_file_id"] = file_id
        
    st.markdown("---")
    st.markdown("### ⚙ OPERATION")
    
    OP_MAP = {
        "Compress Only": "compress_only",
        "Grayscale": "grayscale",
        "Negative": "negative",
        "Sepia": "sepia",
        "Brighten": "brighten",
        "Mirror (Horizontal)": "mirror",
        "Flip (Vertical)": "water",
        "Rotate Left 90°": "rotate_left",
        "Rotate Right 90°": "rotate_right",
        "Blend / Union": "union",
    }
    
    operation = st.selectbox("Operation", list(OP_MAP.keys()), label_visibility="collapsed")
    op_key = OP_MAP[operation]
    
    arr2 = None
    if op_key == "union":
        st.markdown("<div style='font-family: \"DM Mono\", monospace; font-size: 0.85rem;'>Second image for blending:</div>", unsafe_allow_html=True)
        if input_mode == "Upload file":
            up2 = st.file_uploader("Second image", type=["ppm","png","jpg","jpeg"], key="u2")
            if up2:
                try:
                    arr2, _ = load_array(uploaded=up2)
                except ValueError as e:
                    st.error(str(e))
        else:
            sel2 = st.selectbox("Second PPM", list_ppms(), key="sel2")
            if sel2:
                try:
                    arr2, _ = load_array(ppm_name=sel2)
                except ValueError as e:
                    st.error(str(e))

    st.markdown("---")
    st.markdown("### ◎ THRESHOLD")
    threshold = st.slider("Quality vs Compression", 1, 500, 30, label_visibility="collapsed")
    
    if threshold <= 30:
        st.markdown("<div style='color: #00ff87; font-family: \"DM Mono\", monospace; font-weight: bold;'>● LOSSLESS</div>", unsafe_allow_html=True)
    elif threshold <= 100:
        st.markdown("<div style='color: #ffb700; font-family: \"DM Mono\", monospace; font-weight: bold;'>● BALANCED</div>", unsafe_allow_html=True)
    else:
        st.markdown("<div style='color: #ff4444; font-family: \"DM Mono\", monospace; font-weight: bold;'>● LOSSY</div>", unsafe_allow_html=True)
        
    run = st.button("▶ RUN", disabled=(arr1 is None))
    
    st.markdown("<br>", unsafe_allow_html=True)
    download_placeholder = st.empty()


# ── Main Area ────────────────────────────────────────────────────────────────

tab1, tab2, tab3 = st.tabs(["IMAGE VIEW", "ALGORITHM EXPLORER", "C SOURCE"])

with tab1:
    if run and arr1 is not None:
        st.session_state.pop("result", None)
        st.session_state.pop("tree", None)
        
        with st.spinner("Processing..."):
            try:
                if op_key == "compress_only":
                    padded, oh, ow = pad_to_square_pow2(arr1)
                    size = padded.shape[0]
                    tree = compress_image(padded, 0, 0, size, threshold)
                    out = np.zeros((size, size, 3), dtype=np.uint8)
                    decompress_image(tree, out, 0, 0, size)
                    computed = out[:oh, :ow]
                elif op_key == "union":
                    if arr2 is None:
                        st.error("Please provide a second image for union.")
                        computed, tree = None, None
                    else:
                        computed, tree = process_image(arr1, "union", threshold, arr2, return_tree=True)
                else:
                    computed, tree = process_image(arr1, op_key, threshold, return_tree=True)

                if computed is not None:
                    # Bug 1 Fix: Store result and tree together, avoid rebuilding
                    st.session_state["result"] = computed
                    st.session_state["tree"] = tree
                    st.session_state["op_done"] = operation
                    st.session_state["thresh_done"] = threshold
            except RecursionError:
                st.error("RecursionError: image too large for this threshold. Try a higher threshold.")
            except Exception as e:
                st.error(f"Error: {e}")

    result_arr = st.session_state.get("result")
    tree = st.session_state.get("tree")
    
    if arr1 is not None:
        with st.container():
            c1, c2 = st.columns(2)
            with c1:
                st.markdown("<div style='font-family: \"DM Mono\", monospace; margin-bottom: 0.5rem; color: #555;'>BEFORE</div>", unsafe_allow_html=True)
                st.image(arr_to_pil(arr1), use_container_width=True)
                st.markdown(f"<div style='font-family: \"DM Mono\", monospace; color: #555; text-align: center;'>{arr1.shape[1]} × {arr1.shape[0]}</div>", unsafe_allow_html=True)
            with c2:
                st.markdown("<div style='font-family: \"DM Mono\", monospace; margin-bottom: 0.5rem; color: #555;'>AFTER</div>", unsafe_allow_html=True)
                if result_arr is not None:
                    st.image(arr_to_pil(result_arr), use_container_width=True)
                    op_str = st.session_state.get("op_done", "")
                    th_str = st.session_state.get("thresh_done", "")
                    st.markdown(f"<div style='font-family: \"DM Mono\", monospace; color: #555; text-align: center;'>{op_str} · t={th_str}</div>", unsafe_allow_html=True)
                else:
                    st.markdown("<div style='border: 1px solid #2a2a2a; height: 300px; display: flex; align-items: center; justify-content: center; color: #555; font-family: \"DM Mono\", monospace; background-color: #141414;'>[ RESULT IMAGE ]</div>", unsafe_allow_html=True)
        
        if result_arr is not None and tree is not None:
            total_nodes, leaves, max_d = count_nodes_and_leaves(tree)
            total_pixels = arr1.shape[0] * arr1.shape[1]
            ratio = total_pixels / total_nodes if total_nodes > 0 else 0
            
            st.markdown("<div style='font-family: \"DM Mono\", monospace; color: #555; margin-top: 2rem; margin-bottom: 0.5rem;'>QUADTREE STATS</div>", unsafe_allow_html=True)
            
            st.markdown(f"""
            <div style='display: flex; border: 1px solid #2a2a2a; background-color: #141414;'>
                <div style='flex: 1; border-right: 1px solid #2a2a2a; padding: 1rem; text-align: center;'>
                    <div style='color: #555; font-family: "DM Mono", monospace; font-size: 0.8rem; margin-bottom: 0.5rem;'>TOTAL NODES</div>
                    <div style='color: #00ff87; font-family: "DM Mono", monospace; font-size: 1.5rem; font-weight: 700;'>{total_nodes:,}</div>
                    <div style='color: #00ff87; font-size: 0.75rem; font-family: "DM Mono", monospace; margin-top: 4px;'>{ratio:.1f}× fewer nodes</div>
                </div>
                <div style='flex: 1; border-right: 1px solid #2a2a2a; padding: 1rem; text-align: center;'>
                    <div style='color: #555; font-family: "DM Mono", monospace; font-size: 0.8rem; margin-bottom: 0.5rem;'>LEAF NODES</div>
                    <div style='color: #00ff87; font-family: "DM Mono", monospace; font-size: 1.5rem; font-weight: 700;'>{leaves:,}</div>
                </div>
                <div style='flex: 1; border-right: 1px solid #2a2a2a; padding: 1rem; text-align: center;'>
                    <div style='color: #555; font-family: "DM Mono", monospace; font-size: 0.8rem; margin-bottom: 0.5rem;'>MAX DEPTH</div>
                    <div style='color: #00ff87; font-family: "DM Mono", monospace; font-size: 1.5rem; font-weight: 700;'>{max_d}</div>
                </div>
                <div style='flex: 1; border-right: 1px solid #2a2a2a; padding: 1rem; text-align: center;'>
                    <div style='color: #555; font-family: "DM Mono", monospace; font-size: 0.8rem; margin-bottom: 0.5rem;'>TOTAL PIXELS</div>
                    <div style='color: #00ff87; font-family: "DM Mono", monospace; font-size: 1.5rem; font-weight: 700;'>{total_pixels:,}</div>
                </div>
                <div style='flex: 1; padding: 1rem; text-align: center;'>
                    <div style='color: #555; font-family: "DM Mono", monospace; font-size: 0.8rem; margin-bottom: 0.5rem;'>THRESHOLD</div>
                    <div style='color: #00ff87; font-family: "DM Mono", monospace; font-size: 1.5rem; font-weight: 700;'>{st.session_state.get('thresh_done')}</div>
                </div>
            </div>
            """, unsafe_allow_html=True)
            
            with download_placeholder:
                img_pil = arr_to_pil(result_arr)
                buf = io.BytesIO()
                img_pil.save(buf, format="PNG")
                st.download_button("↓ DOWNLOAD RESULT", buf.getvalue(), "result.png", "image/png")
                
    else:
        st.markdown("<div style='border: 1px solid #2a2a2a; padding: 6rem; text-align: center; color: #555; font-family: \"DM Mono\", monospace; background-color: #141414;'>[ AWAITING INPUT ]</div>", unsafe_allow_html=True)


with tab2:
    st.markdown("### 1. What is a Quadtree?")
    c1, c2 = st.columns(2)
    with c1:
        st.code("""
┌────────┬────────┐
│        │        │
│   TL   │   TR   │
│        │        │
├────────┼────────┤
│        │        │
│   BL   │   BR   │
│        │        │
└────────┴────────┘
""", language="text")
    with c2:
        st.markdown("""
- **`red`, `green`, `blue`**: The average color of this region.
- **`area`**: The total pixels covered by this node.
- **Children**: If variance > threshold, the region splits into 4 sub-regions (`topLeft`, `topRight`, `bottomLeft`, `bottomRight`).
- **Leaf Node**: If variance <= threshold, children are null. The region is colored uniformly.
""")

    st.markdown("---")
    st.markdown("### 2. How Compression Works")
    st.markdown("<div style='color: #555; margin-bottom: 1rem;'>See how threshold controls quality vs. compression</div>", unsafe_allow_html=True)
    
    @st.cache_data
    def get_synthetic_image():
        y, x = np.mgrid[0:64, 0:64]
        r = (x * 4) % 255
        g = (y * 4) % 255
        b = ((x + y) * 2) % 255
        noise = np.random.randint(0, 50, (64, 64, 3))
        img = np.stack([r, g, b], axis=-1) + noise
        return np.clip(img, 0, 255).astype(np.uint8)

    synth_img = get_synthetic_image()
    sim_thresh = st.slider("Simulator Threshold", 1, 200, 50, key="sim_t")
    
    sc1, sc2 = st.columns(2)
    with sc1:
        st.image(arr_to_pil(synth_img), caption="Original (64x64)", use_container_width=True)
    with sc2:
        sim_tree = compress_image(synth_img, 0, 0, 64, sim_thresh)
        sim_out = np.zeros((64, 64, 3), dtype=np.uint8)
        decompress_image(sim_tree, sim_out, 0, 0, 64)
        total_sim, leaves_sim, _ = count_nodes_and_leaves(sim_tree)
        st.image(arr_to_pil(sim_out), caption=f"Compressed (Nodes: {total_sim})", use_container_width=True)

    st.markdown("---")
    st.markdown("### 3. Compression Algorithm Step-by-Step")
    st.markdown("""
<div class="step-row"><span class="step-num">STEP 1</span><span>Divide image into 4 quadrants</span></div>
<div class="step-row"><span class="step-num">STEP 2</span><span>Compute mean RGB + variance for each quadrant</span></div>
<div class="step-row"><span class="step-num">STEP 3</span><span>If variance ≤ threshold: <b>LEAF</b> → store avg color, stop subdividing</span></div>
<div class="step-row"><span class="step-num">STEP 4</span><span>If variance > threshold: <b>RECURSE</b> into each quadrant with size/2</span></div>
<div class="step-row"><span class="step-num">STEP 5</span><span>Continue until size=1 pixel or variance ≤ threshold</span></div>
""", unsafe_allow_html=True)

    st.markdown("---")
    st.markdown("### 4. Color Filters — How They Work on the Tree")
    fc1, fc2, fc3, fc4 = st.columns(4)
    with fc1:
        with st.container(border=True):
            st.markdown("**Grayscale**")
            st.code("L = 0.299R + 0.587G + 0.114B\nR'=L, G'=L, B'=L")
            st.caption("Weights green channel most heavily (human eye is most sensitive to green)")
    with fc2:
        with st.container(border=True):
            st.markdown("**Negative**")
            st.code("R' = 255 - R\nG' = 255 - G\nB' = 255 - B")
            st.caption("Inverts each channel — dark becomes light, colors become complementary")
    with fc3:
        with st.container(border=True):
            st.markdown("**Sepia**")
            st.code("R' = 0.393R + 0.769G + 0.189B\nG' = ...\nB' = ...")
            st.caption("Warm brownish tones by mixing channels — mimics aged photographic paper")
    with fc4:
        with st.container(border=True):
            st.markdown("**Brighten**")
            st.code("R' = min(255, R*1.3)\nG' = min(255, G*1.3)\nB' = min(255, B*1.3)")
            st.caption("Scales all channels up — clips at 255 to avoid overflow")

    st.markdown("---")
    st.markdown("### 5. Spatial Transforms — Pointer Swaps")
    st.markdown("<div class='accent-text'><strong>No pixel data is ever copied. Only 4 pointer assignments per node.</strong></div><br>", unsafe_allow_html=True)
    tc1, tc2 = st.columns(2)
    with tc1:
        st.code("""MIRROR (horizontal):
Before:  TL | TR
         -------
         BL | BR

After:   TR | TL
         -------
         BR | BL""", language="text")
    with tc2:
        st.code("""ROTATE LEFT 90°:
Before:  TL | TR
         -------
         BL | BR

After:   TR | BR
         -------
         TL | BL""", language="text")
         
    tc3, tc4 = st.columns(2)
    with tc3:
        st.code("""FLIP (vertical):
Before:  TL | TR
         -------
         BL | BR

After:   BL | BR
         -------
         TL | TR""", language="text")
    with tc4:
        st.code("""ROTATE RIGHT 90°:
Before:  TL | TR
         -------
         BL | BR

After:   BL | TL
         -------
         BR | TR""", language="text")

    st.markdown("---")
    st.markdown("### 6. Union / Blend — Three Cases")
    st.markdown("""
- **Case 1:** Both nodes are internal → recurse into all 4 child pairs.
- **Case 2:** t1 is a leaf, t2 has children → blend t1's solid color with each of t2's children.
- **Case 3:** t2 is a leaf, t1 has children → blend t2's solid color with each of t1's children.

Averaging formula: `result.R = (t1.R + t2.R) / 2`
*Note: This produces a pixel-perfect 50/50 blend without ever decompressing either image to a pixel buffer.*
""")

    st.markdown("---")
    st.markdown("### 7. Complexity Analysis")
    st.markdown("""
| Operation | Time Complexity | Space Complexity | Notes |
|---|---|---|---|
| Compress | O(n log n) | O(n) | n = total pixels |
| Decompress | O(n) | O(n) | Linear tree traversal |
| Filter | O(k) | O(k) | k = tree nodes, k ≪ n |
| Rotate/Mirror | O(k) | O(1) | Only pointer swaps |
| Union | O(min(k1,k2)) | O(min(k1,k2)) | Bounded by smaller tree |

<div style="border-left: 4px solid #00ff87; padding-left: 1rem; margin-top: 1rem; color: #f0f0f0;">
<strong>Key Takeaway:</strong> Filters and transforms run on the compressed tree — they're O(nodes) not O(pixels). At threshold=100, a 512×512 image (262K pixels) may have fewer than 5,000 nodes.
</div>
""", unsafe_allow_html=True)


with tab3:
    st.markdown("### C Source vs Python")
    st.markdown("""
    <div style="background-color: #141414; border: 1px solid #2a2a2a; padding: 1rem; margin-bottom: 1rem;">
    🔗 View the full C implementation on GitHub:<br>
    <a href="https://github.com/Harshwardhan-Deshmukh03/Image-Manipulation-QuadTree.git" target="_blank"
       style="color:#00ff87;font-weight:600;text-decoration:none;">
       github.com/Harshwardhan-Deshmukh03/Image-Manipulation-QuadTree
    </a>
    </div>
    """, unsafe_allow_html=True)
    st.markdown("""
| C File | Python Equivalent |
|---|---|
| `compress.c` | `compress_image()` |
| `filters.c` | `apply_grayscale()`, `apply_negative()`, etc. |
| `rotate.c` | `rotate_left()`, `get_mirror_image()`, etc. |
| `union.c` | `union_of_images()` |
| `decompress.c` | `decompress_image()` |
""")
    
    c_files = {
        "main.c": "CLI entry point — parses flags, orchestrates the full pipeline",
        "compress.c": "Quadtree construction from pixel matrix using variance threshold",
        "decompress.c": "Quadtree → pixel matrix reconstruction",
        "filters.c": "Color filters: grayscale, negative, sepia, brighten",
        "rotate.c": "Spatial transforms: mirror, flip, rotate L/R/180",
        "union.c": "Pixel-level blending of two Quadtrees",
        "suppl.c": "PPM I/O, getMean(), tree serialization helpers",
        "suppl.h": "All struct definitions: pixels, qtNode, qtInfo",
    }

    st.markdown("<br>", unsafe_allow_html=True)
    for fname, desc in c_files.items():
        fpath = os.path.join(PPM_DIR, fname)
        with st.expander(f"`{fname}` — {desc}"):
            if os.path.isfile(fpath):
                with open(fpath) as f:
                    st.code(f.read(), language="c")
            else:
                st.info(f"File not found at: {fpath}")