"""
afm_client.py -- a small Python client for the AFM digital twin's web API.

    pip install requests numpy
    from afm_client import Twin

    twin = Twin("https://afm.yourdomain.com", key="afm_...")     # the key: web page > API key
    run = twin.image(mode="FFM", points=64, setpoint_f=2e-9, gain=20, hold=True)
    for line in run.lines(curves=True):                          # each line as soon as it is scanned
        print(line.number, line.metrics["error_rms"])
        run.next(scan={"gain": 25})                              # settings for the next line, then go on
    print(run.metrics())
    run.download("forcevolume.h5", "my_volume.h5")

Everything is in SI units (m, N, Hz, V, s).  Settings use the same names as the web page; twin.schema() lists them
with their units and limits.  Only the numbers and files of your own runs come back: the simulator runs on the server.
"""
from __future__ import annotations

import base64
import time

import numpy as np
import requests

__all__ = ["Twin", "Run", "Line", "TwinError"]


class TwinError(RuntimeError):
    """The server refused a request; the message says why."""


def _f32(b64, shape=None):
    a = np.frombuffer(base64.b64decode(b64), dtype=np.float32) if b64 else np.zeros(0, np.float32)
    return a.reshape(shape) if shape is not None else a


def _arr(values):
    return np.array([np.nan if v is None else v for v in values], dtype=float)


class Twin:
    """A connection to the server, as one user (the API key's)."""

    def __init__(self, url: str, key: str = "", timeout: float = 60.0):
        self.url = url.rstrip("/")
        self.session = requests.Session()
        if key:
            self.session.headers["Authorization"] = f"Bearer {key}"
        self.timeout = timeout

    # ---- plumbing -----------------------------------------------------------------------------------------------
    def _call(self, method, path, json=None, params=None, timeout=None, raw=False):
        r = self.session.request(method, self.url + path, json=json, params=params, timeout=timeout or self.timeout)
        if r.status_code >= 400:
            try:
                msg = r.json().get("error", r.text)
            except ValueError:
                msg = r.text
            raise TwinError(f"{r.status_code}: {msg}")
        return r if raw else r.json()

    # ---- settings -----------------------------------------------------------------------------------------------
    def schema(self) -> dict:
        """Every setting with its label, unit and limits, the image channels, the sweepable parameters."""
        return self._call("GET", "/api/me")

    def settings(self) -> dict:
        """The present settings: {'scan': {...}, 'params': {...}, 'force': {...}, 'tune': {...}}."""
        return self._call("GET", "/api/settings")["settings"]

    def set(self, scan=None, params=None, force=None, tune=None) -> dict:
        """Change some settings (they also act on a running image).  Returns the server's answer."""
        body = {k: v for k, v in dict(scan=scan, params=params, force=force, tune=tune).items() if v}
        return self._call("PUT", "/api/settings", json=body)

    def reset_parameters(self):
        return self._call("POST", "/api/settings/reset", json={"part": "params"})

    # ---- runs ---------------------------------------------------------------------------------------------------
    def image(self, hold: bool = False, params: dict | None = None, **scan) -> "Run":
        """Start an image.  scan settings as keywords (mode='Contact' | 'AC Mode' | 'FFM', points, lines, aspect, scan_size,
        scan_rate, setpoint_c, setpoint_ac, setpoint_f, gain, ...), Parameters as params={...}.
        hold=True: the scan waits after every line until run.next()."""
        body = dict(kind="image", hold=hold, settings=dict(scan=scan, params=params or {}))
        return Run(self, self._call("POST", "/api/runs", json=body)["id"])

    def force_curve(self, params: dict | None = None, **force) -> "Run":
        """Start a force curve: mode ('Contact' | 'AC Mode'), force_dist, velocity, trig_force, trig_ratio, x, y, ..."""
        body = dict(kind="force", settings=dict(force=force, params=params or {}))
        return Run(self, self._call("POST", "/api/runs", json=body)["id"])

    def runs(self) -> list[dict]:
        return self._call("GET", "/api/runs")["runs"]

    def run(self, run_id: str) -> "Run":
        return Run(self, run_id)

    def tune(self, auto: bool = False) -> dict:
        """AC Mode: a frequency sweep (auto=True: also set drive frequency and amplitude).  Arrays as numpy."""
        r = self._call("POST", "/api/tune", json={"kind": "auto" if auto else "single"}, timeout=180)
        t = r["tune"]
        return dict(frequency=_f32(t["f"]), amplitude=_f32(t["amplitude"]), phase=_f32(t["phase"]), f0=t["f0"], q=t["q"],
                    drive_frequency=t["drive_freq"], settings=r["settings"])

    def force_law(self, tip_radius: float = 0.0) -> list[dict]:
        r = self._call("GET", "/api/forcelaw", params={"radius": tip_radius})
        return [dict(m, gap=_f32(m["gap"]), approach=_f32(m["approach"]), retract=_f32(m["retract"])) for m in r["materials"]]

    # ---- sweeps -------------------------------------------------------------------------------------------------
    def sweep(self, x: dict, y: dict | None = None, single: bool = False, **scan) -> dict:
        """A sweep (x) or grid (x and y) of images, run one after another on the server.  An axis:
        {'key': 'gain', 'start': 5, 'stop': 40, 'steps': 4, 'spacing': 'log'} or {'key': 'setpoint_f', 'values': [...]}.
        single=True: all the steps of x within ONE image, in equal blocks of lines (scan order); each line then
        carries its value (Line.sweep) and twin.batch(id)['runs'] has a row per step.
        Keywords: scan settings for all of them.  Returns the sweep; twin.batch(id) follows it."""
        body = dict(x=x, settings=dict(scan=scan))
        if y:
            body["y"] = y
        if single:
            body["single"] = True
        return self._call("POST", "/api/batches", json=body)

    # ---- favourite settings -------------------------------------------------------------------------------------
    def favorites(self) -> list[dict]:
        """Your saved settings: [{name, created, mode, summary}]."""
        return self._call("GET", "/api/favorites")["favorites"]

    def save_favorite(self, name: str, note: str = "") -> list[dict]:
        """Save all the present settings (mode, scan, Parameters, force curve, tune) under a name."""
        return self._call("POST", "/api/favorites", json=dict(name=name, note=note))["favorites"]

    def load_favorite(self, name: str) -> dict:
        """Make a favourite's settings the present ones; returns them."""
        from urllib.parse import quote
        return self._call("POST", f"/api/favorites/{quote(name, safe='')}/load", json={})["settings"]

    def delete_favorite(self, name: str) -> list[dict]:
        from urllib.parse import quote
        return self._call("DELETE", f"/api/favorites/{quote(name, safe='')}")["favorites"]

    def batch(self, batch_id: str) -> dict:
        """A sweep with each run's values, status and metrics."""
        return self._call("GET", f"/api/batches/{batch_id}")

    def wait_batch(self, batch_id: str, poll: float = 5.0) -> dict:
        while True:
            b = self.batch(batch_id)
            if b["ended"] and all(r["status"] not in ("queued", "running") for r in b["runs"]):
                return b
            time.sleep(poll)


