#!/usr/bin/env python3
"""Generate presentation plots from density evolution results."""

import json
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from pathlib import Path

RESULTS_PATH = Path(__file__).parent.parent / 'data' / 'de_results.json'
RESULTS_Z128_PATH = Path(__file__).parent.parent / 'data' / 'de_results_z128.json'
SC_RESULTS_PATH = Path(__file__).parent.parent / 'data' / 'sc_ldpc_results.json'
OUT_DIR = Path(__file__).parent.parent / 'data' / 'plots'

# Presentation-friendly style
plt.rcParams.update({
    'font.size': 13,
    'axes.titlesize': 15,
    'axes.labelsize': 14,
    'legend.fontsize': 11,
    'xtick.labelsize': 12,
    'ytick.labelsize': 12,
    'figure.dpi': 150,
    'savefig.bbox': 'tight',
    'savefig.pad_inches': 0.15,
})

COLORS = {
    'original': '#d62728',   # red
    'peg_ring': '#1f77b4',   # blue
    'optimized': '#2ca02c',  # green
}

LABELS = {
    'original': 'Original staircase [7,2,2,2,2,2,2,1]',
    'peg_ring': 'PEG ring [7,3,3,3,2,2,2,2]',
    'optimized': 'DE-optimized [7,4,4,4,4,3,3,3]',
}

MARKERS = {
    'original': 's',
    'peg_ring': '^',
    'optimized': 'o',
}


def load_results():
    with open(RESULTS_PATH) as f:
        return json.load(f)


def plot_fer_comparison(data):
    """FER vs lambda_s for all three matrices."""
    fig, ax = plt.subplots(figsize=(8, 5.5))

    fer_data = data['fer_comparison']
    lam_s = fer_data['lam_s_points']

    for key in ['original', 'peg_ring', 'optimized']:
        fer_vals = [fer_data[key][str(l)]['fer'] for l in lam_s]
        # Replace 0 with small value for log scale
        fer_plot = [max(f, 2e-3) for f in fer_vals]
        ax.semilogy(lam_s, fer_plot,
                     color=COLORS[key], marker=MARKERS[key], markersize=8,
                     linewidth=2, label=LABELS[key],
                     markeredgecolor='white', markeredgewidth=0.8)

    ax.set_xlabel(r'Signal photons/slot ($\lambda_s$)')
    ax.set_ylabel('Frame Error Rate (FER)')
    ax.set_title('FER Comparison: Rate-1/8 QC-LDPC (n=256, k=32)')
    ax.legend(loc='upper right', framealpha=0.9)
    ax.set_xlim(1.5, 10.5)
    ax.set_ylim(1e-3, 1.0)
    ax.grid(True, alpha=0.3, which='both')
    ax.set_xticks([2, 3, 4, 5, 7, 10])
    ax.set_xticklabels(['2', '3', '4', '5', '7', '10'])

    # Annotate the improvement
    ax.annotate('13x fewer\nframe errors',
                xy=(5, 0.01), xytext=(6.5, 0.06),
                fontsize=11, color=COLORS['optimized'],
                arrowprops=dict(arrowstyle='->', color=COLORS['optimized'], lw=1.5),
                ha='center')

    fig.savefig(OUT_DIR / 'fer_comparison.png')
    plt.close(fig)
    print(f'  Saved fer_comparison.png')


