from __future__ import annotations

import argparse
import json
import random
import sys
import time
from pathlib import Path

import torch


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("ckpt_dir")
    ap.add_argument("--widths", type=int, nargs="*", default=[35, 67])
    ap.add_argument("--sub", default="submission_a")
    ap.add_argument("--problems-per-round", type=int, default=40)
    args = ap.parse_args()

    sys.path.insert(0, str(Path(__file__).resolve().parent))
    from model import (make_reduce_cell, make_add_cell, reduce_features,
                       add_features, shift_bits, _bits_of, PAD_HEAD)

    torch.set_num_threads(4)
    rng = random.Random()
    ck_path = Path(args.ckpt_dir) / "latest.pt"
    out_r = Path(args.ckpt_dir) / "mined_reduce.jsonl"
    out_a = Path(args.ckpt_dir) / "mined_add.jsonl"

    def is_pp(n):
        if n < 2:
            return False
        for sp in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31):
            if n % sp == 0:
                return n == sp
        d, r = n - 1, 0
        while d % 2 == 0:
            d //= 2
            r += 1
        for _ in range(20):
            a = rng.randrange(2, n - 1)
            x = pow(a, d, n)
            if x in (1, n - 1):
                continue
            for _ in range(r - 1):
                x = x * x % n
                if x == n - 1:
                    break
            else:
                return False
        return True

    def to_bits(v, w):
        t = torch.zeros(1, w)
        b = _bits_of(v)
        t[0, w - len(b):] = torch.tensor(b, dtype=torch.float32)
        return t

    def val(t):
        return int("".join(str(int(x)) for x in t[0].tolist()), 2)

    R = make_reduce_cell()
    A = make_add_cell()
    last_load = 0.0
    n_mined = 0

    while True:
        if time.time() - last_load > 180:
            try:
                ck = torch.load(ck_path, map_location="cpu",
                                weights_only=True)
                R.load_state_dict(ck.get("reduce_ema_state_dict",
                                         ck["reduce_state_dict"]))
                A.load_state_dict(ck.get("add_ema_state_dict",
                                         ck["add_state_dict"]))
                R.eval()
                A.eval()
                last_load = time.time()
            except Exception:
                time.sleep(10)
                continue

        N = rng.choice(args.widths)
        pb_hi = N - PAD_HEAD
        pb_lo = max(2, pb_hi // 2 + 1)
        def draw_pb():
            r = rng.random()
            if r < 0.45:
                return pb_hi
            if r < 0.70:
                return max(pb_lo, pb_hi - 1)
            return rng.randint(pb_lo, pb_hi)
        L = 3 * pb_hi
        mr, ma = [], []
        with torch.no_grad():
            for _ in range(args.problems_per_round):
                pb = draw_pb()
                if rng.random() < 0.35 and pb >= 9:
                    p = 0
                    for c in range(1, 400, 2):
                        cand = (1 << pb) - c
                        if cand > 2 and is_pp(cand):
                            p = cand
                            break
                    if not p:
                        p = (1 << (pb - 1)) | 1
                        while not is_pp(p):
                            p += 2
                else:
                    while True:
                        p = rng.getrandbits(pb - 1) | (1 << (pb - 1)) | 1
                        if p > 2 and is_pp(p):
                            break
                a = rng.getrandbits(rng.randint(1, L))
                b = rng.getrandbits(rng.randint(1, L))
                pt, p3t = to_bits(p, N), to_bits(3 * p, N)
                residues = []
                for op in (a, b):
                    ob = _bits_of(op)
                    if len(ob) % 2:
                        ob = [0] + ob
                    Xv = 0
                    for t in range(0, len(ob), 2):
                        xv = 4 * Xv + 2 * ob[t] + ob[t + 1]
                        x = torch.cat(
                            [to_bits(Xv, N)[:, 2:],
                             torch.tensor([[float(ob[t]),
                                            float(ob[t + 1])]])], dim=1)
                        got = val((R(reduce_features(x, pt, p3t)) > 0).float())
                        want = xv % p
                        if got != want:
                            mr.append({"n": N, "m": p, "x": xv})
                        Xv = want
                    residues.append(Xv)
                ra, rb = residues
                rab = _bits_of(ra)
                rab = [0] * (N - PAD_HEAD - len(rab)) + rab
                yt = to_bits(rb, N)
                Zv = 0
                for g in rab:
                    sv = 2 * Zv + g * rb
                    got = val((A(add_features(
                        shift_bits(to_bits(Zv, N), 1), yt,
                        torch.tensor([float(g)]))) > 0).float())
                    if got != sv:
                        ma.append({"n": N, "x": 2 * Zv, "y": rb, "g": g})
                    got2 = val((R(reduce_features(
                        to_bits(sv, N), pt, p3t)) > 0).float())
                    want = sv % p
                    if got2 != want:
                        mr.append({"n": N, "m": p, "x": sv})
                    Zv = want
        if mr:
            with open(out_r, "a") as f:
                for row in mr:
                    f.write(json.dumps(row) + "\n")
        if ma:
            with open(out_a, "a") as f:
                for row in ma:
                    f.write(json.dumps(row) + "\n")
        n_mined += len(mr) + len(ma)
        print(f"mined so far: {n_mined} (+{len(mr)}r +{len(ma)}a @N={N})",
              flush=True)


if __name__ == "__main__":
    main()