UI_testing.py

"""
Correlation Matrix Scaling Study — Interactive Dashboard
=========================================================
User enters their RAM. Dashboard shows:
  - Theoretical memory/time at any N
  - Actual measured values (run live from imported modules)
  - Memory wall: at what N does it exceed their RAM
  - Method comparison dropdown

All computation imported from existing modules — nothing hardcoded.

Run from repo root:  python -m streamlit run UI_testing.py
"""

import sys, os, gc, time
import streamlit as st
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pandas as pd
import psutil

# ── Path setup: import from existing project modules ─────────────────
ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(ROOT, "CoreProblem"))
sys.path.insert(0, os.path.join(ROOT, "AdditionalMilestones", "Blockwise-study"))
sys.path.insert(0, os.path.join(ROOT, "AdditionalMilestones", "Streaming-study"))
sys.path.insert(0, os.path.join(ROOT, "AdditionalMilestones", "standardisation"))
sys.path.insert(0, os.path.join(ROOT, "AdditionalMilestones", "correlation_approximation"))

from utils import generate_data, measure_correlation
from blockUtils import measure_correlation_blockwise
from streamUtils import measure_correlation_streaming
from standutils import correlation_on_the_fly
from corrutils import approximate_correlation_sampling

# ── Method registry ──────────────────────────────────────────────────
METHODS = {
    "Baseline (np.corrcoef)": {
        "fn": lambda data: measure_correlation(data),
        "needs_full_nxn": True,
        "color": "#1f77b4",
        "desc": "Standard NumPy — computes the full NxN matrix in one call. "
                "Fastest (uses optimised BLAS), but internally allocates ~3-4 NxN "
                "temporary arrays (centered data, covariance, normalisation). "
                "This is what hits the memory wall first.",
    },
    "Blockwise (Tiled)": {
        "fn": lambda data: measure_correlation_blockwise(data, tile_size=1000),
        "needs_full_nxn": True,
        "color": "#2ca02c",
        "desc": "Z-scores all rows once, then fills the NxN output tile-by-tile. "
                "Only one small tile is in memory at a time (besides the output). "
                "Reduces peak memory from ~3-4× to ~1-2× the NxN matrix size.",
    },
    "Streaming (Chunked)": {
        "fn": lambda data: measure_correlation_streaming(data, chunk_size=1000),
        "needs_full_nxn": True,
        "color": "#ff7f0e",
        "desc": "Two-pass approach: computes the global mean in pass 1, then "
                "accumulates covariance chunk-by-chunk along the sample axis in pass 2. "
                "Reduces intermediate memory when T (samples) is large.",
    },
    "On-the-fly Standardisation": {
        "fn": lambda data: correlation_on_the_fly(data),
        "needs_full_nxn": True,
        "color": "#d62728",
        "desc": "Z-scores each row only when needed for a dot product, then discards it. "
                "Eliminates the NxT z-score matrix entirely. Trades memory for speed — "
                "very slow due to Python double loop (O(N²) iterations).",
    },
    "Approximate (Random Sampling)": {
        "fn": lambda data: approximate_correlation_sampling(data, num_samples=10000),
        "needs_full_nxn": False,
        "color": "#9467bd",
        "desc": "Instead of computing all N×(N-1)/2 pairs, randomly samples K=10,000 "
                "pairs. Memory is O(K), completely independent of N. Gives the "
                "distribution of correlations, not the full matrix.",
    },
}

T_DEFAULT = 500  # fixed T used across all benchmarks


# ── Helper: theoretical memory (bytes) for NxN float64 ───────────────
def theory_mem_bytes(n):
    return 8 * n * n

def theory_mem_mb(n):
    return theory_mem_bytes(n) / (1024 ** 2)

def theory_mem_gb(n):
    return theory_mem_bytes(n) / (1024 ** 3)

# ── Helper: find wall N (largest N before exceeding RAM) ─────────────
def find_wall_n(ram_gb, safety=0.85):
    """Return the N where NxN matrix alone hits safety*RAM."""
    usable_bytes = ram_gb * (1024 ** 3) * safety
    # 8*N^2 = usable_bytes => N = sqrt(usable_bytes/8)
    return int(np.sqrt(usable_bytes / 8))


