#!/usr/bin/env python3
"""Search for optimal batch sizes across a range of -log2(p) values.

Usage:
  tune_search.py create FILE BITS LOW HIGH   initialize FILE with given range
  tune_search.py run FILE                    continue refining FILE
  tune_search.py emit FILE [FILE ...]        print C++ TuneEntry arrays for elias_sim.cpp

LOW must be at least 1. All raw observations are kept in memory. For each
batch value with n>=2 observations, the low-half and high-half of the sorted
observations are averaged to form an interval [lo, hi] (singletons have
lo == hi). The file stores bits on the first line, then one `lo hi batch`
triple per line. `run` picks points in regions where coverage count != 1 --
gaps (no interval covers the range) or overlaps (two+ intervals cover it) --
evaluates there, and appends the new observation. State is rewritten every
10s via write-to-temp + rename.
"""
import argparse
import math
import os
import random
import sys
import time
from collections import defaultdict

from mpmath import mp

mp.prec = 256


def elias_output(r):
    out = 0
    while r > 0:
        b = r.bit_length() - 1
        out += b * (1 << b)
        r -= (1 << b)
    return out

def eval_ours(batch, p, bits):
    """Evaluate entropy rate (bits per toss) extracted using our extractor."""
    q = 1 - p
    ret = 0
    mask = (1 << bits) - 1
    for k in range(batch + 1):
        prob = p**k * q**(batch - k)
        cmb = math.comb(batch, k)
        overflows = cmb >> bits
        r = cmb - (overflows << bits)
        if r > 0:
            # Account for cases where v < r.
#            ret += mp.log(r, 2) * r * (overflows + 1) * prob
            ret += elias_output(r) * (overflows + 1) * prob
        if overflows > 0:
            # Account for cases where r <= v <= mask.
#            ret += mp.log(mask - r + 1, 2) * (mask - r + 1) * overflows * prob
            ret += elias_output(mask - r + 1) * overflows * prob
    return ret / batch

BATCH_CACHE = {}

def find_best_batch(p, bits, batch_range=None):
    """Find optimal batch size for given p and bits.

    batch_range: optional (lo, hi) tuple restricting the search. Either side
    may be None to leave that side unbounded. When provided, the endpoints
    are added to the initial set of candidates.
    """
    if batch_range is None:
        range_lo, range_hi = None, None
    else:
        range_lo, range_hi = batch_range
    min_lo = max(2, range_lo) if range_lo is not None else 2
    if range_hi is not None and range_hi < min_lo:
        raise ValueError(f"batch_range {batch_range} is empty")

    def in_range(v):
        if v < min_lo:
            return False
        if range_hi is not None and v > range_hi:
            return False
        return True

    if bits not in BATCH_CACHE:
        BATCH_CACHE[bits] = []
    bcache = BATCH_CACHE[bits]

    res = []
    tried = set()
    def add(v):
        if in_range(v) and v not in tried:
            res.append((eval_ours(v, p, bits), v))
            tried.add(v)

    # Formula-based guess, clamped into range.
    guess = int(bits / (p * mp.log(p, 0.5) + (1 - p) * mp.log(1 - p, 0.5)) + 0.5)
    guess = max(min_lo, guess)
    if range_hi is not None:
        guess = min(range_hi, guess)
    add(guess)

    # Explicit range endpoints (if provided).
    if range_lo is not None:
        add(min_lo)
    if range_hi is not None:
        add(range_hi)

    # Nearest cached p value (if any in range).
    found = None
    for pval, batch in bcache:
        if not in_range(batch):
            continue
        if found is None or abs(p - pval) < abs(p - found[0]):
            found = (pval, batch)
    if found is not None:
        add(found[1])

    while True:
        _, best_batch = max(res)
        win_lo = max(min_lo, min(best_batch - 30, int(best_batch * 0.49 + 0.5)))
        win_hi = max(best_batch + 31, int(best_batch * 2.04 + 1.5))
        if range_hi is not None:
            win_hi = min(range_hi + 1, win_hi)
        todo = [v for v in range(win_lo, win_hi) if v not in tried]
        if not todo:
            break
        batch = random.choice(todo)
        add(batch)
    _, best_batch = max(res)
    bcache.append((p, best_batch))
    return best_batch


def batch_for_log2p(log2p, bits, batch_range=None):
    p = mp.mpf(2) ** (-mp.mpf(log2p))
    return find_best_batch(p, bits, batch_range=batch_range)


