"""Loading, S1 augmentation with ground truth, sanity checks and splits.

Unchanged from model-2 except: load_data() takes an explicit paths dict so the
same code serves train and test, and test loading needs no ground truth.
"""
import os
import numpy as np
import pandas as pd

S1_COLS = ["entity_id", "business_name", "business_address", "country", "country_norm",
           "name_norm", "name_core", "legal_suffix", "address_abbrev_norm", "address_numbers"]
S23_COLS = S1_COLS


def read_any(path, columns=None):
    if path.endswith(".parquet"):
        df = pd.read_parquet(path, columns=columns)
    else:
        df = pd.read_csv(path, sep="\t", dtype=str, keep_default_na=False, usecols=columns)
    for c in df.columns:
        df[c] = df[c].fillna("").astype(str)
    return df


def split_ids(s):
    return [x for x in str(s).split(",") if x]


def load_data(paths, with_gt=True):
    s1 = read_any(paths["source1"], S1_COLS)
    s2 = read_any(paths["source2"], S23_COLS)
    s3 = read_any(paths["source3"], S23_COLS)
    s23 = pd.concat([s2, s3], ignore_index=True)
    if not with_gt:
        return s1, s23, None
    gt = read_any(paths["ground_truth"])
    return s1, s23, gt


def augment_s1(s1, gt, out_path=None):
    """Attach ground-truth labels to Source 1 as column `matched_entity_ids`."""
    assert s1["entity_id"].is_unique, "duplicate entity_id in source1"
    assert gt["source1_entity_id"].is_unique, "duplicate source1_entity_id in ground truth"
    s1 = s1.merge(gt.rename(columns={"source1_entity_id": "entity_id"}), on="entity_id",
                  how="left", validate="one_to_one")
    missing = (s1["matched_entity_ids"].isna()).sum()
    assert missing == 0, f"{missing} source1 entities have no ground-truth row"
    s1["matched_entity_ids"] = s1["matched_entity_ids"].fillna("").astype(str)
    if out_path:
        s1.to_parquet(out_path, index=False)
    return s1


def sanity_checks(s1, s23, rng, sample=200_000):
    rep = {}
    rep["n_s1"] = len(s1)
    rep["n_s23"] = len(s23)
    rep["s23_ids_unique"] = bool(s23["entity_id"].is_unique)
    lists = s1["matched_entity_ids"].map(split_ids)
    n = lists.map(len)
    rep["singleton_rate"] = float((n == 0).mean())
    rep["mean_matches"] = float(n.mean())
    rep["max_matches"] = int(n.max())
    flat = pd.Series([x for l in lists for x in l])
    rep["n_matched_ids"] = int(len(flat))
    rep["multi_owner_ids"] = int(flat.duplicated().sum())
    idx = pd.Index(s23["entity_id"])
    pos = idx.get_indexer(flat)
    rep["gt_ids_missing_in_s23"] = int((pos < 0).sum())
    # same-country rate over a sample of true pairs
    owner = np.repeat(np.arange(len(s1)), n.values)
    take = rng.choice(len(flat), size=min(sample, len(flat)), replace=False)
    ok = pos[take] >= 0
    q_c = s1["country_norm"].values[owner[take][ok]]
    c_c = s23["country_norm"].values[pos[take][ok]]
    rep["true_pairs_same_country_rate"] = float((q_c == c_c).mean()) if ok.any() else None
    rep["s1_country_counts"] = s1["country_norm"].value_counts().to_dict()
    return rep


def make_split(s1, cfg, rng):
    """Return three disjoint arrays of S1 row positions: ce, stack, val."""
    n_ce, n_st, n_va = cfg["n_s1_ce"], cfg["n_s1_stack"], cfg["n_s1_val"]
    if cfg["split"] == "A":
        perm = rng.permutation(len(s1))
        need = n_ce + n_st + n_va
        if need > len(perm):
            raise ValueError(f"requested {need} S1 entities but only {len(perm)} exist")
        return perm[:n_ce], perm[n_ce:n_ce + n_st], perm[n_ce + n_st:need]
    if cfg["split"] == "B":
        c = s1["country_norm"].values
        tr = rng.permutation(np.where(c == cfg["split_b_train_country"])[0])
        va = rng.permutation(np.where(c == cfg["split_b_val_country"])[0])
        if n_ce + n_st > len(tr) or n_va > len(va):
            raise ValueError("not enough entities for split B with these sizes")
        return tr[:n_ce], tr[n_ce:n_ce + n_st], va[:n_va]
    raise ValueError(cfg["split"])


def ensure_dir(p):
    os.makedirs(p, exist_ok=True)
    return p