# ── Page config ──────────────────────────────────────────────────────
st.set_page_config(page_title="Correlation Scaling Study", page_icon="📊", layout="wide")
st.title("📊 Correlation Matrix Scaling Study")
st.markdown("""
**The problem:** Computing a Pearson correlation matrix for N time series produces an N×N matrix.  
Memory needed = **8 × N² bytes** (float64). This grows quadratically — at N=10,000 it's 763 MB,  
at N=50,000 it's 18.6 GB. At some point your RAM simply runs out. That's the **memory wall**.

This dashboard lets you explore:
1. **Where is your memory wall?** — Enter your RAM and see the maximum N you can handle  
2. **How do different methods compare?** — Run benchmarks live or view pre-computed results  
3. **Are the alternative methods correct?** — Validate each method against NumPy's baseline
""")

# ── Session state ────────────────────────────────────────────────────
if "run_results" not in st.session_state:
    st.session_state.run_results = None

st.divider()

# ════════════════════════════════════════════════════════════════════
# SIDEBAR — user inputs
# ════════════════════════════════════════════════════════════════════
st.sidebar.header("⚙️ Configuration")

user_ram_gb = st.sidebar.selectbox(
    "Your device RAM (GB)",
    options=[2, 4, 8, 16, 24, 32, 64, 128],
    index=2,  # default 8 GB
    help="How much RAM your machine has. We use 85% as the safe usable limit."
)

user_n = st.sidebar.number_input(
    "Number of series (N)",
    min_value=10, max_value=100000, value=5000, step=500,
    help="The N for the NxN correlation matrix."
)

T = st.sidebar.slider("Time samples (T)", 100, 2000, T_DEFAULT, 100)

selected_methods = st.sidebar.multiselect(
    "Methods to compare",
    options=list(METHODS.keys()),
    default=["Baseline (np.corrcoef)", "Blockwise (Tiled)", "Streaming (Chunked)", "Approximate (Random Sampling)"],
)

if "On-the-fly Standardisation" in selected_methods and user_n > 500:
    st.sidebar.warning("⚠️ On-the-fly at N > 500 is very slow (Python double-loop).")

run_btn = st.sidebar.button("▶ Run Benchmark at this N", type="primary")

# Show method descriptions
st.sidebar.divider()
st.sidebar.subheader("📖 Method Descriptions")
for name in selected_methods:
    st.sidebar.markdown(f"**{name}**")
    st.sidebar.caption(METHODS[name]["desc"])

st.sidebar.divider()
st.sidebar.markdown(f"**Available RAM now:** {psutil.virtual_memory().available / (1024**3):.1f} GB free")

# ════════════════════════════════════════════════════════════════════
# SECTION 1 — THEORETICAL ANALYSIS
# ════════════════════════════════════════════════════════════════════
st.header("📐 Theoretical Analysis")
st.markdown("""  
Before running any code, we can predict memory usage from first principles:  
an N×N correlation matrix stores N² entries, each 8 bytes (float64). So **theoretical minimum memory = 8N² bytes**.  
In practice, `np.corrcoef` allocates additional temporary arrays internally, so actual peak is ~3-4× this minimum.
""")

wall_n = find_wall_n(user_ram_gb)
t_mem_mb = theory_mem_mb(user_n)
t_mem_display = f"{t_mem_mb:.1f} MB" if t_mem_mb < 1024 else f"{t_mem_mb/1024:.2f} GB"
fits = theory_mem_gb(user_n) <= user_ram_gb * 0.85

c1, c2, c3 = st.columns(3)
c1.metric("Theoretical NxN memory", t_mem_display, help="8 × N² bytes (float64)")
c2.metric("Memory wall N (85% of RAM)", f"{wall_n:,}", help="Largest N where the NxN matrix alone fits in your RAM")
c3.metric("Your RAM", f"{user_ram_gb} GB", delta="safe limit: " + f"{user_ram_gb*0.85:.1f} GB")

