#!/usr/bin/env python3
"""prime_sidon_explore.py — Integer-only Sidon sum-degeneracy exploration.

Tests whether prime-based label sets satisfy the Sidon property (all pairwise
sums distinct) across scales from small primes (2,3,5,7,11) up to quintillion
(10^18).

All arithmetic is pure integer. No float, no numpy eigenvalue decompositions,
no SLOS, no tensor networks. The metric is the actual Sidon property:
  A set A is Sidon iff for all i,j,k,l, a_i + a_j = a_k + a_l ⇒ {i,j} = {k,l}

For each label set we compute:
  - Total pairs: n^2 (including order, so (i,j) counts separately from (j,i))
  - Distinct sums: number of unique values in the sum matrix
  - Sidon score: distinct_sums / n^2 (1.0 = perfect Sidon)
  - Max degeneracy: highest multiplicity of any single sum
  - Collision entropy: Shannon entropy of the degeneracy distribution
  - Collision count: number of sum values that appear ≥ 2 times

Each computation step is a DAG node, checkpointed to disk for resume.
The DAG structure allows adversarial review of each step.
"""

import sys
import math
import json
import time
import hashlib
import itertools
from pathlib import Path
from collections import Counter

REPO_ROOT = Path(__file__).resolve().parent.parent
ARTIFACTS_DIR = REPO_ROOT / ".openresearch" / "artifacts"
DAG_PATH = ARTIFACTS_DIR / "prime_sidon_dag.json"
OUTPUT_PATH = ARTIFACTS_DIR / "prime_sidon_explore.json"

# ── Recoverable DAG ──────────────────────────────────────────────────────────

class ComputationDAG:
    def __init__(self):
        self.nodes = []
        self.edges = []
        ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)

    def add_node(self, node_id, node_type, inputs, result, status, elapsed):
        node = {
            "id": node_id, "type": node_type,
            "inputs": inputs, "result": result,
            "status": status, "elapsed_s": round(elapsed, 4),
            "timestamp": time.time(),
        }
        self.nodes.append(node)
        return node

    def add_edge(self, from_id, to_id, edge_type="depends_on"):
        self.edges.append({"from": from_id, "to": to_id, "type": edge_type})

    def has_node(self, node_id):
        return any(n["id"] == node_id for n in self.nodes)

    def get_node(self, node_id):
        for n in self.nodes:
            if n["id"] == node_id:
                return n["result"]
        return None

    def save(self):
        dag_state = {"nodes": self.nodes, "edges": self.edges}
        DAG_PATH.write_text(json.dumps(dag_state, default=str, indent=2))

    def save_report(self):
        report_path = ARTIFACTS_DIR / "prime_sidon_dag.md"
        lines = [
            "# Prime Sidon Exploration DAG\n",
            f"**Nodes:** {len(self.nodes)}",
            f"**Edges:** {len(self.edges)}\n",
            "## Nodes\n",
            "| ID | Type | Status | Elapsed | Inputs |",
            "|----|------|--------|---------|--------|",
        ]
        for n in self.nodes:
            inp = json.dumps(n["inputs"], default=str)[:60]
            lines.append(f"| {n['id']} | {n['type']} | {n['status']} | {n['elapsed_s']}s | {inp} |")
        lines.append("\n## Edges\n| From | To | Type |\n|------|----|------|")
        for e in self.edges:
            lines.append(f"| {e['from']} | {e['to']} | {e['type']} |")
        report_path.write_text("\n".join(lines))
        return report_path

# ── Prime generation (integer only) ──────────────────────────────────────────