def plot_threshold_bars(data):
    """Bar chart of DE thresholds."""
    fig, ax = plt.subplots(figsize=(7, 5))

    thresholds = {
        'Original\nstaircase': data['reference_thresholds']['Original staircase [7,2,2,2,2,2,2,1]'],
        'PEG\nring': data['reference_thresholds']['PEG ring [7,3,3,3,2,2,2,2]'],
        'DE-\noptimized': data['best_threshold'],
    }
    # Shannon limit for rate 1/8 Poisson channel ~ 0.47 photons/slot
    shannon_limit = 0.47

    names = list(thresholds.keys())
    vals = list(thresholds.values())
    colors = [COLORS['original'], COLORS['peg_ring'], COLORS['optimized']]

    bars = ax.bar(names, vals, color=colors, width=0.55, edgecolor='white', linewidth=1.5)

    # Value labels on bars
    for bar, val in zip(bars, vals):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.1,
                f'{val:.1f}', ha='center', va='bottom', fontweight='bold', fontsize=14)

    # Shannon limit line
    ax.axhline(y=shannon_limit, color='black', linestyle='--', linewidth=1.5, alpha=0.6)
    ax.text(2.42, shannon_limit + 0.12, f'Shannon limit ({shannon_limit})',
            fontsize=11, ha='right', va='bottom', style='italic', alpha=0.7)

    ax.set_ylabel(r'DE Threshold ($\lambda_s^*$, photons/slot)')
    ax.set_title('Decoding Threshold Comparison')
    ax.set_ylim(0, 6.5)
    ax.grid(True, alpha=0.2, axis='y')

    # Improvement annotation
    improvement_db = 10 * np.log10(thresholds['Original\nstaircase'] / thresholds['DE-\noptimized'])
    ax.annotate(f'{improvement_db:.1f} dB\nimprovement',
                xy=(2, thresholds['DE-\noptimized']),
                xytext=(2.35, 4.2),
                fontsize=12, fontweight='bold', color=COLORS['optimized'],
                arrowprops=dict(arrowstyle='->', color=COLORS['optimized'], lw=1.5),
                ha='center')

    fig.savefig(OUT_DIR / 'threshold_comparison.png')
    plt.close(fig)
    print(f'  Saved threshold_comparison.png')


def plot_degree_distribution(data):
    """VN degree distribution comparison."""
    fig, axes = plt.subplots(1, 3, figsize=(12, 4), sharey=True)

    distributions = {
        'Original staircase': [7, 2, 2, 2, 2, 2, 2, 1],
        'PEG ring': [7, 3, 3, 3, 2, 2, 2, 2],
        'DE-optimized': data['best_degrees'],
    }
    color_keys = ['original', 'peg_ring', 'optimized']
    col_labels = [f'c{i}' for i in range(8)]

    for ax, (name, degrees), ckey in zip(axes, distributions.items(), color_keys):
        colors_bar = [COLORS[ckey] if i > 0 else '#888888' for i in range(8)]
        ax.bar(col_labels, degrees, color=colors_bar, edgecolor='white', linewidth=1)
        ax.set_title(name, fontsize=12, fontweight='bold')
        ax.set_xlabel('Base matrix column')
        ax.set_ylim(0, 8)
        ax.grid(True, alpha=0.2, axis='y')

        # Annotate average parity degree
        avg_parity = np.mean(degrees[1:])
        ax.text(4, 7.2, f'avg parity dv = {avg_parity:.1f}',
                ha='center', fontsize=10, style='italic')

    axes[0].set_ylabel('Variable node degree')
    fig.suptitle('VN Degree Distributions (col 0 = info, cols 1-7 = parity)',
                 fontsize=13, y=1.02)
    fig.tight_layout()

    fig.savefig(OUT_DIR / 'degree_distributions.png')
    plt.close(fig)
    print(f'  Saved degree_distributions.png')


def plot_base_matrix_heatmap(data):
    """Heatmap of the constructed base matrix."""
    fig, ax = plt.subplots(figsize=(7, 5))

    H = np.array(data['constructed_matrix'])
    m_base, n_base = H.shape

    # Create display matrix: shift values where connected, NaN where not
    display = np.full_like(H, dtype=float, fill_value=np.nan)
    for r in range(m_base):
        for c in range(n_base):
            if H[r, c] >= 0:
                display[r, c] = H[r, c]

    im = ax.imshow(display, cmap='YlGnBu', aspect='auto', vmin=0, vmax=31)

    # Add text annotations
    for r in range(m_base):
        for c in range(n_base):
            if H[r, c] >= 0:
                ax.text(c, r, str(H[r, c]), ha='center', va='center',
                        fontsize=11, fontweight='bold', color='black')
            else:
                ax.text(c, r, '-', ha='center', va='center',
                        fontsize=11, color='#cccccc')

    ax.set_xticks(range(n_base))
    ax.set_xticklabels([f'c{i}' for i in range(n_base)])
    ax.set_yticks(range(m_base))
    ax.set_yticklabels([f'r{i}' for i in range(m_base)])
    ax.set_xlabel('Column (c0=info, c1-c7=parity)')
    ax.set_ylabel('Row (check node group)')
    ax.set_title(f'Optimized Base Matrix (Z=32, girth={data["matrix_checks"]["girth"]})')

    cbar = fig.colorbar(im, ax=ax, shrink=0.8, label='Circulant shift')

    fig.savefig(OUT_DIR / 'base_matrix_heatmap.png')
    plt.close(fig)
    print(f'  Saved base_matrix_heatmap.png')


