"""Deduplicate before stratified splitting; ambiguous labels are removed."""
import html
import re
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import pandas as pd
from sklearn.model_selection import train_test_split
from common.io import ROOT, dump, sha256

def normalize(text):
    text = html.unescape(str(text)).lower()
    text = re.sub(r"https?://\S+|www\.\S+", " URL ", text)
    text = re.sub(r"@\w+", " USER ", text)
    return " ".join(text.split())

def main():
    raw = ROOT / "data/raw/disaster_tweets.csv"
    df = pd.read_csv(raw, keep_default_na=False)
    assert len(df) == 7613 and set(df.target.unique()) == {0, 1}
    df["group"] = df.text.map(normalize)
    conflicting = df.groupby("group").target.nunique()
    bad = set(conflicting[conflicting > 1].index)
    ambiguous_rows = int(df.group.isin(bad).sum())
    clean = df[~df.group.isin(bad)].drop_duplicates("group").copy()
    train, rest = train_test_split(clean, test_size=0.3, stratify=clean.target, random_state=42)
    val, test = train_test_split(rest, test_size=0.5, stratify=rest.target, random_state=42)
    out = ROOT / "data/splits"
    out.mkdir(exist_ok=True)
    stats = {}
    groups = []
    for name, part in [("train", train), ("validation", val), ("test", test)]:
        part = part.sort_values("id")
        part.to_csv(out / f"{name}.csv", index=False, encoding="utf-8")
        stats[name] = dict(n=len(part), positive=int(part.target.sum()), negative=int((part.target == 0).sum()), sha256=sha256(out / f"{name}.csv"))
        groups.append(set(part.group))
    assert not (groups[0] & groups[1] or groups[0] & groups[2] or groups[1] & groups[2])
    dump(ROOT / "results/data_audit.json", dict(raw_n=len(df), clean_n=len(clean), conflicting_groups=len(bad),
         removed_ambiguous_rows=ambiguous_rows, removed_duplicate_rows=len(df)-ambiguous_rows-len(clean),
         seed=42, splits=stats, rule="HTML unescape, lowercase, URL/user normalization, whitespace; remove conflicts then deduplicate; 70/15/15 stratified"))
    print(stats)

if __name__ == "__main__":
    main()