if not fits:
    st.error(f"⚠️ At N={user_n:,}, the NxN matrix alone needs **{t_mem_display}** — exceeds your {user_ram_gb} GB RAM. Only **Approximate (sampling)** can run here.")
else:
    st.success(f"✅ N={user_n:,} fits in {user_ram_gb} GB RAM. Theoretical minimum: {t_mem_display}.")

# ── Theoretical scaling curve ────────────────────────────────────────
st.subheader("Theoretical Memory vs N (across all methods)")

n_range = np.array([10, 50, 100, 500, 1000, 3000, 5000, 6000, 8000, 10000])
n_smooth = np.linspace(10, max(user_n * 1.2, 10000), 400)

fig, ax = plt.subplots(figsize=(11, 5))

# Theory curve
ax.plot(n_smooth, [theory_mem_mb(n) for n in n_smooth],
        'k--', linewidth=1.5, alpha=0.5, label="Theory: 8N² (float64)")

# RAM limit line
ram_limit_mb = user_ram_gb * 1024 * 0.85
ax.axhline(y=ram_limit_mb, color='red', linestyle='--', linewidth=1.5, alpha=0.7,
           label=f"Your RAM limit ({user_ram_gb} GB × 85% = {ram_limit_mb/1024:.1f} GB)")

# Wall N marker
ax.axvline(x=wall_n, color='orange', linestyle=':', linewidth=2,
           label=f"Memory wall N={wall_n:,}")

# User's N marker
ax.axvline(x=user_n, color='blue', linestyle=':', linewidth=2,
           label=f"Your N={user_n:,}")

ax.set_xlabel("N (number of series)", fontsize=12)
ax.set_ylabel("Theoretical Memory (MB)", fontsize=12)
ax.set_title("O(N²) Memory Scaling", fontsize=13)
ax.legend(fontsize=9)
ax.grid(True, alpha=0.3)
st.pyplot(fig)
plt.close()

st.divider()

# ════════════════════════════════════════════════════════════════════
# SECTION 2 — BENCHMARK RESULTS (tabbed: live vs pre-computed)
# ════════════════════════════════════════════════════════════════════
st.header("🔬 Benchmark Results")

tab_live, tab_precomp = st.tabs(["🔴 Live Compute", "📦 Pre-computed Results"])