def load_z128_results():
    """Load Z=128 results if available."""
    if RESULTS_Z128_PATH.exists():
        with open(RESULTS_Z128_PATH) as f:
            return json.load(f)
    return None


def plot_z128_comparison(data_z32, data_z128):
    """
    FER vs lambda_s comparison between Z=32 and Z=128 for the
    DE-optimized degree distribution [7,4,4,4,4,3,3,3].

    Shows both lifting factors on the same axes, with the optimized
    matrix in each case plus the PEG ring reference.
    """
    fig, ax = plt.subplots(figsize=(9, 6))

    # Colors and styles for Z=32 vs Z=128
    z32_color = '#2ca02c'    # green (matches optimized)
    z128_color = '#9467bd'   # purple
    z32_peg_color = '#1f77b4'  # blue (matches peg_ring)
    z128_peg_color = '#ff7f0e' # orange

    # Z=32 optimized data (from de_results.json)
    fer_z32 = data_z32['fer_comparison']
    lam_s_z32 = fer_z32['lam_s_points']
    fer_z32_opt = [fer_z32['optimized'][str(l)]['fer'] for l in lam_s_z32]
    fer_z32_peg = [fer_z32['peg_ring'][str(l)]['fer'] for l in lam_s_z32]

    # Z=128 data
    fer_z128 = data_z128['fer_results']
    lam_s_z128 = fer_z128['lam_s_points']
    fer_z128_opt = [fer_z128['optimized'][str(l)]['fer'] for l in lam_s_z128]
    fer_z128_peg = [fer_z128['peg_ring'][str(l)]['fer'] for l in lam_s_z128]

    # Replace 0 with small value for log scale
    floor = 5e-3

    # Plot Z=32 curves (solid lines)
    ax.semilogy(lam_s_z32, [max(f, floor) for f in fer_z32_opt],
                color=z32_color, marker='o', markersize=8, linewidth=2,
                label='Optimized Z=32 (n=256, offset MS)',
                markeredgecolor='white', markeredgewidth=0.8)
    ax.semilogy(lam_s_z32, [max(f, floor) for f in fer_z32_peg],
                color=z32_peg_color, marker='^', markersize=8, linewidth=2,
                label='PEG ring Z=32 (n=256, offset MS)',
                markeredgecolor='white', markeredgewidth=0.8,
                linestyle='--', alpha=0.7)

    # Plot Z=128 curves (dashed lines, larger markers)
    ax.semilogy(lam_s_z128, [max(f, floor) for f in fer_z128_opt],
                color=z128_color, marker='D', markersize=9, linewidth=2.5,
                label=f'Optimized Z=128 (n=1024, norm MS, a={data_z128["alpha"]})',
                markeredgecolor='white', markeredgewidth=0.8)
    ax.semilogy(lam_s_z128, [max(f, floor) for f in fer_z128_peg],
                color=z128_peg_color, marker='v', markersize=9, linewidth=2.5,
                label=f'PEG ring Z=128 (n=1024, norm MS, a={data_z128["alpha"]})',
                markeredgecolor='white', markeredgewidth=0.8,
                linestyle='--', alpha=0.7)

    ax.set_xlabel(r'Signal photons/slot ($\lambda_s$)')
    ax.set_ylabel('Frame Error Rate (FER)')
    ax.set_title('FER Comparison: Z=32 vs Z=128\nDE-optimized [7,4,4,4,4,3,3,3]')
    ax.legend(loc='upper right', framealpha=0.9, fontsize=10)
    ax.set_xlim(1.0, 11.0)
    ax.set_ylim(floor / 2, 1.1)
    ax.grid(True, alpha=0.3, which='both')
    ax.set_xticks([1.5, 2, 3, 4, 5, 7, 10])
    ax.set_xticklabels(['1.5', '2', '3', '4', '5', '7', '10'])

    # Add Z=128 matrix info annotation
    girth_128 = data_z128.get('girth', '?')
    ax.text(0.02, 0.02,
            f'Z=128: girth={girth_128}, n=1024, k=128\n'
            f'Z=32: girth={data_z32["matrix_checks"]["girth"]}, n=256, k=32',
            transform=ax.transAxes, fontsize=9, verticalalignment='bottom',
            bbox=dict(boxstyle='round,pad=0.3', facecolor='wheat', alpha=0.5))

    fig.savefig(OUT_DIR / 'z128_fer_comparison.png')
    plt.close(fig)
    print(f'  Saved z128_fer_comparison.png')