def compute_intervals(observations):
    """Derive one (lo, hi, batch) tuple per distinct batch from raw observations.

    For batches with n>=2 observations, lo is the average of the bottom n//2
    sorted log2p values and hi is the average of the top n//2 (the middle one
    is dropped when n is odd). Singletons produce lo == hi == the single
    observation.
    """
    groups = defaultdict(list)
    for log2p, batch in observations:
        groups[batch].append(log2p)
    intervals = []
    for batch, xs in groups.items():
        xs.sort()
        n = len(xs)
        if n == 1:
            intervals.append((xs[0], xs[0], batch))
        else:
            half = n // 2
            lower = xs[:half]
            upper = xs[-half:]
            lo = sum(lower) / len(lower)
            hi = sum(upper) / len(upper)
            intervals.append((lo, hi, batch))
    intervals.sort(key=lambda t: (t[0], t[1], t[2]))
    return intervals


def find_work_regions(full_lo, full_hi, covered):
    """Sweep-line: return disjoint (lo, hi) regions in [full_lo, full_hi] where
    coverage count != 1 (i.e. gaps with count 0, or overlaps with count >= 2).

    `covered` is a list of (lo, hi) with lo < hi.
    """
    events = []
    for lo, hi in covered:
        events.append((lo, +1))
        events.append((hi, -1))
    events.sort()
    segments = []
    pos = full_lo
    count = 0
    for x, delta in events:
        xx = max(full_lo, min(full_hi, x))
        if xx > pos:
            segments.append((pos, xx, count))
            pos = xx
        count += delta
    if pos < full_hi:
        segments.append((pos, full_hi, count))
    work = [(lo, hi) for lo, hi, c in segments if c != 1]
    merged = []
    for lo, hi in work:
        if merged and merged[-1][1] == lo:
            merged[-1] = (merged[-1][0], hi)
        else:
            merged.append([lo, hi])
    return [tuple(r) for r in merged]


def sample_work_point(regions):
    """Pick a point uniformly over the union of (lo, hi) regions by total length."""
    weights = [hi - lo for lo, hi in regions]
    total = sum(weights)
    r = random.uniform(0, total)
    for (lo, hi), w in zip(regions, weights):
        if r <= w:
            while True:
                pt = random.uniform(lo, hi)
                if lo < pt < hi:
                    return pt
        r -= w
    lo, hi = regions[-1]
    while True:
        pt = random.uniform(lo, hi)
        if lo < pt < hi:
            return pt


def batch_range_for_point(mid, intervals):
    """Pick a batch_range hint for find_best_batch at `mid`.

    If any intervals cover mid (overlap case), use (min, max) of covering
    batches. Otherwise (gap) use the batches of the nearest intervals on
    each side.
    """
    covering = [b for lo, hi, b in intervals if lo <= mid <= hi]
    if covering:
        return (min(covering), max(covering))
    left = [(hi, b) for lo, hi, b in intervals if hi < mid]
    right = [(lo, b) for lo, hi, b in intervals if lo > mid]
    candidates = []
    if left:
        _, lb = max(left, key=lambda p: p[0])
        candidates.append(lb)
    if right:
        _, rb = min(right, key=lambda p: p[0])
        candidates.append(rb)
    if candidates:
        return (min(candidates), max(candidates))
    return None


def save(path, bits, observations):
    """Write the derived intervals (lo hi batch per line) atomically."""
    intervals = compute_intervals(observations)
    tmp = path + '.tmp'
    with open(tmp, 'w') as f:
        f.write(f"bits {bits}\n")
        for lo, hi, batch in intervals:
            f.write(f"{lo!r} {hi!r} {batch}\n")
    os.replace(tmp, path)


def load(path):
    """Parse file into (bits, observations); each interval reconstitutes as 1 or 2 obs."""
    with open(path) as f:
        lines = [ln.strip() for ln in f if ln.strip()]
    if not lines:
        sys.exit(f"{path}: file is empty")
    head = lines[0].split()
    if len(head) != 2 or head[0] != 'bits':
        sys.exit(f"{path}: expected first line 'bits N'")
    bits = int(head[1])
    observations = []
    for ln in lines[1:]:
        parts = ln.split()
        if len(parts) != 3:
            sys.exit(f"{path}: expected 'lo hi batch' per line, got {ln!r}")
        lo, hi, batch = float(parts[0]), float(parts[1]), int(parts[2])
        observations.append([lo, batch])
        if hi != lo:
            observations.append([hi, batch])
    return bits, observations


def seed_cache(bits, observations):
    """Prime BATCH_CACHE with loaded (log2p, batch) pairs."""
    BATCH_CACHE.setdefault(bits, [])
    for log2p, batch in observations:
        p = mp.mpf(2) ** (-mp.mpf(log2p))
        BATCH_CACHE[bits].append((p, batch))


