#!/usr/bin/env python3
"""Generate report diagrams for LDPC optical decoder project.

Creates:
  data/plots/system_architecture.png - Block diagram of the system architecture
  data/plots/channel_model.png       - Poisson channel model and LLR computation
"""

import os
import sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch
from scipy.stats import poisson

# Global style
plt.rcParams.update({
    "font.size": 13,
    "figure.dpi": 150,
    "savefig.bbox": "tight",
    "font.family": "sans-serif",
    "axes.spines.top": False,
    "axes.spines.right": False,
})

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_DIR = os.path.dirname(SCRIPT_DIR)
PLOT_DIR = os.path.join(PROJECT_DIR, "data", "plots")
os.makedirs(PLOT_DIR, exist_ok=True)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def rounded_box(ax, xy, width, height, label, facecolor="white",
                edgecolor="black", fontsize=11, fontweight="normal",
                linewidth=1.5, text_color="black", pad=0.02, zorder=2):
    """Draw a rounded rectangle with centered text."""
    box = FancyBboxPatch(
        xy, width, height,
        boxstyle="round,pad=0.02",
        facecolor=facecolor,
        edgecolor=edgecolor,
        linewidth=linewidth,
        zorder=zorder,
    )
    ax.add_patch(box)
    cx = xy[0] + width / 2
    cy = xy[1] + height / 2
    ax.text(cx, cy, label, ha="center", va="center",
            fontsize=fontsize, fontweight=fontweight, color=text_color,
            zorder=zorder + 1)
    return box


def arrow(ax, xy_from, xy_to, color="black", lw=1.8, style="->",
          connectionstyle="arc3,rad=0", zorder=3):
    """Draw an arrow between two points."""
    a = FancyArrowPatch(
        xy_from, xy_to,
        arrowstyle=style,
        connectionstyle=connectionstyle,
        color=color,
        linewidth=lw,
        mutation_scale=18,
        zorder=zorder,
    )
    ax.add_patch(a)
    return a


# ---------------------------------------------------------------------------
# Diagram 1: System Architecture
# ---------------------------------------------------------------------------