def load_sc_results():
    """Load SC-LDPC results if available."""
    if SC_RESULTS_PATH.exists():
        with open(SC_RESULTS_PATH) as f:
            return json.load(f)
    return None


def plot_threshold_progression(data_z32, data_sc):
    """
    Threshold progression: original -> DE-optimized -> normalized -> SC-LDPC,
    with Shannon limit line. The key summary plot.
    """
    fig, ax = plt.subplots(figsize=(10, 5.5))

    shannon_limit = 0.47

    # Build progression data
    stages = []
    values = []
    colors_bar = []

    # Stage 1: Original staircase (offset)
    orig_thresh = data_z32['reference_thresholds'].get(
        'Original staircase [7,2,2,2,2,2,2,1]', 5.23)
    stages.append('Original\nstaircase\n(offset)')
    values.append(float(orig_thresh))
    colors_bar.append('#d62728')

    # Stage 2: DE-optimized (offset)
    opt_thresh = data_z32['best_threshold']
    stages.append('DE-optimized\n(offset)')
    values.append(float(opt_thresh))
    colors_bar.append('#ff7f0e')

    # Stage 3: Normalized min-sum
    if data_sc and 'uncoupled_thresholds' in data_sc:
        norm_thresh = data_sc['uncoupled_thresholds'].get('optimized_normalized', opt_thresh)
        stages.append('DE-optimized\n(normalized)')
        values.append(float(norm_thresh))
        colors_bar.append('#2ca02c')

    # Stage 4: SC-LDPC
    if data_sc and 'sc_thresholds' in data_sc:
        sc_thresh = data_sc['sc_thresholds'].get('sc_optimized',
                    data_sc['sc_thresholds'].get('sc_original', None))
        if sc_thresh is not None:
            stages.append('SC-LDPC\n(normalized)')
            values.append(float(sc_thresh))
            colors_bar.append('#9467bd')

    x = np.arange(len(stages))
    bars = ax.bar(x, values, color=colors_bar, width=0.6, edgecolor='white', linewidth=1.5)

    # Value labels on bars
    for bar, val in zip(bars, values):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.08,
                f'{val:.2f}', ha='center', va='bottom', fontweight='bold', fontsize=13)

    # Shannon limit line
    ax.axhline(y=shannon_limit, color='black', linestyle='--', linewidth=2, alpha=0.7)
    ax.text(len(stages) - 0.5, shannon_limit + 0.08,
            f'Shannon limit ({shannon_limit})',
            fontsize=11, ha='right', va='bottom', style='italic', alpha=0.7)

    # dB annotations between bars
    for i in range(1, len(values)):
        if values[i] > 0 and values[i-1] > 0:
            gain_db = 10 * np.log10(values[i-1] / values[i])
            mid_y = (values[i-1] + values[i]) / 2
            ax.annotate(f'{gain_db:.1f} dB',
                       xy=(i, values[i] + 0.02),
                       xytext=(i - 0.5, mid_y + 0.3),
                       fontsize=10, color='#555555',
                       arrowprops=dict(arrowstyle='->', color='#999999', lw=1),
                       ha='center')

    # Total gap annotation
    if len(values) >= 2:
        total_gap_db = 10 * np.log10(values[-1] / shannon_limit)
        ax.text(len(stages) / 2, max(values) * 0.85,
                f'Remaining gap to Shannon: {total_gap_db:.1f} dB',
                ha='center', fontsize=12, fontweight='bold',
                bbox=dict(boxstyle='round,pad=0.3', facecolor='lightyellow', alpha=0.8))

    ax.set_xticks(x)
    ax.set_xticklabels(stages, fontsize=11)
    ax.set_ylabel(r'DE Threshold ($\lambda_s^*$, photons/slot)')
    ax.set_title('Shannon Limit Roadmap: Threshold Progression')
    ax.set_ylim(0, max(values) * 1.2)
    ax.grid(True, alpha=0.2, axis='y')

    fig.savefig(OUT_DIR / 'threshold_progression.png')
    plt.close(fig)
    print(f'  Saved threshold_progression.png')


