"""Turn token predictions into argument parts: (start_char, end_char, type, effectiveness, confidence)."""
import numpy as np

from coach.data import BIO, EFFECT, TYPES

ONCE = {"Lead", "Position", "Concluding Statement"}  # at most one of each per essay (every Kaggle winner's rule)
MIN_TOKENS = {"Lead": 6, "Position": 3, "Claim": 2, "Counterclaim": 4, "Rebuttal": 3, "Evidence": 8, "Concluding Statement": 6}


def decode(bio_logits, eff_logits, offsets, text):
    """bio_logits (L,15), eff_logits (L,3), offsets list of (s,e) per token (None for specials)."""
    p = np.exp(bio_logits - bio_logits.max(-1, keepdims=True))
    p /= p.sum(-1, keepdims=True)
    lab = p.argmax(-1)
    spans, cur = [], None
    for i, l in enumerate(lab):
        if offsets[i] is None:
            continue
        tag = BIO[l]
        if tag == "O":
            if cur: spans.append(cur); cur = None
            continue
        pre, typ = tag.split("-", 1)
        if cur and cur["type"] == typ and pre == "I":
            cur["toks"].append(i)
        else:  # a B tag, or an I of a different type: start a new part (repairs spans that begin with I)
            if cur: spans.append(cur)
            cur = {"type": typ, "toks": [i]}
    if cur: spans.append(cur)
    out = []
    for s in spans:
        if len(s["toks"]) < MIN_TOKENS[s["type"]]:
            continue
        ti = s["toks"]
        conf = float(np.mean([p[i, BIO.index(f"B-{s['type']}")] + p[i, BIO.index(f"I-{s['type']}")] for i in ti]))
        e = eff_logits[ti].mean(0)
        e = np.exp(e - e.max()); e /= e.sum()
        out.append({"start": offsets[ti[0]][0], "end": offsets[ti[-1]][1], "type": s["type"], "conf": conf,
                    "effectiveness": EFFECT[int(e.argmax())], "eff_probs": e.round(3).tolist()})
    for typ in ONCE:  # keep the most confident one
        same = [o for o in out if o["type"] == typ]
        for o in sorted(same, key=lambda o: -o["conf"])[1:]:
            out.remove(o)
    for o in out:  # trim the whitespace the tokenizer attaches to the first word of a span
        while o["start"] < o["end"] and text[o["start"]].isspace():
            o["start"] += 1
        while o["end"] > o["start"] and text[o["end"] - 1].isspace():
            o["end"] -= 1
        o["text"] = text[o["start"]:o["end"]]
    return sorted(out, key=lambda o: o["start"])


def span_f1(pred, gold):
    """The Feedback Prize metric: a predicted part matches a gold one of the same type if their word overlap is >= 0.5
    both ways. pred/gold: lists of (start, end, type) in chars; we compare on word index sets."""
    def words(text_spans):
        return text_spans

    def match(a, b):
        inter = max(0, min(a[1], b[1]) - max(a[0], b[0]))
        return inter / max(1, a[1] - a[0]) >= 0.5 and inter / max(1, b[1] - b[0]) >= 0.5

    f1s = []
    for typ in TYPES:
        P = [x for x in pred if x[2] == typ]
        G = [x for x in gold if x[2] == typ]
        used, tp = set(), 0
        for x in P:
            for j, y in enumerate(G):
                if j not in used and match(x, y):
                    used.add(j); tp += 1; break
        fp, fn = len(P) - tp, len(G) - tp
        if tp + fp + fn:
            f1s.append(tp / (tp + 0.5 * (fp + fn)))
    return float(np.mean(f1s)) if f1s else 0.0