def generate_system_architecture(save_path):
    fig, ax = plt.subplots(figsize=(16, 9))
    ax.set_xlim(-0.5, 15.5)
    ax.set_ylim(-1.5, 9.5)
    ax.set_aspect("equal")
    ax.axis("off")
    fig.patch.set_facecolor("white")

    # Title
    ax.text(7.5, 9.0, "LDPC Decoder System Architecture",
            ha="center", va="center", fontsize=18, fontweight="bold",
            color="#1a1a2e")
    ax.text(7.5, 8.4, "Photon-Starved Optical Communication  |  Caravel SoC (SkyWater 130nm)",
            ha="center", va="center", fontsize=11, color="#555555")

    # ---- Top-level blocks (left to right) ----

    # Photon Detector
    rounded_box(ax, (-0.3, 5.0), 2.6, 1.6, "Photon\nDetector",
                facecolor="#e8f5e9", edgecolor="#388e3c", fontsize=13,
                fontweight="bold", linewidth=2)
    ax.text(1.0, 4.4, "Single-photon\ndetector array",
            ha="center", va="top", fontsize=9, color="#388e3c", style="italic")

    # LLR Computer
    rounded_box(ax, (3.5, 5.0), 2.6, 1.6, "LLR Computer\n(PicoRV32)",
                facecolor="#fff3e0", edgecolor="#f57c00", fontsize=12,
                fontweight="bold", linewidth=2)
    ax.text(4.8, 4.4, "Quantize to\n6-bit LLR",
            ha="center", va="top", fontsize=9, color="#f57c00", style="italic")

    # LDPC Decoder (large block with sub-blocks)
    decoder_x, decoder_y = 7.0, 0.8
    decoder_w, decoder_h = 6.0, 6.8
    decoder_bg = FancyBboxPatch(
        (decoder_x, decoder_y), decoder_w, decoder_h,
        boxstyle="round,pad=0.05",
        facecolor="#e3f2fd",
        edgecolor="#1565c0",
        linewidth=2.5,
        zorder=1,
    )
    ax.add_patch(decoder_bg)
    ax.text(decoder_x + decoder_w / 2, decoder_y + decoder_h - 0.35,
            "LDPC Decoder", ha="center", va="center",
            fontsize=15, fontweight="bold", color="#0d47a1", zorder=2)
    ax.text(decoder_x + decoder_w / 2, decoder_y + decoder_h - 0.75,
            "Rate 1/8 QC-LDPC  |  n=256, k=32  |  Offset Min-Sum",
            ha="center", va="center", fontsize=9, color="#1565c0", zorder=2)

    # Sub-blocks inside decoder
    sb_color = "#bbdefb"
    sb_edge = "#1565c0"
    sb_fs = 9.5

    # Column 1 (left side of decoder)
    col1_x = decoder_x + 0.3
    rounded_box(ax, (col1_x, 5.8), 2.4, 0.7, "Wishbone Interface",
                facecolor=sb_color, edgecolor=sb_edge, fontsize=sb_fs, linewidth=1.2)
    rounded_box(ax, (col1_x, 4.6), 2.4, 0.7, "LLR RAM\n(256 x 6-bit)",
                facecolor=sb_color, edgecolor=sb_edge, fontsize=sb_fs, linewidth=1.2)
    rounded_box(ax, (col1_x, 3.1), 2.4, 1.0, "VN Update Array\n(Z=32 parallel)",
                facecolor="#c8e6c9", edgecolor="#2e7d32", fontsize=sb_fs, linewidth=1.2)
    rounded_box(ax, (col1_x, 1.6), 2.4, 1.0, "CN Update Array\n(Z=32 parallel)",
                facecolor="#ffccbc", edgecolor="#bf360c", fontsize=sb_fs, linewidth=1.2)

    # Column 2 (right side of decoder)
    col2_x = decoder_x + 3.3
    rounded_box(ax, (col2_x, 5.8), 2.4, 0.7, "Barrel Shifter\n(Z=32)",
                facecolor=sb_color, edgecolor=sb_edge, fontsize=sb_fs, linewidth=1.2)
    rounded_box(ax, (col2_x, 4.6), 2.4, 0.7, "Message RAM\n(1792 x 6-bit)",
                facecolor=sb_color, edgecolor=sb_edge, fontsize=sb_fs, linewidth=1.2)
    rounded_box(ax, (col2_x, 3.1), 2.4, 1.0, "Iteration\nController",
                facecolor="#fff9c4", edgecolor="#f9a825", fontsize=sb_fs, linewidth=1.2)
    rounded_box(ax, (col2_x, 1.6), 2.4, 1.0, "Syndrome\nChecker",
                facecolor="#f3e5f5", edgecolor="#7b1fa2", fontsize=sb_fs, linewidth=1.2)

    # Decoded Data output
    rounded_box(ax, (14.0, 5.0), 1.8, 1.6, "Decoded\nData",
                facecolor="#fce4ec", edgecolor="#c62828", fontsize=13,
                fontweight="bold", linewidth=2)
    ax.text(14.9, 4.4, "32 info bits\nper codeword",
            ha="center", va="top", fontsize=9, color="#c62828", style="italic")

    # ---- Arrows between top-level blocks ----
    # Photon Detector -> LLR Computer
    arrow(ax, (2.3, 5.8), (3.5, 5.8), color="#333333")
    ax.text(2.9, 6.15, "photon\ncounts", ha="center", va="bottom",
            fontsize=8, color="#555555")

    # LLR Computer -> LDPC Decoder (Wishbone Interface)
    arrow(ax, (6.1, 5.8), (7.0, 6.15), color="#333333")
    ax.text(6.55, 6.3, "6-bit\nLLR", ha="center", va="bottom",
            fontsize=8, color="#555555")

    # LDPC Decoder -> Decoded Data
    arrow(ax, (13.0, 5.8), (14.0, 5.8), color="#333333")
    ax.text(13.5, 6.15, "hard\nbits", ha="center", va="bottom",
            fontsize=8, color="#555555")

    # ---- Internal arrows inside decoder ----
    # Wishbone -> LLR RAM
    arrow(ax, (col1_x + 1.2, 5.8), (col1_x + 1.2, 5.3),
          color="#1565c0", lw=1.2, style="->")

    # LLR RAM -> VN Update
    arrow(ax, (col1_x + 1.2, 4.6), (col1_x + 1.2, 4.1),
          color="#1565c0", lw=1.2, style="->")

    # VN Update -> CN Update
    arrow(ax, (col1_x + 1.2, 3.1), (col1_x + 1.2, 2.6),
          color="#2e7d32", lw=1.2, style="->")

    # CN Update -> VN Update (feedback, curved)
    arrow(ax, (col1_x + 2.0, 2.6), (col1_x + 2.0, 3.1),
          color="#bf360c", lw=1.2, style="->",
          connectionstyle="arc3,rad=-0.4")

    # Barrel Shifter <-> VN/CN (horizontal connections)
    arrow(ax, (col1_x + 2.4, 3.6), (col2_x, 6.15),
          color="#666666", lw=1.0, style="->",
          connectionstyle="arc3,rad=-0.2")
    arrow(ax, (col2_x, 5.8), (col1_x + 2.4, 2.1),
          color="#666666", lw=1.0, style="->",
          connectionstyle="arc3,rad=-0.2")

    # Message RAM <-> VN/CN
    arrow(ax, (col2_x + 1.2, 4.6), (col2_x + 1.2, 4.1),
          color="#1565c0", lw=1.0, style="<->")

    # Iteration Controller -> VN/CN
    arrow(ax, (col2_x, 3.6), (col1_x + 2.4, 3.4),
          color="#f9a825", lw=1.0, style="->")
    arrow(ax, (col2_x, 3.3), (col1_x + 2.4, 2.0),
          color="#f9a825", lw=1.0, style="->",
          connectionstyle="arc3,rad=0.15")

    # Syndrome Checker -> Iteration Controller
    arrow(ax, (col2_x + 1.2, 2.6), (col2_x + 1.2, 3.1),
          color="#7b1fa2", lw=1.0, style="->")

    # ---- Performance annotations ----
    perf_y = 0.0
    perf_items = [
        "Target: 150 MHz (Sky130)",
        "Throughput: 7.6 Mbps",
        "Latency: ~4.2 \u00b5s/codeword",
        "Area: ~1.5 mm\u00b2",
        "Max 30 iterations (early termination)",
    ]
    perf_text = "   |   ".join(perf_items)
    ax.text(7.5, perf_y, perf_text,
            ha="center", va="center", fontsize=8.5, color="#444444",
            bbox=dict(boxstyle="round,pad=0.4", facecolor="#f5f5f5",
                      edgecolor="#cccccc", linewidth=1))

    fig.savefig(save_path, dpi=150, facecolor="white")
    plt.close(fig)
    print(f"Saved: {save_path}")


