#!/usr/bin/env python3
"""
analyze-question-tracker.py

Socratic Reply Lab - Offline Sequential Analytics & Visualization

This script parses the exported JSON from your Socratic Question Tracker,
calculates sequential metrics (Survival, Verdict, and Mutation rates,
along with Wilson 95% confidence intervals on Survival and a cumulative score),
prints a scientific ASCII report, and generates publication-quality plots.

Usage:
  python3 analyze-question-tracker.py export.json --out plot.png --csv table.csv
"""

import os
import sys
import json
import math
import argparse
import numpy as np

# Use non-interactive backend for headless environments
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def calculate_wilson_interval(p, n, z=1.96):
    """
    Calculates the Wilson score interval for a binomial proportion.
    """
    if n == 0:
        return 0.0, 1.0
    denominator = 1 + (z**2) / n
    center = (p + (z**2) / (2 * n)) / denominator
    half_width = z * math.sqrt((p * (1 - p) / n) + (z**2) / (4 * (n**2))) / denominator
    lower = max(0.0, center - half_width)
    upper = min(1.0, center + half_width)
    return lower, upper


def parse_args():
    parser = argparse.ArgumentParser(description="Analyze Socratic Question Tracker data.")
    parser.add_argument("json_file", help="Path to the exported JSON file from the tracker.")
    parser.add_argument("--out", default="socratic_sequential_plot.png", help="Path to save the output plot image.")
    parser.add_argument("--csv", default="socratic_sequential_stats.csv", help="Path to save the summary CSV table.")
    return parser.parse_args()


def normalize_entry(entry):
    """Flatten tracker field names so the script tolerates both old and new exports."""
    out = dict(entry)
    # outcome / code
    if "outcome" in out and "code" not in out:
        out["code"] = out["outcome"]
    # date / posted_at
    if "posted_at" not in out and "date" in out:
        out["posted_at"] = out["date"]
    # question / question_as_posted
    if "question_as_posted" not in out and "question" in out:
        out["question_as_posted"] = out["question"]
    return out