class Line:
    """One scan line: line (its row in the image), number (its place in the scan order), data[channel][direction]
    (numpy, one value per pixel), metrics (dict), curves[direction] (FFM, if asked for: dict with deflection and
    zsensor [points, samples] in metres, baseline [points], time [samples], points)."""

    def __init__(self, d: dict, units: dict):
        self.number, self.line, self.setpoint = d["number"], d["line"], d.get("setpoint")
        self.sweep = d.get("sweep")           # a sweep within the image: the value this line was scanned with
        self.metrics = d.get("metrics") or {}
        self.units = units
        self.data = {c: {dd: _arr(v) for dd, v in dirs.items()} for c, dirs in d["data"].items()}
        self.curves = {}
        for dd, c in (d.get("curves") or {}).items():
            if c is None:
                self.curves[dd] = None
                continue
            shape = tuple(c["shape"])
            self.curves[dd] = dict(deflection=_f32(c["deflection"], shape), zsensor=_f32(c["zsensor"], shape),
                                   baseline=_arr(c["baseline"]), start_time=_arr(c["start_time"]),
                                   points=np.array(c["points"]), time=np.arange(shape[1]) / c["sample_rate"],
                                   complete=c["complete"])

    def force(self, direction="Trace", k: float | None = None):
        """FFM: force = k x (deflection - baseline) of every curve of the line [N]; k defaults to the run's."""
        c = self.curves[direction]
        return (c["deflection"] - c["baseline"][:, None]) * (k if k is not None else self._k)

    def __repr__(self):
        return f"<Line {self.line} (#{self.number}) {list(self.data)}>"


