#!/usr/bin/env python3
"""
DAG Tuning & Analysis: map iteration behavior, find optimal moduli.
"""
import sys, math, random, itertools, json
from collections import Counter
sys.path.insert(0, '/home/allaun/SilverSight/scripts')
from verify_wrapping import f_k, is_sidon, sum_collisions, wrapping_criterion, certify_sidon_creation
from iteration_dag import IterationDAG, AdaptiveRule, GeometricRule, DAGNode

# ---------- Modulus Effectiveness Analysis ----------

def test_all_moduli(A, S, max_modulus=20):
    """Test ALL coprime modulus pairs in the valid range, return effectiveness map."""
    results = []
    maxA = max(A)
    for L1 in range(2, max_modulus + 1):
        for L2 in range(L1 + 1, max_modulus + 1):
            if math.gcd(L1, L2) != 1:
                continue
            M = L1 * L2
            if not (maxA < M <= 2 * maxA):
                continue
            guaranteed, FA, reason = certify_sidon_creation(A, S, [L1, L2])
            results.append({
                'moduli': [L1, L2], 'M': M,
                'sidon': is_sidon(FA),
                'guaranteed': guaranteed,
                'FA': FA,
                'reason': reason
            })
    return results

def modulus_heatmap(A, S, max_modulus=20):
    """Generate a heatmap of modulus effectiveness."""
    results = test_all_moduli(A, S, max_modulus)
    if not results:
        print("  No valid moduli in range")
        return {}
    
    sidon_count = sum(1 for r in results if r['sidon'])
    guaranteed_count = sum(1 for r in results if r['guaranteed'])
    
    # Best moduli by Sidon creation
    sidon_mods = [r for r in results if r['sidon']]
    
    stats = {
        'total_moduli': len(results),
        'sidon_success': sidon_count,
        'guaranteed_sidon': guaranteed_count,
        'success_rate': sidon_count / max(len(results), 1),
        'guarantee_rate': guaranteed_count / max(len(results), 1),
        'best_moduli': sidon_mods[:10] if len(sidon_mods) <= 10 else sidon_mods[:10],
        'worst_moduli': [r for r in results if not r['sidon']][:5]
    }
    return stats

# ---------- DAG Depth Analysis ----------

def depth_distribution(A0, S0, max_steps=6):
    """Analyze the distribution of path lengths to Sidon."""
    rule = AdaptiveRule(max_val=16)
    dag = IterationDAG(A0, S0, rule, max_steps=max_steps)
    dag.build()
    
    path_lengths = []
    for path in dag.sidon_paths:
        path_lengths.append(len(path) - 1)  # steps, not nodes
    
    return {
        'total_nodes': len(dag.all_nodes),
        'sidon_paths': len(dag.sidon_paths),
        'path_lengths': dict(Counter(path_lengths)),
        'min_steps': min(path_lengths) if path_lengths else None,
        'max_steps': max(path_lengths) if path_lengths else None,
        'avg_steps': sum(path_lengths) / len(path_lengths) if path_lengths else None
    }

# ---------- Parameter Sweep ----------

def sweep_parameter(target_property="sidon", trials=200, max_modulus=16, max_steps=4):
    """Sweep across random A sets and find optimal tuning strategies."""
    random.seed(42)
    primes = [2,3,5,7,11,13,17,19,23,29,31,37]
    results = []
    
    for trial in range(trials):
        # Generate random A
        maxA = random.randint(5, 30)
        n = random.randint(4, 8)
        A = sorted(random.sample(range(maxA + 1), min(n, maxA + 1)))
        S = random.randint(maxA, maxA + 10)
        
        # Test single-step: find modulus pairs that produce Sidon in one step
        mod_results = test_all_moduli(A, S, max_modulus)
        one_step_sidon = sum(1 for r in mod_results if r['sidon'])
        
        # Test DAG: find multi-step paths to Sidon
        rule = AdaptiveRule(max_val=max_modulus)
        dag = IterationDAG(A, S, rule, max_steps=max_steps)
        dag.build()
        multi_step = len(dag.sidon_paths)
        
        # Find the smallest M that works
        min_sidon_M = None
        for r in mod_results:
            if r['sidon']:
                if min_sidon_M is None or r['M'] < min_sidon_M:
                    min_sidon_M = r['M']
        
        results.append({
            'A': A, 'S': S, 'n': len(A), 'maxA': maxA,
            'one_step_candidates': one_step_sidon,
            'one_step_total': len(mod_results),
            'one_step_rate': one_step_sidon / max(len(mod_results), 1),
            'multi_step_paths': multi_step,
            'min_sidon_M': min_sidon_M
        })
    
    return results

