"""Build the SimCSE training corpus and a labelled eval set from FW_cleaned_data_new_base_15.csv.

Only Make + Model are used, as requested.

Outputs (written to simcse/data/):
  train.jsonl  - {"sentence1": s, "sentence2": s} unsupervised SimCSE pairs
  eval.jsonl   - {"sentence1", "sentence2", "score"} labelled pairs for EmbeddingSimilarityEvaluator
  groups.json  - canonical string -> list of raw surface forms (kept for inspection / later use)

The labels come from the canonical columns already present in the CSV: two raw
"<Make> <Model>" strings that normalize to the same "<MAKE> <NEW_MODELS>" are a
positive pair (score 1.0); strings from different canonical groups are negatives
(score 0.0). Nothing but Make/Model text ever reaches the model -- the canonical
columns are used only to build labels.
"""

import json
import random
import re
from collections import defaultdict
from pathlib import Path

import pandas as pd

CSV = Path(__file__).resolve().parents[1] / "data" / "FW_cleaned_data_new_base_15.csv"
OUT = Path(__file__).resolve().parent / "data"
SEED = 42
N_NEG = 2000  # negative eval pairs; positives come from the data (~1.3k)


def norm(s: str) -> str:
    """Lowercase, strip punctuation to spaces, collapse whitespace."""
    s = str(s).lower().strip()
    s = re.sub(r"[^a-z0-9]+", " ", s)
    return re.sub(r"\s+", " ", s).strip()


def dedupe_make(make: str, model: str) -> str:
    """Join make+model, dropping a make token repeated at the head of the model.

    The raw data contains rows like Make='audi', Model='audi a3'.
    """
    make_t, model_t = norm(make).split(), norm(model).split()
    i = 0
    while i < len(make_t) and i < len(model_t) and make_t[i] == model_t[i]:
        i += 1
    # only strip if the whole make prefix was echoed and something survives
    if i == len(make_t) and len(model_t) > i:
        model_t = model_t[i:]
    return " ".join(make_t + model_t)


def main():
    rng = random.Random(SEED)
    OUT.mkdir(parents=True, exist_ok=True)

    df = pd.read_csv(CSV, low_memory=False)
    df = df.dropna(subset=["Make", "Model", "MAKE", "NEW_MODELS"])

    raw = (df["Make"].map(norm) + " " + df["Model"].map(norm)).str.strip()
    canon = df.apply(lambda r: dedupe_make(r["MAKE"], r["NEW_MODELS"]), axis=1)

    pairs = pd.DataFrame({"raw": raw, "canon": canon}).drop_duplicates()
    pairs = pairs[(pairs["raw"] != "") & (pairs["canon"] != "")]

    groups = defaultdict(set)
    for r, c in pairs.itertuples(index=False):
        groups[c].add(r)
    groups = {c: sorted(v) for c, v in groups.items()}

    # ---- training corpus -------------------------------------------------
    # One sentence per canonical group plus every distinct surface form. Using
    # the canonical string too lets the model see the clean form it should map
    # everything toward. Duplicates are dropped because SimCSE's in-batch
    # negatives would otherwise treat identical strings as negatives.
    corpus = sorted({s for forms in groups.values() for s in forms} | set(groups))
    rng.shuffle(corpus)

    with (OUT / "train.jsonl").open("w") as f:
        for s in corpus:
            f.write(json.dumps({"sentence1": s, "sentence2": s}) + "\n")

    # ---- labelled eval set ----------------------------------------------
    positives = []
    for c, forms in groups.items():
        variants = sorted(set(forms) | {c})
        for i in range(len(variants)):
            for j in range(i + 1, len(variants)):
                positives.append((variants[i], variants[j], 1.0))

    canon_list = sorted(groups)
    negatives, seen = [], set()
    while len(negatives) < N_NEG:
        a, b = rng.sample(canon_list, 2)
        sa = rng.choice(sorted(set(groups[a]) | {a}))
        sb = rng.choice(sorted(set(groups[b]) | {b}))
        key = tuple(sorted((sa, sb)))
        if sa == sb or key in seen:
            continue
        seen.add(key)
        negatives.append((sa, sb, 0.0))

    eval_rows = positives + negatives
    rng.shuffle(eval_rows)
    with (OUT / "eval.jsonl").open("w") as f:
        for s1, s2, score in eval_rows:
            f.write(json.dumps({"sentence1": s1, "sentence2": s2, "score": score}) + "\n")

    with (OUT / "groups.json").open("w") as f:
        json.dump(groups, f, indent=1, sort_keys=True)

    print(f"canonical groups : {len(groups)}")
    print(f"train sentences  : {len(corpus)}  -> {OUT / 'train.jsonl'}")
    print(f"eval pairs       : {len(eval_rows)} ({len(positives)} pos / {len(negatives)} neg)"
          f"  -> {OUT / 'eval.jsonl'}")
    print("\nsample train:", corpus[:5])
    print("sample eval :", eval_rows[:3])


if __name__ == "__main__":
    main()