# ── Pre-computed results tab ─────────────────────────────────────────
with tab_precomp:
    st.markdown("""
    These plots were generated from full benchmark runs on the development machine.  
    They cover a wide range of N values and are useful for seeing scaling trends  
    without waiting for live computation.
    """)

    RESULTS_DIR = os.path.join(ROOT, "results")
    PRECOMP_SECTIONS = {
        "Core Problem (Baseline np.corrcoef)": {
            "dir": RESULTS_DIR,
            "plots": ["scaling_loglog.png", "scaling_memory.png", "memory_breakdown.png"],
            "info": "Baseline using `np.corrcoef`. Shows O(N²) scaling in both memory and runtime. "
                    "The memory breakdown reveals how the NxN output and NxT input matrices contribute.",
        },
        "Blockwise (Tiled) Study": {
            "dir": os.path.join(RESULTS_DIR, "Blockwise-study"),
            "plots": ["scaling_loglog_blockwise.png", "scaling_memory_blockwise.png",
                      "memory_breakdown_blockwise.png", "scaling_memory_comparison.png"],
            "info": "Computes tiles of the NxN matrix one at a time, reducing intermediate memory. "
                    "The comparison plot shows blockwise peak vs baseline peak at each N.",
        },
        "Streaming (Chunked) Study": {
            "dir": os.path.join(RESULTS_DIR, "Streaming-study"),
            "plots": ["scaling_loglog_streaming.png", "scaling_memory_streaming.png",
                      "memory_breakdown_streaming.png"],
            "info": "Processes data in chunks along the sample axis using a two-pass approach. "
                    "Pass 1 computes global means, pass 2 accumulates covariance chunk by chunk.",
        },
        "On-the-fly Standardisation Study": {
            "dir": os.path.join(RESULTS_DIR, "Standardisation-study"),
            "plots": ["scaling_loglog_onthefly.png", "scaling_memory_onthefly.png",
                      "memory_breakdown_onthefly.png"],
            "info": "Z-scores each row on demand instead of pre-computing a full NxT z-score matrix. "
                    "Saves the NxT allocation but is very slow due to a Python double loop.",
        },
        "Approximate (Random Sampling) Study": {
            "dir": os.path.join(RESULTS_DIR, "Approximate-study"),
            "plots": ["scaling_loglog_approximate.png", "scaling_memory_approximate.png",
                      "memory_breakdown_approximate.png", "accuracy_vs_samples_approximate.png",
                      "correlation_distribution_approximate.png"],
            "info": "Randomly samples K pairs instead of computing all N(N-1)/2. "
                    "Memory is O(K), independent of N. Shows how estimation accuracy improves with K.",
        },
        "GC Overhead Study": {
            "dir": os.path.join(RESULTS_DIR, "GarbageCollection-study"),
            "plots": ["scaling_loglog_gc.png", "scaling_memory_gc.png",
                      "memory_breakdown_gc.png"],
            "info": "Compares tracemalloc readings with and without `gc.collect()` before benchmarks. "
                    "Shows that leftover garbage can inflate memory readings in repeated measurements.",
        },
        "psutil Measurement Study": {
            "dir": os.path.join(RESULTS_DIR, "PSUTIL-study"),
            "plots": ["scaling_loglog_psutil.png", "scaling_memory_psutil.png",
                      "memory_breakdown_psutil.png"],
            "info": "Uses psutil (process RSS) instead of tracemalloc. Reports much higher values because "
                    "RSS includes the Python interpreter and all loaded libraries, not just the computation.",
        },
    }

    for section_name, sec in PRECOMP_SECTIONS.items():
        with st.expander(section_name, expanded=False):
            st.caption(sec.get("info", ""))
            found_any = False
            cols = st.columns(2)
            for idx, fname in enumerate(sec["plots"]):
                fpath = os.path.join(sec["dir"], fname)
                if os.path.isfile(fpath):
                    cols[idx % 2].image(fpath, caption=fname.replace("_", " ").replace(".png", "").title())
                    found_any = True
            if not found_any:
                st.info("No pre-computed plots found for this study.")