# ---------- Analysis Reports ----------

def report_sidon_example():
    """Detailed analysis of the known Sidon creation example."""
    A, S = [1, 2, 5, 6], 7
    print("=" * 60)
    print(f"SIDON EXAMPLE ANALYSIS: A={A}, S={S}")
    print("=" * 60)
    
    stats = modulus_heatmap(A, S)
    print(f"\nModulus Analysis ({stats['total_moduli']} coprime pairs in range):")
    print(f"  Sidon creation success: {stats['sidon_success']}/{stats['total_moduli']} ({stats['success_rate']*100:.1f}%)")
    print(f"  Guaranteed Sidon: {stats['guaranteed_sidon']}/{stats['total_moduli']} ({stats['guarantee_rate']*100:.1f}%)")
    print(f"  Best moduli (first 10 Sidon-creating pairs):")
    for r in stats['best_moduli']:
        print(f"    [{r['moduli'][0]}, {r['moduli'][1]}] M={r['M']} FA={r['FA']}")

def report_complex_set():
    """Quick analysis of a more complex set — moduli only, no DAG."""
    A, S = [0, 1, 3, 8, 13], 27
    print("\n" + "=" * 60)
    print(f"COMPLEX SET: A={A}, S={S}")
    print("=" * 60)
    
    stats = modulus_heatmap(A, S, max_modulus=16)
    print(f"\nModulus Analysis ({stats['total_moduli']} coprime pairs in range):")
    pct = stats['success_rate'] * 100
    print(f"  Sidon creation: {stats['sidon_success']}/{stats['total_moduli']} ({pct:.1f}%)")
    print(f"  Guaranteed: {stats['guaranteed_sidon']}/{stats['total_moduli']} ({stats['guarantee_rate']*100:.1f}%)")
    for r in stats['best_moduli'][:5]:
        print(f"    [{r['moduli'][0]}, {r['moduli'][1]}] M={r['M']} FA={r['FA']}")

def report_sweep():
    """Fast sweep — moduli only, no DAG building."""
    print("\n" + "=" * 60)
    print("PARAMETER SWEEP (200 random sets — modulus-only)")
    print("=" * 60)
    
    random.seed(42)
    results = []
    
    for trial in range(200):
        maxA = random.randint(5, 30)
        n = random.randint(4, 8)
        A = sorted(random.sample(range(maxA + 1), min(n, maxA + 1)))
        S = random.randint(maxA, maxA + 10)
        
        mod_results = test_all_moduli(A, S, max_modulus=12)
        one_step_sidon = sum(1 for r in mod_results if r['sidon'])
        guaranteed = sum(1 for r in mod_results if r['guaranteed'])
        
        results.append({
            'n': len(A), 'maxA': maxA,
            'one_step_sidon': one_step_sidon,
            'total_moduli': len(mod_results),
            'guaranteed': guaranteed,
            'success_rate': one_step_sidon / max(len(mod_results), 1) if mod_results else 0,
        })
    
    sr = [r['success_rate'] for r in results]
    print(f"\nResults ({len(results)} sets):")
    print(f"  Sets with >0 valid moduli: {sum(1 for r in results if r['total_moduli'] > 0)}/{len(results)}")
    print(f"  Sets with at least one Sidon-creating modulus: {sum(1 for r in results if r['one_step_sidon'] > 0)}/{len(results)}")
    print(f"  Avg success rate: {sum(sr)/len(sr)*100:.1f}%")
    print(f"  Best success rate: {max(sr)*100:.1f}%")

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

if __name__ == "__main__":
    report_sidon_example()
    report_complex_set()
    report_sweep()
