"""Feed a learner a long stream and watch it learn, inside fixed RAM and time limits.

Stops cleanly at whichever comes first: the end of the data, the memory cap, or the
time limit. Runs the learner inside this process, so only use learner files you trust.

  python feed.py --learner runs/first/cands/00017/learner.py --data data/dl/enwik8 \
                 --max-mb 16 --mem-mb 5000 --minutes 60 --out runs/first/wiki_feed.json
"""
from __future__ import annotations

import argparse
import bz2
import json
import lzma
import time
import zlib
from pathlib import Path

from compsearch.classify import load_learner_factory
from compsearch.limits import confine_own_process
from compsearch.runner import byte_cost_bits

MB = 1024 * 1024
REPORT_EVERY = MB // 4


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--learner", required=True)
    parser.add_argument("--data", required=True)
    parser.add_argument("--max-mb", type=float, default=16.0)
    parser.add_argument("--mem-mb", type=int, default=4000)
    parser.add_argument("--minutes", type=float, default=60.0)
    parser.add_argument("--out")
    args = parser.parse_args()

    with open(args.data, "rb") as handle:
        data = handle.read(int(args.max_mb * MB))
    factory = load_learner_factory(args.learner)
    if not confine_own_process(args.mem_mb):
        raise SystemExit("could not apply the memory cap; not running")

    learner, started = factory(), time.monotonic()
    windows, window_bits, total_bits, done, stopped = [], 0.0, 0.0, 0, "end of data"
    try:
        for byte in data:
            bits = byte_cost_bits(learner.predict(), byte)
            learner.update(byte)
            window_bits += bits
            total_bits += bits
            done += 1
            if done % REPORT_EVERY == 0:
                windows.append(round(window_bits / REPORT_EVERY, 4))
                window_bits = 0.0
                elapsed = time.monotonic() - started
                print(f"{done / MB:5.2f} MB  this quarter-MB {windows[-1]:.3f} bpb  "
                      f"overall {total_bits / done:.3f} bpb  {elapsed / 60:.1f} min", flush=True)
                if elapsed > args.minutes * 60:
                    stopped = f"time limit ({args.minutes:.0f} minutes)"
                    break
    except MemoryError:
        learner = None  # release memory so the report below can be built
        stopped = f"memory cap ({args.mem_mb} MB)"

    seen = data[:done]
    report = {
        "learner": args.learner, "data": args.data, "stopped_by": stopped,
        "megabytes_read": round(done / MB, 2),
        "minutes": round((time.monotonic() - started) / 60, 1),
        "overall_bpb": round(total_bits / max(done, 1), 4),
        "bpb_per_quarter_megabyte": windows,
        "same_bytes_with_standard_compressors": {
            "zlib -9": round(8 * len(zlib.compress(seen, 9)) / max(done, 1), 4),
            "bz2 -9": round(8 * len(bz2.compress(seen, 9)) / max(done, 1), 4),
            "xz -9e": round(8 * len(lzma.compress(seen, preset=9 | lzma.PRESET_EXTREME)) / max(done, 1), 4),
        },
    }
    print(json.dumps(report, indent=1))
    if args.out:
        Path(args.out).write_text(json.dumps(report, indent=1), encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())