"""Tune a two-hidden-layer PyTorch MLP against the hosted TauOpt service with
no package installed: make_moons (CPU, one thread), each setting scored by
validation accuracy (maximize); ``momentum`` exists only when
``optimizer == "sgd"`` (a conditional parameter, absent from the assignments
on the adam branch).

The package-free twin of ``pytorch_mlp.py``: the same network and data (torch
and scikit-learn installed, nothing else), the ask/observe loop over the REST
routes. The seed rides every ask as an explicit ``config`` with ``timebox_mode:
"caps"`` (deterministic work ceilings, so the same seed replays the same
sequence on the service; the service refuses ``"off"`` for a non-exempt user). Needs
``TAUOPT_BASE_URL`` and ``TAUOPT_TOKEN`` (the API https://api.tauopt.com; a token minted on the web app's tokens page); without them it
prints one line and exits 0.

    python pytorch_mlp_rest.py --seed 7
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request


# --- rest helper (identical in every *_rest.py example; tests/test_rest_examples.py pins it) ---
class ApiError(RuntimeError):
    """A non-2xx answer: the service's one error shape,
    {"error": {"type", "message", "details"}}."""

    def __init__(self, status: int, body: dict, headers) -> None:
        err = body.get("error") or {}
        self.status = status
        self.error_type = err.get("type", "http_error")
        self.details = err.get("details")
        ra = headers.get("Retry-After")
        self.retry_after = float(ra) if ra and ra.replace(".", "", 1).isdigit() else None
        super().__init__(f"[{status}] {self.error_type}: {err.get('message', '')}")


class _NoRedirect(urllib.request.HTTPRedirectHandler):
    """Never follow a redirect: a 301 from http:// to https:// would turn a POST into a GET."""

    def redirect_request(self, *args, **kwargs):
        return None


class Api:
    """The hosted REST API over the standard library: a bearer token, JSON in and out."""

    def __init__(
        self, base_url: str, token: str, timeout_s: float = 120.0, wait_s: float = 600.0
    ) -> None:
        self._base = base_url.rstrip("/")
        self._token = token
        self._timeout = timeout_s
        self._wait = wait_s  # the budget per call for the waits a 429 or a 503 names
        self._opener = urllib.request.build_opener(_NoRedirect)
        self.last_exhaustion: dict | None = None  # the 409 space_exhausted details ask() last saw

    def call(
        self, method: str, path: str, body: dict | None = None, wait_s: float | None = None
    ) -> dict | None:
        """One request. A non-2xx answer raises ApiError (a 429 carries retry_after); 204
        returns None; a body that is not JSON raises ApiError whatever the status (the
        service answers JSON or 204; anything else is a proxy's page).

        A 429 that names a wait in Retry-After (``bandwidth_full``, a quota, the rate limit)
        and a 503 ``experiment_busy`` are waited out, the named interval (at least 1 s) in
        sleeps of at most 30 s, and re-sent, within a budget per call: the constructor's
        ``wait_s`` unless the call passes its own (a script with a deadline passes what is
        left of it). A wait that would end past the budget raises the answer instead, with
        no further request. ``experiment_quota_exceeded`` carries no Retry-After and raises
        at once. A failure before the answer is retried twice (see ``_send``); a timeout
        after the request was sent is raised as is.
        """
        budget = self._wait if wait_s is None else wait_s
        started = time.monotonic()
        while True:
            try:
                return self._send(method, path, body)
            except ApiError as e:
                busy = e.status == 503 and e.error_type == "experiment_busy"
                if e.retry_after is None or not (e.status == 429 or busy):
                    raise
                left = max(e.retry_after, 1.0)  # a Retry-After of 0 is never a tight loop
                if time.monotonic() - started + left > budget:
                    raise
                while left > 0:
                    time.sleep(min(left, 30.0))
                    left -= 30.0

    def _send(self, method: str, path: str, body: dict | None) -> dict | None:
        req = urllib.request.Request(
            self._base + path,
            data=None if body is None else json.dumps(body).encode(),
            method=method,
            headers={
                "Authorization": f"Bearer {self._token}",
                "Content-Type": "application/json",
                "Accept": "application/json",
            },
        )
        for attempt in range(3):
            try:
                with self._opener.open(req, timeout=self._timeout) as r:
                    status, raw, headers = r.status, r.read(), r.headers
                break
            except urllib.error.HTTPError as e:  # every non-2xx answer, the 3xx included
                status, raw, headers = e.code, e.read(), e.headers
                break
            except urllib.error.URLError:
                # a failure before the answer: DNS, connect, the TLS handshake, or the send
                # itself (urllib wraps all of them in URLError; whether any byte reached the
                # service is unknowable here, so a resend of a landed POST is possible in
                # principle) -- retried twice; a timeout while WAITING for the answer is a
                # bare TimeoutError, not caught here, because the request landed: list the
                # open suggestions first
                if attempt == 2:
                    raise
                time.sleep(1.0 + attempt)
        if status == 204:
            return None
        try:
            data = json.loads(raw)
        except ValueError:
            data = {"error": {"type": "http_error", "message": raw[:200].decode(errors="replace")}}
            raise ApiError(status, data, headers) from None  # a 2xx that is not JSON too
        if status >= 300:
            raise ApiError(status, data, headers)
        return data