def cmd_create(path, bits, lo_exp, hi_exp):
    if os.path.exists(path) and os.path.getsize(path) > 0:
        sys.exit(f"{path}: already exists and is non-empty")
    if lo_exp < 1:
        sys.exit("LOW must be at least 1")
    if not (lo_exp < hi_exp):
        sys.exit("LOW must be less than HIGH")
    print(f"creating {path}: bits={bits}, -log2(p) range [{lo_exp}, {hi_exp}]", flush=True)
    b_lo = batch_for_log2p(lo_exp, bits)
    print(f"  -log2(p)={lo_exp}: batch={b_lo}", flush=True)
    b_hi = batch_for_log2p(hi_exp, bits)
    print(f"  -log2(p)={hi_exp}: batch={b_hi}", flush=True)
    observations = [[float(lo_exp), b_lo], [float(hi_exp), b_hi]]
    save(path, bits, observations)


def cmd_run(path):
    bits, observations = load(path)
    seed_cache(bits, observations)
    print(f"loaded {path}: bits={bits}, {len(observations)} observations", flush=True)
    last_save = time.time()
    try:
        while True:
            intervals = compute_intervals(observations)
            full_lo = min(o[0] for o in observations)
            full_hi = max(o[0] for o in observations)
            covered = [(lo, hi) for lo, hi, b in intervals if lo < hi]
            regions = find_work_regions(full_lo, full_hi, covered)
            regions = [(lo, hi) for lo, hi in regions
                       if hi > lo and math.nextafter(lo, hi) != hi]
            if not regions:
                print("converged: no work regions remaining", flush=True)
                break
            mid = sample_work_point(regions)
            br = batch_range_for_point(mid, intervals)
            n_cov = sum(1 for lo, hi, _ in intervals if lo < hi and lo <= mid <= hi)
            kind = "gap" if n_cov == 0 else f"overlap({n_cov})"
            batch = batch_for_log2p(mid, bits, batch_range=br)
            observations.append([mid, batch])
            print(f"  -log2(p)={mid:.6f}: batch={batch}  [{kind}, range={br}]", flush=True)
            now = time.time()
            if now - last_save >= 10:
                save(path, bits, observations)
                last_save = now
                print(f"  (saved, {len(observations)} observations across {len(intervals)} batches)", flush=True)
    finally:
        save(path, bits, observations)


def cmd_emit(paths):
    """Print C++ TuneEntry arrays for use in elias_sim.cpp.

    Transitions between consecutive batches are placed at the midpoint between
    the end of one batch's interval and the beginning of the next. The first
    entry uses the low end of the smallest batch's interval.
    """
    bits_to_transitions = {}
    for path in paths:
        bits, observations = load(path)
        intervals = compute_intervals(observations)
        intervals.sort(key=lambda t: t[2])  # by batch
        transitions = []
        for i, (lo, hi, batch) in enumerate(intervals):
            if i == 0:
                t = lo
            else:
                prev_hi = intervals[i - 1][1]
                t = (prev_hi + lo) / 2
            transitions.append((t, batch))
        if bits in bits_to_transitions:
            sys.exit(f"duplicate bits={bits} across files")
        bits_to_transitions[bits] = transitions

    for bits in sorted(bits_to_transitions):
        trans = bits_to_transitions[bits]
        entries = ", ".join(f"{{{log2p!r}, {batch}}}" for log2p, batch in trans)
        print(f"constexpr TuneEntry TUNE_{bits}[] = {{")
        for i in range(0, len(trans), 3):
            chunk = trans[i:i+3]
            line = ", ".join(f"{{{log2p!r}, {batch}}}" for log2p, batch in chunk)
            print(f"    {line},")
        print(f"}};")
        print()


def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    sub = ap.add_subparsers(dest='cmd', required=True)

    c = sub.add_parser('create', help='initialize file')
    c.add_argument('file')
    c.add_argument('bits', type=int)
    c.add_argument('low', type=float, help='low -log2(p) endpoint (>= 1)')
    c.add_argument('high', type=float, help='high -log2(p) endpoint')

    r = sub.add_parser('run', help='continue splitting intervals')
    r.add_argument('file')

    e = sub.add_parser('emit', help='print C++ TuneEntry arrays for elias_sim.cpp')
    e.add_argument('files', nargs='+')

    args = ap.parse_args()
    if args.cmd == 'create':
        cmd_create(args.file, args.bits, args.low, args.high)
    elif args.cmd == 'run':
        cmd_run(args.file)
    elif args.cmd == 'emit':
        cmd_emit(args.files)


if __name__ == "__main__":
    main()