"""Essay datasets on the Modal volume, loaded into plain pandas frames with one column naming scheme.

ELLIPSE (six analytic scores): train csv plus a password-protected test zip (password is in the corpus README).
"""
import os
import zipfile

import pandas as pd

TRAITS = ["cohesion", "syntax", "vocabulary", "phraseology", "grammar", "conventions"]
BR = "[BR]"  # DeBERTa's tokenizer drops newlines; every Kaggle winner put a token in their place


def with_breaks(text):
    return str(text).replace("\r\n", "\n").replace("\n", f" {BR} ")


def ellipse(root="/data/ELLIPSE-Corpus"):
    """Returns (train, test) frames with columns text_id, text, overall, the six traits, prompt."""
    test_csv = os.path.join(root, "ELLIPSE_Final_github_test.csv")
    if not os.path.exists(test_csv):
        with zipfile.ZipFile(os.path.join(root, "ELLIPSE_Final_github_test.zip")) as z:
            names = [n for n in z.namelist() if n.lower().endswith(".csv") and not n.startswith("__MACOSX")]
            with z.open(names[0], pwd=b"ellipse_test") as f, open(test_csv, "wb") as out:
                out.write(f.read())
    frames = []
    for split, name in (("train", "ELLIPSE_Final_github_train.csv"), ("test", "ELLIPSE_Final_github_test.csv")):
        df = pd.read_csv(os.path.join(root, name))
        df.columns = [c.lower() for c in df.columns]
        df = df.rename(columns={"text_id_kaggle": "text_id", "full_text": "text"})
        df = df[["text_id", "text", "overall", *TRAITS, "prompt"]].copy()
        df["split"] = split
        frames.append(df)
    return frames[0], frames[1]


# --- PERSUADE 2.0: argument parts with effectiveness ---------------------------------------------------------------
TYPES = ["Lead", "Position", "Claim", "Counterclaim", "Rebuttal", "Evidence", "Concluding Statement"]
EFFECT = ["Ineffective", "Adequate", "Effective"]
BIO = ["O"] + [f"{p}-{t}" for t in TYPES for p in ("B", "I")]  # 15 labels


def persuade(root="/data/persuade_corpus_2.0", split="train"):
    """One row per essay: essay_id, text, holistic, prompt, parts = [(start, end, type, effectiveness), ...].
    Offsets are checked against the text and repaired with a search when they are off (they often are, slightly)."""
    name = "persuade_2.0_train.csv" if split == "train" else "persuade_corpus_2.0_test.csv"
    cols = ["essay_id", "full_text", "holistic_essay_score", "discourse_start", "discourse_end", "discourse_text",
            "discourse_type", "discourse_effectiveness", "prompt_name", "task"]
    df = pd.read_csv(os.path.join(root, name), usecols=cols)
    rows, bad = [], 0
    for eid, g in df.groupby("essay_id", sort=False):
        text = str(g.full_text.iloc[0])
        parts = []
        for r in g.itertuples(index=False):
            if r.discourse_type not in TYPES:
                continue
            s, e, piece = int(r.discourse_start), int(r.discourse_end), str(r.discourse_text).strip()
            if text[s:e].strip() != piece:
                k = text.find(piece)
                if k < 0:
                    bad += 1
                    continue
                s, e = k, k + len(piece)
            eff = r.discourse_effectiveness if r.discourse_effectiveness in EFFECT else None
            parts.append((s, e, r.discourse_type, eff))
        rows.append({"essay_id": eid, "text": text, "holistic": float(g.holistic_essay_score.iloc[0]),
                     "prompt": g.prompt_name.iloc[0], "task": g.task.iloc[0], "parts": sorted(parts)})
    out = pd.DataFrame(rows)
    out.attrs["dropped_parts"] = bad
    return out