def primes_upto(limit):
    """Segmented sieve. Pure integer arithmetic."""
    if limit < 2:
        return []
    sieve = bytearray(b'\x01') * (limit + 1)
    sieve[0:2] = b'\x00\x00'
    for i in range(2, int(limit ** 0.5) + 1):
        if sieve[i]:
            step = i
            start = i * i
            sieve[start:limit+1:step] = b'\x00' * ((limit - start) // step + 1)
    return [i for i, is_prime in enumerate(sieve) if is_prime]

def prime_cluster_via_sieve(count, offset_start):
    """Get count consecutive primes starting at offset_start-th prime via sieve."""
    limit = max(offset_start * 2 + 100, 10_000_000)
    all_primes = primes_upto(limit)
    if offset_start + count <= len(all_primes):
        return all_primes[offset_start:offset_start + count]
    return None

def prime_cluster_via_sympy(start_val, count):
    """Get count consecutive primes starting from start_val via sympy.nextprime."""
    import sympy
    p = sympy.nextprime(start_val)
    cluster = []
    for _ in range(count):
        cluster.append(p)
        p = sympy.nextprime(p)
    return cluster

# ── Reference label sets ─────────────────────────────────────────────────────

def sidon_pow2_set(size):
    """Powers of 2: perfect Sidon (all pairwise sums distinct)."""
    return [1 << i for i in range(size)]

def nonsidon_consecutive_set(size, start=1):
    """Consecutive integers: maximally non-Sidon for a given size."""
    return list(range(start, start + size))

# ── Integer-only Sidon sum-degeneracy metrics ───────────────────────────────

def sidon_metrics(labels):
    """Compute Sidon sum-degeneracy metrics using only integer arithmetic.
    
    All sums a_i + a_j are computed as exact integers. No floating point.
    The sum matrix M[i][j] = labels[i] + labels[j] has n^2 entries.
    
    Returns:
        n: number of labels
        total_pairs: n^2
        distinct_sums: number of unique sum values
        sidon_score: distinct_sums / n^2 (1.0 = perfect Sidon)
        max_degeneracy: max multiplicity of any single sum
        collision_count: number of sums with multiplicity >= 2
        collision_entropy: Shannon entropy of degeneracy distribution
        sum_histogram: {sum_value: count} for sums with count >= 2
    """
    n = len(labels)
    total_pairs = n * n
    
    # Count occurrences of each sum (pure integer arithmetic)
    sum_counts = Counter()
    for i in range(n):
        for j in range(n):
            sum_counts[labels[i] + labels[j]] += 1
    
    distinct_sums = len(sum_counts)
    sidon_score = distinct_sums / total_pairs
    
    # Degeneracy analysis
    multiplicities = list(sum_counts.values())
    max_degeneracy = max(multiplicities) if multiplicities else 0
    collision_count = sum(1 for c in multiplicities if c >= 2)
    
    # Shannon entropy of degeneracy distribution
    probs = [c / total_pairs for c in multiplicities]
    collision_entropy = -sum(p * math.log2(p) for p in probs if p > 0)
    
    # Report only collisions (sums with count >= 2)
    collision_histogram = {str(k): v for k, v in sum_counts.items() if v >= 2}
    
    return {
        "n": n,
        "total_pairs": total_pairs,
        "distinct_sums": distinct_sums,
        "sidon_score": round(sidon_score, 6),
        "max_degeneracy": max_degeneracy,
        "collision_count": collision_count,
        "collision_entropy": round(collision_entropy, 4),
        "collision_histogram": collision_histogram,
    }

# ── Permutation test (null: random distinct integers at same scale) ─────────

def random_sidon_score(size, range_min, range_max, trials=1000, rng=None):
    """Generate `trials` random n-subsets of distinct integers in [min, max]
    and return the distribution of Sidon scores under the null hypothesis."""
    if rng is None:
        rng = __import__("random").Random(42)
    scores = []
    for _ in range(trials):
        sample = sorted(rng.sample(range(range_min, range_max + 1), size))
        scores.append(sidon_metrics(sample)["sidon_score"])
    scores.sort()
    return {
        "mean": round(sum(scores) / len(scores), 4),
        "std": round((sum((s - sum(scores)/len(scores))**2 for s in scores) / len(scores))**0.5, 4),
        "p5": round(scores[int(len(scores) * 0.05)], 4),
        "p50": round(scores[int(len(scores) * 0.50)], 4),
        "p95": round(scores[int(len(scores) * 0.95)], 4),
        "trials": trials,
    }

# ── Main exploration ────────────────────────────────────────────────────────

def explore():
    print("=" * 60)
    print("  Prime Sidon Sum-Degeneracy Exploration")
    print("  Integer-only arithmetic. Permutation test vs random null.")
    print("=" * 60)
    
    dag = ComputationDAG()
    
    # Scale definitions
    scales = {
        "small": {"offset": 0, "label": "primes 2..19"},
        "kilo": {"offset": 168, "label": "primes near 1000"},
        "million": {"offset": 78498, "label": "primes near 10^6"},
    }
    large_scales = {
        "billion": 10**9,
        "trillion": 10**12,
        "quadrillion": 10**15,
        "quintillion": 10**18,
    }
    
    sizes = [5, 6, 7, 8, 10]
    results = []
    N_TRIALS = 1000
    
    # ── Step 1: Analytical bounds ─────────────────────────────────────────
    print("\n--- Step 1: Analytical Sidon bounds ---")
    print("  For any n distinct integers:")
    print("    Lower bound (consecutive): (2n-1)/n²")
    print("    Upper bound (perfect Sidon): (n+1)/(2n)")
    print("  The 'BETWEEN' claim is a tautology — testing if primes")
    print("  deviate from random expectation within these bounds.")
    print()
    
    bound_nodes = {}
    for size in sizes:
        lower = round((2 * size - 1) / (size * size), 4)
        upper = round((size + 1) / (2 * size), 4)
        node_id = f"bounds_s{size}"
        dag.add_node(node_id, "analytical_bounds",
                     {"size": size, "lower_bound": lower, "upper_bound": upper},
                     {"lower": lower, "upper": upper}, "success", 0.0)
        bound_nodes[size] = {"lower": lower, "upper": upper}
        print(f"  size={size}: lower={lower:.4f} upper={upper:.4f}")
    
    # ── Step 2: Reference sets (baseline verification) ────────────────────
    print("\n--- Step 2: Reference baseline verification ---")
    for size in sizes:
        labels = sidon_pow2_set(size)
        metrics = sidon_metrics(labels)
        node_id = f"sidon_pow2_s{size}"
        dag.add_node(node_id, "reference",
                     {"type": "sidon_pow2", "size": size, "labels": labels},
                     metrics, "success", 0.0)
        bound = bound_nodes[size]
        at_upper = abs(metrics["sidon_score"] - bound["upper"]) < 1e-6
        print(f"  {node_id:25s} score={metrics['sidon_score']:.4f} (upper={bound['upper']:.4f}, at_bound={at_upper})")
        
        labels = nonsidon_consecutive_set(size)
        metrics = sidon_metrics(labels)
        node_id = f"nonsidon_seq_s{size}"
        dag.add_node(node_id, "reference",
                     {"type": "nonsidon_consecutive", "size": size, "labels": labels},
                     metrics, "success", 0.0)
        at_lower = abs(metrics["sidon_score"] - bound["lower"]) < 1e-6
        print(f"  {node_id:25s} score={metrics['sidon_score']:.4f} (lower={bound['lower']:.4f}, at_bound={at_lower})")
    
    # ── Step 3: Prime sets with permutation test ──────────────────────────
    print("\n--- Step 3: Prime sets vs random null ---")
    import random as _random
    rng = _random.Random(42)
    
    for scale_name, scale_info in scales.items():
        offset = scale_info["offset"]
        for size in sizes:
            cluster = prime_cluster_via_sieve(size, offset)
            if cluster is None:
                cluster = prime_cluster_via_sympy(2, offset + size)
                cluster = cluster[offset:offset + size]
            
            node_id = f"primes_{scale_name}_s{size}"
            t0 = time.time()
            metrics = sidon_metrics(cluster)
            
            # Permutation test: random n-subsets in [min(cluster), max(cluster)]
            range_min = min(cluster)
            range_max = max(cluster)
            null_dist = random_sidon_score(size, range_min, range_max, N_TRIALS, rng)
            
            # p-value: proper permutation test
            null_scores = []
            for _ in range(N_TRIALS):
                sample = sorted(rng.sample(range(range_min, range_max + 1), size))
                null_scores.append(sidon_metrics(sample)["sidon_score"])
            null_scores.sort()
            p_less = sum(1 for s in null_scores if s <= metrics["sidon_score"]) / N_TRIALS
            p_greater = sum(1 for s in null_scores if s >= metrics["sidon_score"]) / N_TRIALS
            p_two_sided = min(1.0, 2 * min(p_less, p_greater))
            
            effect = round(metrics["sidon_score"] - null_dist["mean"], 4)
            significant = p_two_sided < 0.05
            
            elapsed = time.time() - t0
            
            dag.add_node(node_id, "prime_set",
                         {"size": size, "scale": scale_name, "labels": cluster,
                          "range": [range_min, range_max], "null_trials": N_TRIALS},
                         {**metrics, "null_distribution": null_dist,
                          "p_value_two_sided": round(p_two_sided, 4),
                          "effect_size": effect, "significant": significant},
                         "success", elapsed)
            
            results.append({
                "desc": node_id, "scale": scale_name, "n_labels": size,
                "labels": cluster, "metrics": metrics,
                "null_distribution": null_dist,
                "p_value": round(p_two_sided, 4),
                "effect_size": effect,
                "significant": significant,
            })
            
            sig_str = "***" if significant else ""
            dir_str = "LESS Sidon-like" if effect < 0 else "MORE Sidon-like"
            print(f"  {node_id:30s} score={metrics['sidon_score']:.4f} "
                  f"null={null_dist['mean']:.4f}±{null_dist['std']:.4f} "
                  f"Δ={effect:+.4f} p={p_two_sided:.4f} {sig_str} ({dir_str})")
    
    # ── Step 4: Large prime scales (via sympy.nextprime) ──────────────────
    print("\n--- Step 4: Large prime scales vs random null ---")
    for scale_name, start_val in large_scales.items():
        try:
            import sympy
        except ImportError:
            print(f"  sympy not available, skipping {scale_name}")
            continue
        
        for size in sizes:
            node_id = f"primes_{scale_name}_s{size}"
            t0 = time.time()
            cluster = prime_cluster_via_sympy(start_val, size)
            metrics = sidon_metrics(cluster)
            
            range_min = min(cluster)
            range_max = max(cluster)
            null_dist = random_sidon_score(size, range_min, range_max, N_TRIALS, rng)
            
            null_scores = []
            for _ in range(N_TRIALS):
                sample = sorted(rng.sample(range(range_min, range_max + 1), size))
                null_scores.append(sidon_metrics(sample)["sidon_score"])
            null_scores.sort()
            p_less = sum(1 for s in null_scores if s <= metrics["sidon_score"]) / N_TRIALS
            p_greater = sum(1 for s in null_scores if s >= metrics["sidon_score"]) / N_TRIALS
            p_two_sided = min(1.0, 2 * min(p_less, p_greater))
            
            effect = round(metrics["sidon_score"] - null_dist["mean"], 4)
            significant = p_two_sided < 0.05
            
            elapsed = time.time() - t0
            
            dag.add_node(node_id, "prime_set",
                         {"size": size, "scale": scale_name, "labels": cluster,
                          "range": [range_min, range_max], "null_trials": N_TRIALS},
                         {**metrics, "null_distribution": null_dist,
                          "p_value_two_sided": round(p_two_sided, 4),
                          "effect_size": effect, "significant": significant},
                         "success", elapsed)
            
            results.append({
                "desc": node_id, "scale": scale_name, "n_labels": size,
                "labels": cluster, "metrics": metrics,
                "null_distribution": null_dist,
                "p_value": round(p_two_sided, 4),
                "effect_size": effect,
                "significant": significant,
            })
            
            sig_str = "***" if significant else ""
            dir_str = "LESS Sidon-like" if effect < 0 else "MORE Sidon-like"
            print(f"  {node_id:30s} score={metrics['sidon_score']:.4f} "
                  f"null={null_dist['mean']:.4f}±{null_dist['std']:.4f} "
                  f"Δ={effect:+.4f} p={p_two_sided:.4f} {sig_str} ({dir_str})")
    
    # ── Step 5: Cross-scale analysis ──────────────────────────────────────
    print("\n--- Step 5: Cross-scale analysis ---")
    
    analysis = []
    for size in sizes:
        size_results = [r for r in results if r["n_labels"] == size]
        
        # How many are significant?
        n_sig = sum(1 for r in size_results if r["significant"])
        n_total = len(size_results)
        all_sig = n_sig == n_total
        
        # Direction of effect
        neg_effects = sum(1 for r in size_results if r["effect_size"] < 0)
        pos_effects = sum(1 for r in size_results if r["effect_size"] > 0)
        
        # Score trend with scale
        scores_by_scale = {r["scale"]: r["metrics"]["sidon_score"] for r in size_results}
        ascending_scales = [s for s in ["small", "kilo", "million", "billion", "trillion", "quadrillion", "quintillion"] if s in scores_by_scale]
        score_trend = [scores_by_scale[s] for s in ascending_scales]
        increasing = all(score_trend[i] <= score_trend[i+1] for i in range(len(score_trend)-1))
        
        analysis.append({
            "size": size,
            "n_significant": n_sig,
            "n_total": n_total,
            "all_significant": all_sig,
            "negative_effects": neg_effects,
            "positive_effects": pos_effects,
            "score_trend_increasing": increasing,
            "scores_by_scale": scores_by_scale,
        })
        
        print(f"  size={size}: {n_sig}/{n_total} significant, "
              f"{neg_effects} less Sidon-like, {pos_effects} more Sidon-like, "
              f"trend={'↑' if increasing else '~'}")
    
    # ── Step 6: Summary ─────────────────────────────────────────────────
    print("\n--- Summary ---")
    total_sig = sum(1 for r in results if r["significant"])
    total_cases = len(results)
    
    print(f"  Total test cases: {total_cases}")
    print(f"  Significant at p<0.05: {total_sig}/{total_cases}")
    print()
    print(f"  Key question: are prime Sidon scores different from random")
    print(f"  n-subsets of integers at the same scale?")
    print(f"  Answer: {total_sig}/{total_cases} cases show significant deviation.")
    
    # Summary node
    summary_node = {
        "schema": "prime_sidon_explore_v2",
        "claim_boundary": "prime-sidon-vs-random-null:permutation-test",
        "config": {"sizes": sizes, "scales": list(scales.keys()) + list(large_scales.keys()),
                    "null_trials": N_TRIALS},
        "analysis": analysis,
        "total_test_cases": total_cases,
        "total_significant": total_sig,
    }
    dag.add_node("summary", "summary", summary_node, summary_node, "success", 0.0)
    for r in results:
        dag.add_edge(f"primes_{r['scale']}_s{r['n_labels']}", "summary", "feeds_summary")
    
    # ── Save ──────────────────────────────────────────────────────────────
    dag.save()
    report_path = dag.save_report()
    
    with open(OUTPUT_PATH, "w") as f:
        json.dump(summary_node, f, indent=2, default=str)
    
    print(f"\n  DAG nodes: {len(dag.nodes)}")
    print(f"  DAG edges: {len(dag.edges)}")
    print(f"  DAG: {DAG_PATH}")
    print(f"  Report: {report_path}")
    print("=" * 60)

if __name__ == "__main__":
    explore()