def plot_sc_threshold_bars(data_sc):
    """SC threshold vs uncoupled threshold bar chart."""
    if not data_sc:
        return

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

    shannon_limit = data_sc.get('shannon_limit', 0.47)

    thresholds = {}
    if 'uncoupled_thresholds' in data_sc:
        ut = data_sc['uncoupled_thresholds']
        thresholds['Original\n(offset)'] = (float(ut['original_offset']), '#d62728')
        thresholds['DE-opt\n(offset)'] = (float(ut['optimized_offset']), '#ff7f0e')
        thresholds['DE-opt\n(normalized)'] = (float(ut['optimized_normalized']), '#2ca02c')
    if 'sc_thresholds' in data_sc:
        st = data_sc['sc_thresholds']
        thresholds['SC\noriginal'] = (float(st['sc_original']), '#17becf')
        thresholds['SC\nDE-opt'] = (float(st['sc_optimized']), '#9467bd')

    names = list(thresholds.keys())
    vals = [v[0] for v in thresholds.values()]
    colors = [v[1] for v in thresholds.values()]

    bars = ax.bar(names, vals, color=colors, width=0.55, edgecolor='white', linewidth=1.5)

    for bar, val in zip(bars, vals):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.05,
                f'{val:.2f}', ha='center', va='bottom', fontweight='bold', fontsize=12)

    ax.axhline(y=shannon_limit, color='black', linestyle='--', linewidth=1.5, alpha=0.6)
    ax.text(len(names) - 0.5, shannon_limit + 0.05,
            f'Shannon ({shannon_limit})', fontsize=10, ha='right', style='italic', alpha=0.7)

    ax.set_ylabel(r'DE Threshold ($\lambda_s^*$, photons/slot)')
    ax.set_title('SC-LDPC Threshold Saturation Effect')
    ax.set_ylim(0, max(vals) * 1.25)
    ax.grid(True, alpha=0.2, axis='y')

    fig.savefig(OUT_DIR / 'sc_threshold_comparison.png')
    plt.close(fig)
    print(f'  Saved sc_threshold_comparison.png')


def plot_sc_fer_comparison(data_sc):
    """FER comparison: SC-LDPC vs uncoupled."""
    if not data_sc or 'fer_comparison' not in data_sc:
        return

    fig, ax = plt.subplots(figsize=(8, 5.5))

    fer_data = data_sc['fer_comparison']
    lam_s_points = fer_data['lam_s_points']
    sc_fer = fer_data['sc_fer']

    # SC-LDPC FER
    fer_vals = [sc_fer[str(l)]['fer'] for l in lam_s_points]
    floor = 5e-3
    ax.semilogy(lam_s_points, [max(f, floor) for f in fer_vals],
                color='#9467bd', marker='D', markersize=9, linewidth=2.5,
                label=f'SC-LDPC (L={fer_data["params"]["L"]}, windowed)',
                markeredgecolor='white', markeredgewidth=0.8)

    ax.set_xlabel(r'Signal photons/slot ($\lambda_s$)')
    ax.set_ylabel('Frame Error Rate (FER)')
    L = fer_data['params']['L']
    ax.set_title(f'SC-LDPC FER (L={L}, w=2, Z=32, normalized alpha=0.875)')
    ax.legend(loc='upper right', framealpha=0.9)
    ax.set_xlim(1.5, 10.5)
    ax.set_ylim(floor / 2, 1.1)
    ax.grid(True, alpha=0.3, which='both')

    fig.savefig(OUT_DIR / 'sc_fer.png')
    plt.close(fig)
    print(f'  Saved sc_fer.png')


def main():
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    data = load_results()

    print('Generating plots...')
    plot_fer_comparison(data)
    plot_threshold_bars(data)
    plot_degree_distribution(data)
    plot_base_matrix_heatmap(data)

    # Z=128 comparison plot
    data_z128 = load_z128_results()
    if data_z128 is not None:
        plot_z128_comparison(data, data_z128)
    else:
        print('  Skipping Z=128 comparison (data/de_results_z128.json not found)')

    # SC-LDPC plots
    data_sc = load_sc_results()
    if data_sc is not None:
        plot_sc_threshold_bars(data_sc)
        plot_sc_fer_comparison(data_sc)
        plot_threshold_progression(data, data_sc)
    else:
        print('  Skipping SC-LDPC plots (data/sc_ldpc_results.json not found)')

    print(f'\nAll plots saved to {OUT_DIR}/')


if __name__ == '__main__':
    main()