"""The coach as a web endpoint on Modal (CPU, scales to zero, so idle costs nothing).

    MODAL_PROFILE=labs3 modal deploy scripts/serve.py     # prints the URL; the page POSTs {"text": ...} to <url>/read
    MODAL_PROFILE=labs3 modal serve scripts/serve.py      # hot-reloading dev copy

Response: {"scores": {...six 1-5...}, "scores_raw": {...}, "parts": [{start,end,type,effectiveness,eff_probs,conf,text}],
"notes": [{kind,text,...}], "band": {"value": 6.5, "note": "..."} or null, "ms": 1234}.
"""
import modal

app = modal.App("essay-coach")
vol = modal.Volume.from_name("essay-data")
cache = modal.Volume.from_name("essay-hf-cache", create_if_missing=True)
image = (modal.Image.debian_slim(python_version="3.11")
         .pip_install("torch==2.4.1", index_url="https://download.pytorch.org/whl/cpu")
         .pip_install("transformers==4.46.3", "sentencepiece", "protobuf", "tiktoken", "safetensors", "pandas", "numpy<2.3",
                      "fastapi[standard]")
         .add_local_python_source("coach"))
SCORES, PARTS, BAND = "scores-v1", "parts-v1", "band-v1"


@app.cls(image=image, volumes={"/data": vol, "/root/.cache/huggingface": cache}, cpu=4, memory=8192,
         scaledown_window=300, timeout=120)
@modal.concurrent(max_inputs=4)
class Server:
    @modal.enter()
    def load(self):
        import json
        import os

        import torch

        from coach.infer import Coach
        torch.set_num_threads(4)
        self.coach = Coach(f"/data/runs/{SCORES}", f"/data/runs/{PARTS}", device="cpu")
        p = f"/data/runs/{BAND}/band.json"
        self.band = json.load(open(p)) if os.path.exists(p) else None

    def band_of(self, read, text, task2=True):
        import numpy as np

        from coach.data import TRAITS, TYPES
        if not self.band:
            return None
        ps = read["parts"]
        f = [read["scores_raw"][t] for t in TRAITS] + [np.log1p(len(text.split()))] + [sum(p["type"] == t for p in ps) for t in TYPES]
        f += [np.mean([p["effectiveness"] == "Effective" for p in ps]) if ps else 0.0, float(task2)]
        x = (np.array(f) - np.array(self.band["mean"])) / np.array(self.band["std"])
        v = float(x @ np.array(self.band["coef"]) + self.band["intercept"])
        cv = self.band["cv"]
        return {"value": float(np.clip(round(v * 2) / 2, 1, 9)), "raw": v,
                "note": f"An estimate, not a grade: on {cv['n']} human-banded essays it lands within half a band {cv['within_half']*100:.0f}% of the time."}

    @modal.fastapi_endpoint(method="POST", docs=True)
    def read(self, body: dict):
        import time
        t0 = time.time()
        text = str(body.get("text", ""))[:12000]
        if len(text.split()) < 20:
            return {"error": "Paste at least a few sentences."}
        out = self.coach.read(text)
        out["band"] = self.band_of(out, text, task2=bool(body.get("task2", True)))
        out["ms"] = int((time.time() - t0) * 1000)
        return out