def main():
    args = parse_args()

    if not os.path.exists(args.json_file):
        print(f"Error: File '{args.json_file}' not found.", file=sys.stderr)
        sys.exit(1)

    try:
        with open(args.json_file, "r") as f:
            data = json.load(f)
    except Exception as e:
        print(f"Error reading JSON file: {e}", file=sys.stderr)
        sys.exit(1)

    # Ensure data is a list
    if not isinstance(data, list):
        # Some platforms might wrap the export in an object, try to extract list
        if isinstance(data, dict) and "entries" in data:
            data = data["entries"]
        else:
            print("Error: JSON structure must be a list of entries.", file=sys.stderr)
            sys.exit(1)

    data = [normalize_entry(e) for e in data]

    # Sort by date if available, otherwise preserve order
    try:
        data.sort(key=lambda e: e.get("posted_at") or e.get("date") or "")
    except Exception:
        pass

    print(f"Loaded {len(data)} logged questions.")

    # Sequential calculation over ENGAGED entries
    engaged_entries = []
    for entry in data:
        outcome = entry.get("outcome", entry.get("code", "O")).strip().upper()
        # O means Open, which are censored and excluded from sequential analysis
        if outcome != "O":
            engaged_entries.append(entry)

    n_engaged = len(engaged_entries)
    print(f"Engaged entries (excluding Open): {n_engaged}")

    # Initialize trackers
    counts = {"A": 0, "P": 0, "M": 0, "D": 0, "R": 0, "X": 0}

    # Lists to store metrics at each step n
    steps_n = []
    survival_rates = []
    verdict_rates = []
    mutation_rates = []
    deltas = []
    wilson_lowers = []
    wilson_uppers = []
    cumulative_scores = []

    running_score = 0

    # Process sequentially
    for idx, entry in enumerate(engaged_entries, 1):
        outcome = entry.get("outcome", entry.get("code", "")).strip().upper()
        if outcome in counts:
            counts[outcome] += 1
        else:
            # Fallback
            counts["D"] += 1

        # Cumulative score calculation: score = 3*replies + 1*likes + 5*follows
        replies = int(entry.get("replies", 0) or 0)
        likes = int(entry.get("likes", 0) or 0)
        follows = int(entry.get("follows", 0) or 0)
        entry_score = (3 * replies) + (1 * likes) + (5 * follows)
        running_score += entry_score

        # Calculate rates at step idx
        # S = (P + M + R) / Neng
        # V = (A + D) / Neng
        # μ = M / Neng
        p_count = counts["P"] + counts["M"] + counts["R"]
        v_count = counts["A"] + counts["D"]
        m_count = counts["M"]

        s_rate = p_count / idx
        v_rate = v_count / idx
        m_rate = m_count / idx
        delta = s_rate - v_rate

        # Wilson 95% on S
        w_lower, w_upper = calculate_wilson_interval(s_rate, idx)

        # Append to lists
        steps_n.append(idx)
        survival_rates.append(s_rate)
        verdict_rates.append(v_rate)
        mutation_rates.append(m_rate)
        deltas.append(delta)
        wilson_lowers.append(w_lower)
        wilson_uppers.append(w_upper)
        cumulative_scores.append(running_score)

    # Write summary CSV
    try:
        import csv
        with open(args.csv, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow([
                "Engaged_n", "Outcome", "Survival_S", "Verdict_V", "Delta_S_V",
                "Mutation_mu", "Wilson_S_Lower", "Wilson_S_Upper", "Cumulative_Score",
                "Question"
            ])
            for i in range(n_engaged):
                entry = engaged_entries[i]
                writer.writerow([
                    steps_n[i],
                    entry.get("outcome", entry.get("code", "")),
                    f"{survival_rates[i]:.4f}",
                    f"{verdict_rates[i]:.4f}",
                    f"{deltas[i]:.4f}",
                    f"{mutation_rates[i]:.4f}",
                    f"{wilson_lowers[i]:.4f}",
                    f"{wilson_uppers[i]:.4f}",
                    cumulative_scores[i],
                    entry.get("question_as_posted", entry.get("question", ""))[:50] + "..."
                ])
        print(f"Saved sequential table to CSV: {args.csv}")
    except Exception as e:
        print(f"Warning: Failed to write CSV file: {e}", file=sys.stderr)

    # Output elegant ASCII report
    print("\n" + "="*70)
    print("                 SOCRATIC REPLY LAB FIELD ANALYSIS REPORT")
    print("="*70)
    print(f"Total Engaged Trials (n): {n_engaged}")
    if n_engaged > 0:
        final_s = survival_rates[-1]
        final_v = verdict_rates[-1]
        final_mu = mutation_rates[-1]
        final_delta = deltas[-1]
        w_l, w_u = wilson_lowers[-1], wilson_uppers[-1]

        print(f"Current Survival Rate (S):  {final_s:.1%} (95% CI: {w_l:.1%} to {w_u:.1%})")
        print(f"Current Verdict Rate (V):   {final_v:.1%}")
        print(f"Current Mutation Rate (μ):  {final_mu:.1%}")
        print(f"Drift Margin (Δ = S - V):   {final_delta:+.2f}")
        print(f"Total Carried Forward Score: {running_score}")
        print("-"*70)
        print("Outcome Distribution:")
        for code, count in counts.items():
            pct = (count / n_engaged) if n_engaged > 0 else 0
            print(f"  {code}: {count:2d} ({pct:.1%})")
        print("-"*70)

        # Falsifier check at n = 20
        print("Hypothesis Evaluation:")
        if n_engaged < 20:
            print(f"  STATUS: MONITORING ({n_engaged}/20 engaged trials).")
            if final_s > final_v:
                print("  TREND: Positive (S > V). Socratic questions are maintaining momentum.")
            else:
                print("  TREND: Critical (S <= V). The environment is resisting inquiry; adapt questions.")
        else:
            if final_s > final_v:
                print("  STATUS: ✅ WORKING HYPOTHESIS SUPPORTED.")
                print("  Socratic questions are scientifically shown to travel further than counter-assertions.")
            else:
                print("  STATUS: ❌ HYPOTHESIS FAILS AS STATED.")
                print("  The verdict rate has caught or exceeded the survival rate. Question design needs revision.")
    else:
        print("No engaged entries to analyze yet. Keep logging on X and Substack!")
    print("="*70 + "\n")

    # Generate Publication-Quality Plots
    if n_engaged > 0:
        available_styles = plt.style.available
        if 'seaborn-v0_8-whitegrid' in available_styles:
            plt.style.use('seaborn-v0_8-whitegrid')
        elif 'seaborn-whitegrid' in available_styles:
            plt.style.use('seaborn-whitegrid')
        else:
            plt.style.use('default')

        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(10, 12), sharex=True)

        # Deep palette colors matching the Socratic theme
        c_navy = '#1A365D'
        c_teal = '#0D9488'
        c_coral = '#E11D48'
        c_gray = '#64748B'

        # Panel 1: Survival vs. Verdict Rates
        ax1.plot(steps_n, survival_rates, label='Survival Rate (S)', color=c_teal, linewidth=2.5, marker='o', markersize=4)
        ax1.plot(steps_n, verdict_rates, label='Verdict Rate (V)', color=c_coral, linewidth=2.0, linestyle='--', marker='x', markersize=4)
        ax1.set_ylabel('Rate (0.0 to 1.0)', fontsize=11, fontweight='bold', color=c_navy)
        ax1.set_title('Socratic Survival vs. Verdict Rates Over Sequential Trials', fontsize=12, fontweight='bold', color=c_navy, pad=10)
        ax1.set_ylim(-0.05, 1.05)
        ax1.legend(loc='upper right', frameon=True, facecolor='white', edgecolor='#E2E8F0')
        ax1.tick_params(colors=c_navy)

        # Panel 2: Drift Margin (S - V) with Wilson Confidence Interval of S
        ax2.axhline(0, color=c_gray, linestyle=':', linewidth=1.2, label='Neutral Drift (S = V)')
        ax2.plot(steps_n, deltas, label='Drift Margin (Δ = S - V)', color=c_navy, linewidth=2.0, marker='s', markersize=4)

        # Show Wilson Confidence Interval of Survival Rate centered around Survival - Verdict
        # For visualization, we plot the lower and upper bounds of S minus V (shifted by the same margin)
        ci_lower = [l - verdict_rates[i] for i, l in enumerate(wilson_lowers)]
        ci_upper = [u - verdict_rates[i] for i, u in enumerate(wilson_uppers)]
        ax2.fill_between(steps_n, ci_lower, ci_upper, color=c_teal, alpha=0.15, label='S-Rate Wilson 95% Confidence Band')

        ax2.set_ylabel('Drift Margin (Δ)', fontsize=11, fontweight='bold', color=c_navy)
        ax2.set_title('Sequential Drift Margin & Statistical Confidence', fontsize=12, fontweight='bold', color=c_navy, pad=10)
        ax2.set_ylim(-1.05, 1.05)
        ax2.legend(loc='lower left', frameon=True, facecolor='white', edgecolor='#E2E8F0')
        ax2.tick_params(colors=c_navy)

        # Panel 3: Cumulative Score (3R + L + 5F)
        ax3.fill_between(steps_n, cumulative_scores, color=c_teal, alpha=0.08)
        ax3.plot(steps_n, cumulative_scores, label='Carried Forward Score (3R + L + 5F)', color=c_teal, linewidth=2.5, marker='^', markersize=4)
        ax3.set_xlabel('Sequential Engaged Trial Number (n)', fontsize=11, fontweight='bold', color=c_navy)
        ax3.set_ylabel('Cumulative Score', fontsize=11, fontweight='bold', color=c_navy)
        ax3.set_title('Socratic Kinetic Energy Score over Sequential Trials', fontsize=12, fontweight='bold', color=c_teal, pad=10)
        ax3.legend(loc='upper left', frameon=True, facecolor='white', edgecolor='#E2E8F0')
        ax3.tick_params(colors=c_navy)

        # Format axes
        for ax in (ax1, ax2, ax3):
            ax.set_xlim(0.8, n_engaged + 0.2)
            # Standardizing grid lines
            ax.grid(True, linestyle=':', alpha=0.6, color='#CBD5E1')

        plt.tight_layout()
        plt.savefig(args.out, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Generated sequential analytical plot: {args.out}")
    else:
        print("No engaged entries to generate plots.")


if __name__ == "__main__":
    main()