"""Frequency Threshold Experiment — Test different min_coverage values.

The current pipeline uses min_coverage=0.2 fixed. This tests what happens
at 0.01, 0.05, 0.10, 0.15, 0.20 on all 4 codebases.
"""

import json
import time
from collections import Counter, defaultdict
from pathlib import Path

from bex.tag_preprocessor.analyze import (
    _preprocess_files, _file_to_package, scan_directory, frequency_filter
)
from bex.tag_preprocessor.code import _extract_call_tokens, _extract_coarsened_tokens
from bex.twotinf import build_soa
from bex.rwr0 import rwr0


RESULTS_DIR = Path(__file__).parent / "results"


CODEBASES = {
    "ragsak": {"name": "RAGSAK", "path": "/home/tobi/Desktop/kesai/RAGSAK", "ext": ".kt"},
    "flask": {"name": "Flask", "path": "/home/tobi/Desktop/dervish/external_refs/flask", "ext": ".py"},
    "coroutines": {"name": "Kotlin Coroutines", "path": "/home/tobi/Desktop/dervish/projects/grammar-inference-engine/external_refs/kotlinx.coroutines", "ext": ".kt"},
    "fastapi": {"name": "FastAPI", "path": "/home/tobi/Desktop/dervish/projects/grammar-inference-engine/external_refs/fastapi", "ext": ".py"},
}

THRESHOLDS = [0.0, 0.01, 0.02, 0.05, 0.10, 0.15, 0.20]


def load_data(base, ext):
    groups = scan_directory(base)
    files = groups.get(ext, [])
    sequences, seq_files = _preprocess_files(files)
    raw_seqs = [_extract_call_tokens(seq) for seq in sequences]
    packages = [_file_to_package(fp, base) for fp in seq_files]
    return raw_seqs, packages, seq_files


def measure_at_threshold(raw_seqs, packages, threshold):
    """Measure what happens at a given frequency threshold."""
    # Count symbol frequency
    sym_file = Counter()
    for seq in raw_seqs:
        for sym in set(seq):
            sym_file[sym] += 1

    total = len(raw_seqs)
    n_keep = sum(1 for c in sym_file.values() if c >= max(1, int(total * threshold)))
    keep_syms = {s for s, c in sym_file.items() if c >= max(1, int(total * threshold))}
    surviving = sum(1 for seq in raw_seqs if any(s in keep_syms for s in seq))

    # Filter sequences
    wrapped = [[(j, s, 0) for j, s in enumerate(seq)] for seq in raw_seqs]
    filtered = frequency_filter(wrapped, min_coverage=threshold)
    clean = [[text for _, text, _ in r] for r in filtered]
    clean = [s for s in clean if s]

    # Group by first 3 symbols for inference
    contexts = defaultdict(list)
    for i, seq in enumerate(clean):
        if not seq:
            contexts[("_eps",)].append(seq)
        else:
            ctx = tuple(seq[:3])
            contexts[ctx].append(seq)

    sore_successes = 0
    methods_in_good = 0
    total_methods = 0
    top_patterns = []

    for ctx, seqs in sorted(contexts.items(), key=lambda x: -len(x[1])):
        n = len(seqs)
        total_methods += n
        if n < 3:
            continue
        unique = len(set(tuple(s) for s in seqs))
        if unique / n > 0.9:
            continue
        alphabet = set()
        for s in seqs:
            alphabet.update(s)
        if len(alphabet) > 20:
            continue
        try:
            soa = build_soa(seqs)
            sore = rwr0(soa)
            if sore != "∅":
                sore_successes += 1
                methods_in_good += n
                top_patterns.append({
                    "context": str(ctx),
                    "methods": n,
                    "unique": unique,
                    "sore": sore[:150],
                })
        except Exception:
            pass

    return {
        "threshold": threshold,
        "symbols_total": len(sym_file),
        "symbols_kept": n_keep,
        "seqs_total": total,
        "seqs_surviving": surviving,
        "contexts": len(contexts),
        "sore_successes": sore_successes,
        "total_methods": total_methods,
        "methods_in_good": methods_in_good,
        "coverage": round(methods_in_good / total_methods * 100, 1) if total_methods else 0,
        "top_patterns": top_patterns[:5],
    }


def run_codebase(key):
    cfg = CODEBASES[key]
    print(f"\n{'=' * 60}")
    print(f"  {cfg['name']}")
    print(f"{'=' * 60}")

    raw_seqs, packages, seq_files = load_data(cfg["path"], cfg["ext"])
    print(f"  {len(raw_seqs)} methods from {len(set(packages))} packages")

    results = []
    for thresh in THRESHOLDS:
        t0 = time.time()
        r = measure_at_threshold(raw_seqs, packages, thresh)
        elapsed = time.time() - t0
        r["elapsed"] = round(elapsed, 2)
        results.append(r)
        print(f"  thresh={thresh:.2f}  syms={r['symbols_kept']:4d}  seqs={r['seqs_surviving']:5d}  "
              f"ctxs={r['contexts']:4d}  SOREs={r['sore_successes']:3d}  cov={r['coverage']:5.1f}%  "
              f"({elapsed:.1f}s)")

    out_path = RESULTS_DIR / f"freq_{key}.json"
    with open(out_path, "w") as f:
        json.dump({"codebase": cfg["name"], "results": results}, f, indent=2)
    print(f"  Saved to {out_path}")

    return results


def main():
    all_results = {}
    for key in CODEBASES:
        all_results[key] = run_codebase(key)

    # Summary table
    print(f"\n{'=' * 80}")
    print("  SUMMARY")
    print(f"{'=' * 80}")
    print(f"  {'Codebase':20s} {'Thresh':7s} {'Syms':5s} {'Seqs':6s} {'SOREs':6s} {'Cov':7s}")
    print(f"  {'-'*55}")
    for key, results in all_results.items():
        name = CODEBASES[key]["name"]
        for r in results:
            print(f"  {name:20s} {r['threshold']:7.2f} {r['symbols_kept']:5d} "
                  f"{r['seqs_surviving']:6d} {r['sore_successes']:6d} {r['coverage']:6.1f}%")

    return all_results


if __name__ == "__main__":
    main()