class Run:
    """A run on the server."""

    def __init__(self, twin: Twin, run_id: str):
        self.twin, self.id = twin, run_id
        self._k = None

    def state(self) -> dict:
        return self.twin._call("GET", f"/api/runs/{self.id}")

    @property
    def status(self) -> str:
        return self.state()["status"]

    def wait(self, poll: float = 2.0) -> dict:
        """Until the run has ended; returns its state."""
        while True:
            st = self.state()
            if st["status"] not in ("queued", "running"):
                return st
            time.sleep(poll)

    def lines(self, after: int = 0, channels=None, curves: bool = False, points: str = "", curve_step: int = 1,
              follow: bool = True):
        """The scan lines from number `after` on, each as soon as it is scanned (follow=True: until the run ends).
        channels: list of channel names (default all).  curves (FFM): the force curves of each line, for the pixels
        `points` ('start:stop:step' or 'i,j,k'), every `curve_step`-th sample."""
        if self._k is None:
            self._k = self.twin.settings()["params"]["k"]
        n = after
        while True:
            r = self.twin._call("GET", f"/api/runs/{self.id}/lines", timeout=60, params=dict(
                after=n, wait=25 if follow else 0, channels=",".join(channels or []), curves=curves, points=points,
                curve_step=curve_step))
            for d in r["items"]:
                line = Line(d, r["units"])
                line._k = self._k
                n = d["number"] + 1
                yield line
            ended = r["status"] not in ("queued", "running")
            if not follow or (ended and n >= r["lines_done"]):
                return

    def set(self, scan=None, params=None) -> dict:
        """Change settings of this running image (acting at once, from the next line, or restarting it)."""
        return self.twin._call("POST", f"/api/runs/{self.id}/settings", json=dict(scan=scan or {}, params=params or {}))

    def next(self, scan=None, params=None) -> dict:
        """With hold=True: apply these settings (optional) and scan the next line."""
        return self.twin._call("POST", f"/api/runs/{self.id}/next", json=dict(scan=scan or {}, params=params or {}))

    def new_tip(self) -> dict:
        """Fit a fresh tip (the Parameters' radius) at the start of the next line of this running image."""
        return self.twin._call("POST", "/api/new_tip", json={})

    def hold(self, on: bool = True) -> dict:
        return self.twin._call("POST", f"/api/runs/{self.id}/hold", json={"hold": on})

    def stop(self):
        return self.twin._call("POST", f"/api/runs/{self.id}/stop", json={})

    def metrics(self) -> dict:
        return self.twin._call("GET", f"/api/runs/{self.id}/metrics")["metrics"]

    def image(self, channel="Height", direction="Trace") -> np.ndarray:
        """A whole channel [line, point] (NaN where not scanned)."""
        r = self.twin._call("GET", f"/api/runs/{self.id}/image", params=dict(channel=channel, direction=direction))
        return _f32(r["data"], (r["lines"], r["points"])).astype(float)

    def curve(self, line: int, point: int, direction="Trace") -> dict:
        """FFM: the curve at a pixel (time, zsensor, deflection, ...); also during the scan, for lines scanned so far."""
        r = self.twin._call("GET", f"/api/runs/{self.id}/ffm_curve", params=dict(line=line, point=point, direction=direction))
        out = {k.lower(): _f32(v) for k, v in r["data"].items()}
        out.update(r["meta"])
        return out

    def force_curve(self) -> dict:
        """A force-curve run: its columns as numpy arrays."""
        r = self.twin._call("GET", f"/api/runs/{self.id}/curve")
        return {k: _f32(v) for k, v in r["data"].items()}

    def download(self, name: str, path: str | None = None) -> str:
        """images.h5, forcevolume.h5, curve.csv or settings.json, saved to path (default: <run id>_<name>)."""
        path = path or f"{self.id}_{name}"
        r = self.twin._call("GET", f"/api/runs/{self.id}/files/{name}", raw=True, timeout=600)
        with open(path, "wb") as f:
            f.write(r.content)
        return path

    def __repr__(self):
        return f"<Run {self.id}>"