# ---------------------------------------------------------------------------
# Diagram 2: Channel Model
# ---------------------------------------------------------------------------

def generate_channel_model(save_path):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6.5))
    fig.patch.set_facecolor("white")
    fig.suptitle("Poisson Photon-Counting Channel Model",
                 fontsize=17, fontweight="bold", y=0.98, color="#1a1a2e")

    lambda_s = 3.0
    lambda_b = 0.1

    # ---- Left panel: Poisson distributions ----
    y_max = 12
    y_vals = np.arange(0, y_max + 1)

    # P(y | bit=0): Poisson(lambda_b)
    pmf_0 = poisson.pmf(y_vals, lambda_b)
    # P(y | bit=1): Poisson(lambda_s + lambda_b)
    pmf_1 = poisson.pmf(y_vals, lambda_s + lambda_b)

    bar_width = 0.35
    bars0 = ax1.bar(y_vals - bar_width / 2, pmf_0, bar_width,
                    color="#ef5350", alpha=0.85, edgecolor="#b71c1c",
                    linewidth=0.8, label=r"$P(y\,|\,\mathrm{bit}=0)$: Poisson($\lambda_b$)",
                    zorder=3)
    bars1 = ax1.bar(y_vals + bar_width / 2, pmf_1, bar_width,
                    color="#42a5f5", alpha=0.85, edgecolor="#0d47a1",
                    linewidth=0.8, label=r"$P(y\,|\,\mathrm{bit}=1)$: Poisson($\lambda_s + \lambda_b$)",
                    zorder=3)

    ax1.set_xlabel("Photon count $y$", fontsize=13)
    ax1.set_ylabel("Probability", fontsize=13)
    ax1.set_title("Channel Output Distributions", fontsize=14, fontweight="bold",
                  pad=10)
    ax1.legend(fontsize=10, loc="upper right", framealpha=0.9)
    ax1.set_xticks(y_vals)
    ax1.set_xlim(-0.8, y_max + 0.8)
    ax1.grid(axis="y", alpha=0.3, zorder=0)

    # Annotate the parameters
    param_text = (
        f"$\\lambda_s = {lambda_s:.1f}$ signal photons/slot\n"
        f"$\\lambda_b = {lambda_b:.1f}$ background photons/slot"
    )
    ax1.text(0.97, 0.70, param_text, transform=ax1.transAxes,
             ha="right", va="top", fontsize=10.5,
             bbox=dict(boxstyle="round,pad=0.4", facecolor="#fffde7",
                       edgecolor="#f9a825", alpha=0.95))

    # Channel model text annotations
    ax1.annotate("bit = 0\n(background only)",
                 xy=(0, pmf_0[0]), xytext=(2.5, pmf_0[0] + 0.15),
                 fontsize=9, color="#b71c1c", fontweight="bold",
                 arrowprops=dict(arrowstyle="->", color="#b71c1c", lw=1.2),
                 ha="center")
    ax1.annotate("bit = 1\n(signal + background)",
                 xy=(3, pmf_1[3]), xytext=(6, pmf_1[3] + 0.08),
                 fontsize=9, color="#0d47a1", fontweight="bold",
                 arrowprops=dict(arrowstyle="->", color="#0d47a1", lw=1.2),
                 ha="center")

    # ---- Right panel: LLR vs photon count ----
    # LLR(y) = log(P(y|1) / P(y|0)) = log(Poisson(y; ls+lb) / Poisson(y; lb))
    #        = -(ls+lb) + lb + y * log((ls+lb)/lb)
    #        = -ls + y * log((ls+lb)/lb)
    # Note: using natural log, then scale doesn't matter for shape

    y_cont = np.linspace(0, y_max, 500)
    log_ratio = np.log((lambda_s + lambda_b) / lambda_b)
    llr_cont = -lambda_s + y_cont * log_ratio

    # Compute LLR at integer counts for scatter plot
    llr_int = -lambda_s + y_vals * log_ratio

    # Quantization: clip to [-32, 31] (6-bit signed)
    q_min, q_max = -32, 31
    llr_quantized = np.clip(np.round(llr_int), q_min, q_max)

    ax2.plot(y_cont, llr_cont, color="#1565c0", linewidth=2.5,
             label="Exact LLR", zorder=4)
    ax2.scatter(y_vals, llr_int, color="#1565c0", s=60, zorder=5,
                edgecolors="white", linewidth=1.0)

    # Show quantized values as step markers
    for i, (yv, lq) in enumerate(zip(y_vals, llr_quantized)):
        ax2.plot([yv - 0.3, yv + 0.3], [lq, lq], color="#e65100",
                 linewidth=2.0, zorder=4, alpha=0.8)
        if i == 0:
            ax2.plot([], [], color="#e65100", linewidth=2.0,
                     label="Quantized (6-bit)")

    # Decision boundary at LLR=0
    ax2.axhline(y=0, color="#c62828", linewidth=1.5, linestyle="--",
                alpha=0.7, zorder=3, label="Decision boundary (LLR=0)")

    # Quantization range shading
    ax2.axhspan(q_min, q_max, alpha=0.06, color="#1565c0", zorder=1)
    ax2.axhline(y=q_min, color="#888888", linewidth=0.8, linestyle=":",
                alpha=0.5, zorder=2)
    ax2.axhline(y=q_max, color="#888888", linewidth=0.8, linestyle=":",
                alpha=0.5, zorder=2)
    ax2.text(y_max + 0.3, q_max, f"+{q_max}", fontsize=8, color="#888888",
             va="center")
    ax2.text(y_max + 0.3, q_min, f"{q_min}", fontsize=8, color="#888888",
             va="center")

    # Annotate quantization region
    ax2.annotate("6-bit range\n[$-32$, $+31$]",
                 xy=(y_max - 0.5, q_max), xytext=(y_max - 2.5, q_max + 8),
                 fontsize=9, color="#555555",
                 arrowprops=dict(arrowstyle="->", color="#888888", lw=1.0),
                 ha="center",
                 bbox=dict(boxstyle="round,pad=0.3", facecolor="#f5f5f5",
                           edgecolor="#cccccc"))

    # Mark where LLR crosses zero (decision boundary)
    y_cross = lambda_s / log_ratio
    ax2.axvline(x=y_cross, color="#c62828", linewidth=0.8, linestyle=":",
                alpha=0.5, zorder=2)
    ax2.text(y_cross, ax2.get_ylim()[0] if ax2.get_ylim()[0] > -40 else -38,
             f"$y^* = {y_cross:.2f}$",
             ha="center", va="bottom", fontsize=9, color="#c62828")

    # LLR formula
    formula = (
        r"$\mathrm{LLR}(y) = -\lambda_s + y \cdot "
        r"\ln\!\left(\frac{\lambda_s + \lambda_b}{\lambda_b}\right)$"
    )
    ax2.text(0.03, 0.97, formula, transform=ax2.transAxes,
             ha="left", va="top", fontsize=11,
             bbox=dict(boxstyle="round,pad=0.4", facecolor="#e8f5e9",
                       edgecolor="#388e3c", alpha=0.95))

    ax2.set_xlabel("Photon count $y$", fontsize=13)
    ax2.set_ylabel("Log-Likelihood Ratio (LLR)", fontsize=13)
    ax2.set_title("LLR Computation & Quantization", fontsize=14,
                  fontweight="bold", pad=10)
    ax2.legend(fontsize=10, loc="lower right", framealpha=0.9)
    ax2.set_xticks(y_vals)
    ax2.set_xlim(-0.5, y_max + 0.8)
    ax2.grid(alpha=0.3, zorder=0)

    fig.tight_layout(rect=[0, 0, 1, 0.94])
    fig.savefig(save_path, dpi=150, facecolor="white")
    plt.close(fig)
    print(f"Saved: {save_path}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    print("Generating LDPC optical decoder report diagrams...")
    print(f"Output directory: {PLOT_DIR}")
    print()

    arch_path = os.path.join(PLOT_DIR, "system_architecture.png")
    generate_system_architecture(arch_path)

    channel_path = os.path.join(PLOT_DIR, "channel_model.png")
    generate_channel_model(channel_path)

    print()
    print("Done. All diagrams generated successfully.")


if __name__ == "__main__":
    main()