def ask(
    api: Api, exp_id, body: dict | None = None, wait_s: float | None = None
) -> dict | None:
    """One suggestion, or None on proven exhaustion (a completion, not an error;
    the 409's details stay on ``api.last_exhaustion``).

    ``body`` is the optional SuggestionIn (a ``config``, ``max_wait_s``). A 429 that
    names a wait and a 503 ``experiment_busy`` are waited out by ``Api.call`` (``wait_s``
    is this call's budget); ``experiment_quota_exceeded`` carries no wait: archive an
    experiment first.
    """
    try:
        return api.call(
            "POST", f"/v1/experiments/{exp_id}/suggestions", body or {}, wait_s=wait_s
        )
    except ApiError as e:
        if e.status == 409 and e.error_type == "space_exhausted":
            api.last_exhaustion = e.details
            return None
        raise


def archive(api: Api, exp_id, wait_s: float | None = None) -> None:
    """DELETE the experiment; a 404 means it is already archived (the call is idempotent).
    ``wait_s`` is this call's budget for a wait the answer names."""
    try:
        api.call("DELETE", f"/v1/experiments/{exp_id}", wait_s=wait_s)
    except ApiError as e:
        if e.error_type != "not_found":
            raise


# --- end rest helper ---


def show(a: dict) -> str:
    """One line of assignments, floats rounded for display."""
    return " ".join(f"{k}={v:.6g}" if isinstance(v, float) else f"{k}={v}"
                    for k, v in a.items())


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.strip().split("\n\n")[0])
    ap.add_argument(
        "--seed", type=int, default=7,
        help="the seed every ask's config carries, so a rerun replays the same points"
             " (default %(default)s)",
    )
    ap.add_argument(
        "--budget", type=int, default=15,
        help="ask/observe rounds before the result is read (default %(default)s)",
    )
    ap.add_argument("--keep", action="store_true", help="leave the experiment active at the end")
    args = ap.parse_args()
    base_url, token = os.environ.get("TAUOPT_BASE_URL"), os.environ.get("TAUOPT_TOKEN")
    if not base_url or not token:
        pair = (("TAUOPT_BASE_URL", base_url), ("TAUOPT_TOKEN", token))
        missing = "/".join(name for name, value in pair if not value)
        print(f"{missing} not set; skipping"
              " (export TAUOPT_BASE_URL=https://api.tauopt.com and a TAUOPT_TOKEN minted at app.tauopt.com/settings/tokens)")
        return 0

    try:
        import torch
        import torch.nn.functional as F
    except ImportError as e:
        raise ImportError("this example needs torch: pip install torch (CPU build)") from e
    try:
        from sklearn.datasets import make_moons
    except ImportError as e:
        raise ImportError(
            "this example needs scikit-learn: pip install scikit-learn"
        ) from e

    # one thread: the run time is then a function of the settings alone
    torch.set_num_threads(1)
    X, y = make_moons(n_samples=600, noise=0.3, random_state=0)
    X = torch.tensor(X, dtype=torch.float32)
    y = torch.tensor(y, dtype=torch.long)
    X_train, y_train, X_val, y_val = X[:400], y[:400], X[400:], y[400:]

    def val_accuracy(a: dict) -> float:
        # a fixed init seed keeps the objective a pure function of the
        # assignments, so the optimizer sees the setting, not init luck
        torch.manual_seed(0)
        width = int(a["width"])
        net = torch.nn.Sequential(
            torch.nn.Linear(2, width), torch.nn.ReLU(),
            torch.nn.Linear(width, width), torch.nn.ReLU(),
            torch.nn.Linear(width, 2),
        )
        if a["optimizer"] == "adam":
            opt = torch.optim.Adam(net.parameters(), lr=float(a["lr"]))
        else:
            # momentum is present exactly when the sgd branch is active
            opt = torch.optim.SGD(net.parameters(), lr=float(a["lr"]),
                                  momentum=float(a["momentum"]))
        for _ in range(200):
            opt.zero_grad()
            F.cross_entropy(net(X_train), y_train).backward()
            opt.step()
        with torch.no_grad():
            return float((net(X_val).argmax(1) == y_val).float().mean())

    api = Api(base_url, token)
    exp = api.call(
        "POST",
        "/v1/experiments",
        {
            "name": f"mlp-moons-rest-seed{args.seed}",
            "parameters": [
                {"name": "width", "type": "integer", "grid": [8, 16, 32, 64]},
                {"name": "lr", "type": "double", "bounds": [1e-4, 1e-1], "log_scale": True},
                {"name": "optimizer", "type": "categorical", "categorical_values": ["adam", "sgd"]},
                # active only on the sgd branch; adam has no momentum knob
                {"name": "momentum", "type": "double", "bounds": [0.0, 0.99],
                 "conditions": {"optimizer": ["sgd"]}},
            ],
            "metric": {"name": "val_accuracy", "objective": "maximize"},
            "observation_budget": args.budget,
            "metadata": {"seed": args.seed},
        },
    )
    exp_id = exp["id"]
    try:
        # caps: deterministic work ceilings, so same-seed runs match on the service
        config = {"strategy": "auto", "seed": args.seed, "timebox_mode": "caps"}
        for i in range(args.budget):
            sug = ask(api, exp_id, body={"config": config})
            if sug is None:
                break
            acc = val_accuracy(sug["assignments"])
            api.call(
                "POST",
                f"/v1/experiments/{exp_id}/observations",
                {"suggestion_id": sug["id"], "value": acc},
            )
            # momentum shows up on the sgd rows only: the engine drops inactive
            # conditional parameters from the assignments
            print(f"[{i:02d}] {sug['provenance']['strategy']:>6} val_accuracy={acc:.4f}"
                  f"  {show(sug['assignments'])}")

        best = api.call("GET", f"/v1/experiments/{exp_id}/best_assignments")
        assert best["data"], "no successful observations"
        top = best["data"][0]
        print(f"best val_accuracy={top['value']:.4f}  {show(top['assignments'])}")
    finally:
        if not args.keep:
            failing = sys.exc_info()[1]  # the error above, if any, stays the one reported
            try:
                archive(api, exp_id)
                print(f"experiment {exp_id} archived")
            except Exception as e:  # noqa: BLE001 -- never hide the error above
                if failing is None:
                    raise
                print(f"experiment {exp_id} not archived: {e}", file=sys.stderr)
    return 0


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