# ── Live compute tab ─────────────────────────────────────────────────
with tab_live:
    st.warning("""
    ⚡ **Live Compute Disclaimer**  
    This mode runs the actual benchmark functions in real time on **your machine**.  
    - Results depend on your hardware, current system load, and available RAM.  
    - Large N values (>5000) may take significant time or cause memory errors.  
    - For consistent reference results, see the **📦 Pre-computed Results** tab.
    """)
    st.markdown(f"Runs selected methods at **N={user_n:,}, T={T}** using the real imported functions from `CoreProblem/` and `AdditionalMilestones/`.")

    if run_btn:
        if not selected_methods:
            st.warning("Select at least one method.")
        else:
            available_gb = psutil.virtual_memory().available / (1024 ** 3)
            needed_data_gb = (user_n * T * 8) / (1024 ** 3)
            needed_nxn_gb = theory_mem_gb(user_n)

            rows = []
            progress = st.progress(0)

            for i, name in enumerate(selected_methods):
                progress.progress((i + 1) / len(selected_methods), text=f"Running {name}...")
                method = METHODS[name]

                # Safety check
                if method["needs_full_nxn"] and needed_nxn_gb > available_gb * 0.85:
                    rows.append({
                        "Method": name,
                        "Actual Memory (MB)": "🛑 Blocked",
                        "Actual Runtime (s)": "—",
                        "Theoretical Memory (MB)": f"{theory_mem_mb(user_n):.1f}",
                        "Fits in RAM?": "❌",
                    })
                    continue

                try:
                    data = generate_data(user_n, T, 42)
                    result, peak_bytes, elapsed = method["fn"](data)
                    del result, data
                    gc.collect()

                    rows.append({
                        "Method": name,
                        "Actual Memory (MB)": round(peak_bytes / (1024 ** 2), 2),
                        "Actual Runtime (s)": round(elapsed, 4),
                        "Theoretical Memory (MB)": round(theory_mem_mb(user_n), 2),
                        "Fits in RAM?": "✅" if needed_nxn_gb <= user_ram_gb * 0.85 else "❌",
                    })
                except MemoryError:
                    rows.append({
                        "Method": name,
                        "Actual Memory (MB)": "❌ OOM",
                        "Actual Runtime (s)": "—",
                        "Theoretical Memory (MB)": f"{theory_mem_mb(user_n):.1f}",
                        "Fits in RAM?": "❌",
                    })
                    gc.collect()
                except Exception as e:
                    rows.append({
                        "Method": name,
                        "Actual Memory (MB)": f"Error",
                        "Actual Runtime (s)": str(e)[:50],
                        "Theoretical Memory (MB)": f"{theory_mem_mb(user_n):.1f}",
                        "Fits in RAM?": "❌",
                    })

            progress.empty()
            st.session_state.run_results = pd.DataFrame(rows)

    # Display results
    if st.session_state.run_results is not None:
        df = st.session_state.run_results
        st.dataframe(df, width='stretch')

        # Plot actual vs theoretical
        numeric_df = df[df["Actual Memory (MB)"].apply(lambda x: isinstance(x, (int, float)))].copy()
        if not numeric_df.empty:
            fig2, axes = plt.subplots(1, 2, figsize=(14, 5))

            colors = [METHODS.get(m, {}).get("color", "#333") for m in numeric_df["Method"]]
            methods_list = numeric_df["Method"].tolist()
            actual_mem = numeric_df["Actual Memory (MB)"].tolist()
            theory_vals = [theory_mem_mb(user_n)] * len(methods_list)
            runtimes = numeric_df["Actual Runtime (s)"].tolist()

            x = np.arange(len(methods_list))
            w = 0.35

            axes[0].bar(x - w/2, actual_mem, w, label="Actual (measured)", color=colors, alpha=0.85)
            axes[0].bar(x + w/2, theory_vals, w, label="Theoretical (8N²)", color="lightgray", edgecolor="black")
            axes[0].axhline(y=user_ram_gb * 1024 * 0.85, color='red', linestyle='--', linewidth=1.5,
                            label=f"RAM limit ({user_ram_gb}GB × 85%)")
            axes[0].set_xticks(x)
            axes[0].set_xticklabels(methods_list, rotation=20, ha='right', fontsize=8)
            axes[0].set_ylabel("Memory (MB)")
            axes[0].set_title(f"Actual vs Theoretical Memory @ N={user_n:,}")
            axes[0].legend(fontsize=8)
            axes[0].grid(True, axis='y', alpha=0.3)

            axes[1].bar(x, runtimes, color=colors, alpha=0.85)
            axes[1].set_xticks(x)
            axes[1].set_xticklabels(methods_list, rotation=20, ha='right', fontsize=8)
            axes[1].set_ylabel("Runtime (s)")
            axes[1].set_title(f"Runtime @ N={user_n:,}")
            axes[1].grid(True, axis='y', alpha=0.3)

            fig2.tight_layout()
            st.pyplot(fig2)
            plt.close()
    else:
        st.info("👈 Configure options in the sidebar and click **▶ Run Benchmark** to see actual measured values.")

st.divider()

# ════════════════════════════════════════════════════════════════════
# SECTION 3 — MEMORY WALL ACROSS N (scaling view)
# ════════════════════════════════════════════════════════════════════
st.header("🧱 Memory Wall: Scaling Across N")
st.markdown(f"Shows where your **{user_ram_gb} GB RAM** becomes the bottleneck as N grows.")

n_plot = np.array([10, 50, 100, 500, 1000, 3000, 5000, 6000, 8000, 10000])
theory_curve = np.array([theory_mem_mb(n) for n in n_plot])
ram_limit_mb_line = user_ram_gb * 1024 * 0.85

