"""Fit the IELTS band estimate on top of the two trained models, on the ~1.2k human-banded IELTS essays.

    MODAL_PROFILE=labs3 modal run scripts/band.py --scores scores-v1 --parts parts-v1

Features per essay: the six raw scores, word count, number of parts of each type, share of effective parts. Model:
ridge regression, 5-fold cross-validation, reported as exact-band and within-half-band accuracy and RMSE in bands.
The fitted weights go to /data/runs/band-v1 (band.json) so inference is a dot product.
"""
import modal

app = modal.App("essay-coach-band")
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/cu124")
         .pip_install("transformers==4.46.3", "sentencepiece", "protobuf", "tiktoken", "safetensors", "pandas", "numpy<2.3",
                      "scikit-learn", "huggingface_hub")
         .add_local_python_source("coach"))


@app.function(image=image, gpu="A10G", volumes={"/data": vol, "/root/.cache/huggingface": cache}, timeout=3600)
def fit(scores: str, parts: str, tag: str):
    import json
    import os

    import numpy as np
    import pandas as pd
    from sklearn.linear_model import Ridge
    from sklearn.model_selection import KFold

    from coach.data import TRAITS, TYPES
    from coach.infer import Coach

    df = pd.read_csv("/data/ielts-writing-scored-essays-dataset/ielts_writing_dataset.csv")
    df = df[df.Overall.notna() & df.Essay.notna()].copy()
    df["Task_Type"] = df.Task_Type.astype(int)
    print(f"IELTS essays with a band: {len(df)} (task 2: {(df.Task_Type == 2).sum()})", flush=True)
    coach = Coach(f"/data/runs/{scores}", f"/data/runs/{parts}")
    feats = []
    for i, r in enumerate(df.itertuples(index=False)):
        read = coach.read(str(r.Essay))
        raw = read["scores_raw"]
        ps = read["parts"]
        f = [raw[t] for t in TRAITS] + [np.log1p(len(str(r.Essay).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(r.Task_Type == 2)]
        feats.append(f)
        if i % 200 == 0:
            print(i, flush=True)
    X, y = np.array(feats, dtype=np.float32), df.Overall.values.astype(np.float32)
    names = TRAITS + ["log_words"] + [f"n_{t}" for t in TYPES] + ["share_effective", "task2"]
    mu, sd = X.mean(0), X.std(0) + 1e-6
    Xn = (X - mu) / sd
    pred = np.zeros_like(y)
    for trn, val in KFold(5, shuffle=True, random_state=0).split(Xn):
        m = Ridge(alpha=3.0).fit(Xn[trn], y[trn])
        pred[val] = m.predict(Xn[val])
    snapped = np.clip(np.round(pred * 2) / 2, 1, 9)
    res = {"n": int(len(y)), "rmse_bands": float(np.sqrt(((pred - y) ** 2).mean())),
           "exact": float((snapped == y).mean()), "within_half": float((np.abs(snapped - y) <= 0.5).mean()),
           "within_one": float((np.abs(snapped - y) <= 1.0).mean()), "baseline_rmse_mean": float(y.std()),
           "corr": float(np.corrcoef(pred, y)[0, 1])}
    print(res, flush=True)
    m = Ridge(alpha=3.0).fit(Xn, y)
    run = f"/data/runs/{tag}"
    os.makedirs(run, exist_ok=True)
    json.dump({"features": names, "mean": mu.tolist(), "std": sd.tolist(), "coef": m.coef_.tolist(),
               "intercept": float(m.intercept_), "cv": res, "scores_run": scores, "parts_run": parts}, open(f"{run}/band.json", "w"), indent=1)
    vol.commit()
    return res


@app.local_entrypoint()
def main(scores: str = "scores-v1", parts: str = "parts-v1", tag: str = "band-v1"):
    print(fit.remote(scores, parts, tag))