#!/usr/bin/env python3
"""Generate modulus heatmap visualization data (JSON)."""
import sys, math, json
sys.path.insert(0, '/home/allaun/SilverSight/scripts')
from verify_wrapping import f_k, is_sidon, certify_sidon_creation

def generate_heatmap_data(A, S, max_mod=16):
    """Generate heatmap of all coprime pairs in [2,max_mod]."""
    maxA = max(A)
    data = {'A': A, 'S': S, 'maxA': maxA, 'cells': []}
    
    for L1 in range(2, max_mod + 1):
        for L2 in range(2, max_mod + 1):
            if L1 == L2: continue
            if math.gcd(L1, L2) != 1: continue
            M = L1 * L2
            if not (maxA < M <= 2 * maxA):
                # Still record but mark as out-of-range
                status = "out_of_range"
            else:
                FA = [f_k(a, S, [L1, L2]) for a in A]
                sidon = is_sidon(FA)
                _, _, reason = certify_sidon_creation(A, S, [L1, L2])
                status = "sidon" if sidon else "fail"
            data['cells'].append({
                'L1': L1, 'L2': L2, 'M': M,
                'status': status
            })
    return data

# Known examples
examples = [
    ([1,2,5,6], 7, "Sidon example"),
    ([0,1,3,8,13], 27, "Complex set"),
]

for A, S, name in examples:
    data = generate_heatmap_data(A, S)
    with open(f'/home/allaun/SilverSight/docs/diagrams/heatmap_{name.replace(" ","_")}.json', 'w') as f:
        json.dump(data, f, indent=2)
    sidon_count = sum(1 for c in data['cells'] if c['status'] == 'sidon')
    fail_count = sum(1 for c in data['cells'] if c['status'] == 'fail')
    out_count = sum(1 for c in data['cells'] if c['status'] == 'out_of_range')
    print(f"{name}: {sidon_count} sidon, {fail_count} fail, {out_count} out-of-range")