fig3, ax3 = plt.subplots(figsize=(11, 5))
ax3.plot(n_plot, theory_curve, 'ko-', linewidth=2, markersize=5, label="Theoretical NxN memory (8N²)")
ax3.axhline(y=ram_limit_mb_line, color='red', linestyle='--', linewidth=2,
            label=f"Your RAM limit ({user_ram_gb} GB × 85% = {ram_limit_mb_line/1024:.1f} GB)")
ax3.axvline(x=wall_n, color='orange', linestyle=':', linewidth=2,
            label=f"Wall N = {wall_n:,}")
ax3.fill_between(n_plot, theory_curve,
                 where=theory_curve >= ram_limit_mb_line,
                 color='red', alpha=0.12, label="Exceeds RAM")

# Mark user's N
if user_n <= n_plot[-1]:
    ax3.axvline(x=user_n, color='blue', linestyle=':', linewidth=1.5,
                label=f"Your N = {user_n:,}")

ax3.set_xlabel("N", fontsize=12)
ax3.set_ylabel("Memory (MB)", fontsize=12)
ax3.set_title("Memory Wall — When does your RAM become the bottleneck?", fontsize=13)
ax3.legend(fontsize=9)
ax3.grid(True, alpha=0.3)
st.pyplot(fig3)
plt.close()

# Reference table
st.subheader(f"Memory at common N values on your {user_ram_gb} GB machine")
ref_rows = []
for n in [100, 500, 1000, 3000, 5000, 10000, 20000, 50000, 100000]:
    mem_mb = theory_mem_mb(n)
    mem_gb = mem_mb / 1024
    fits_flag = "✅ fits" if mem_gb <= user_ram_gb * 0.85 else "❌ exceeds"
    ref_rows.append({
        "N": f"{n:,}",
        "N² entries": f"{n*n:,}",
        "Matrix memory": f"{mem_mb:.1f} MB" if mem_mb < 1024 else f"{mem_gb:.2f} GB",
        f"On {user_ram_gb}GB RAM?": fits_flag,
    })
st.table(pd.DataFrame(ref_rows))

st.divider()

# ════════════════════════════════════════════════════════════════════
# SECTION 4 — CORRECTNESS VALIDATION
# ════════════════════════════════════════════════════════════════════
st.header("✅ Correctness Validation")
st.markdown("Verify each method matches `np.corrcoef` — using the actual imported functions.")

val_n = st.number_input("Validation N", min_value=10, max_value=2000, value=100, step=10)
val_T = st.number_input("Validation T", min_value=100, max_value=2000, value=500, step=100)
validate_methods = st.multiselect(
    "Methods to validate",
    options=["Blockwise (Tiled)", "Streaming (Chunked)", "On-the-fly Standardisation"],
    default=["Blockwise (Tiled)", "Streaming (Chunked)"],
)

if st.button("▶ Run Validation"):
    data = generate_data(val_n, val_T, 42)
    c_base, _, _ = measure_correlation(data)

    val_rows = []
    for name in validate_methods:
        c_test, _, elapsed = METHODS[name]["fn"](data)
        max_diff = np.max(np.abs(c_base - c_test))
        val_rows.append({
            "Method": name,
            "Max |diff| vs np.corrcoef": f"{max_diff:.2e}",
            "np.allclose": "✅ PASS" if np.allclose(c_base, c_test) else "❌ FAIL",
            "Runtime (s)": round(elapsed, 4),
        })
        del c_test
        gc.collect()

    del data, c_base
    gc.collect()

    st.dataframe(pd.DataFrame(val_rows), width='stretch')

# ── Footer ────────────────────────────────────────────────────────────
st.divider()
st.caption("""
**Source modules:** `CoreProblem/utils.py` · `Blockwise-study/blockUtils.py` · 
`Streaming-study/streamUtils.py` · `standardisation/standutils.py` · 
`correlation_approximation/corrutils.py`  
All computation is live — zero hardcoded numbers. Theoretical values use 8×N² bytes (float64).  
**Repo:** [github.com/so-slay/DS_CorrelationEstimation](https://github.com/so-slay/DS_CorrelationEstimation)
""")