"""One essay in, the coach's reading out: six scores, argument parts with effectiveness, and plain-words notes.

    coach = Coach("/data/runs/scores-v1", "/data/runs/parts-v1")
    coach.read(text)  # dict
"""
import json
import os

import numpy as np
import torch
from safetensors.torch import load_file
from transformers import AutoTokenizer

from coach.data import TRAITS, TYPES, with_breaks
from coach.decode import decode
from coach.model import PartsModel, ScoreModel

TRAIT_WORDS = {
    "cohesion": "how the ideas link up and flow from one to the next",
    "syntax": "sentence structure and variety",
    "vocabulary": "the range of words used",
    "phraseology": "natural phrasing and word combinations",
    "grammar": "grammatical accuracy",
    "conventions": "spelling, punctuation and capitalisation",
}


class Coach:
    def __init__(self, scores_dir, parts_dir, device=None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        sc = json.load(open(os.path.join(scores_dir, "config.json")))
        self.score_tok = AutoTokenizer.from_pretrained(os.path.join(scores_dir, "tokenizer"))
        self.score_net = ScoreModel(sc["backbone"], n_tokens=len(self.score_tok)).to(self.device).eval()
        self.score_net.load_state_dict(load_file(os.path.join(scores_dir, "model.safetensors")))
        self.score_len = sc["max_len"]
        pc = json.load(open(os.path.join(parts_dir, "config.json")))
        self.parts_tok = AutoTokenizer.from_pretrained(os.path.join(parts_dir, "tokenizer"))
        self.parts_net = PartsModel(pc["backbone"], n_tokens=len(self.parts_tok)).to(self.device).eval()
        self.parts_net.load_state_dict(load_file(os.path.join(parts_dir, "model.safetensors")))
        self.parts_len = pc["max_len"]

    @torch.no_grad()
    def scores(self, text):
        enc = self.score_tok(with_breaks(text), truncation=True, max_length=self.score_len, return_tensors="pt").to(self.device)
        out = self.score_net(enc["input_ids"], enc["attention_mask"])[0].float().cpu().numpy()
        return {t: float(np.clip(round(v * 2) / 2, 1.0, 5.0)) for t, v in zip(TRAITS, out)}, {t: float(v) for t, v in zip(TRAITS, out)}

    @torch.no_grad()
    def parts(self, text):
        enc = self.parts_tok(text, truncation=True, max_length=self.parts_len, return_offsets_mapping=True, return_tensors="pt")
        offs = [None if (s == e == 0) else (int(s), int(e)) for s, e in enc.pop("offset_mapping")[0].tolist()]
        offs[0] = None
        enc = enc.to(self.device)
        lb, le = self.parts_net(enc["input_ids"], enc["attention_mask"])
        return decode(lb[0].float().cpu().numpy(), le[0].float().cpu().numpy(), offs, text)

    def read(self, text):
        text = text.replace("\r\n", "\n").strip()
        snapped, raw = self.scores(text)
        parts = self.parts(text)
        return {"scores": snapped, "scores_raw": raw, "parts": parts, "notes": notes(snapped, parts, text)}


def notes(scores, parts, text):
    """Plain-words feedback from the numbers: the two weakest skills, and what is missing or weak in the argument."""
    out = []
    weakest = sorted(scores, key=scores.get)[:2]
    for t in weakest:
        out.append({"kind": "skill", "trait": t, "score": scores[t], "text": f"Weakest area: {TRAIT_WORDS[t]} ({scores[t]:.1f} of 5)."})
    have = {p["type"] for p in parts}
    if "Position" not in have:
        out.append({"kind": "structure", "text": "No clear position: say early on which side you take, in one sentence."})
    if "Lead" not in have:
        out.append({"kind": "structure", "text": "No lead: the first lines should pull the reader in and point at the question."})
    if "Counterclaim" not in have:
        out.append({"kind": "structure", "text": "No counterclaim: name the strongest view against yours, then answer it. Examiners look for this."})
    elif "Rebuttal" not in have:
        out.append({"kind": "structure", "text": "You raise the other side but never answer it: add a rebuttal after the counterclaim."})
    if "Concluding Statement" not in have:
        out.append({"kind": "structure", "text": "No conclusion: finish by restating your position and the main reason for it."})
    claims = [p for p in parts if p["type"] == "Claim"]
    for c in claims:
        nxt = [p for p in parts if p["start"] > c["start"]]
        if not nxt or nxt[0]["type"] != "Evidence":
            out.append({"kind": "part", "start": c["start"], "end": c["end"],
                        "text": "This claim has no evidence right after it: add an example, a fact or a reason."})
    for p in parts:
        if p["effectiveness"] == "Ineffective":
            out.append({"kind": "part", "start": p["start"], "end": p["end"],
                        "text": f"This {p['type'].lower()} reads as ineffective: make it specific and tie it to your position."})
    words = len(text.split())
    if words < 250:
        out.append({"kind": "length", "text": f"{words} words. IELTS Task 2 needs at least 250; short essays lose marks before they are read."})
    return out