"""Stages: manifest (integrity), prep (normalised views + legal forms), splits."""
import random
from collections import Counter
from multiprocessing import Pool

import numpy as np
import pandas as pd
import pyarrow.parquet as pq

import text_utils as tu
from common import get_logger, load_json, save_json, save_npz, save_parquet, n_workers, Progress
from config import CFG, EMB, EMB_MODELS, PATHS, PREP, WORK
from embeddings import emb_ids, load_emb

log = get_logger("prep")
RAW_COLS = ["entity_id", "business_name", "business_address", "country"]
SIDES = [("train", "s1"), ("train", "pool"), ("test", "s1"), ("test", "pool")]


def read_raw(split, side):
    files = [PATHS["data"][split][1]] if side == "s1" else [PATHS["data"][split][2], PATHS["data"][split][3]]
    df = pd.concat([pq.read_table(f, columns=RAW_COLS).to_pandas() for f in files], ignore_index=True)
    for c in RAW_COLS:
        df[c] = df[c].fillna("").astype(str)
    return df


def prep_path(split, side):
    return PREP / f"{split}_{side}.parquet"


def load_prep(split, side, columns=None):
    return pd.read_parquet(prep_path(split, side), columns=columns)


# ----------------------------------------------------------------------------------------------- manifest

def run_manifest():
    rep, problems = {}, []
    ids = {}
    for split, side in SIDES:
        df = read_raw(split, side)
        n = len(df)
        dup = int(df.entity_id.duplicated().sum())
        pref = df.entity_id.str[:3]
        ok_pref = pref.isin(["S1-"] if side == "s1" else ["S2-", "S3-"]).all()
        rep[f"{split}_{side}"] = {"rows": n, "dup_ids": dup, "prefix_ok": bool(ok_pref),
                                  "countries": df.country.value_counts().to_dict(),
                                  "empty_name": int((df.business_name == "").sum()),
                                  "empty_addr": int((df.business_address == "").sum())}
        if dup or not ok_pref:
            problems.append(f"{split}_{side}: dup={dup} prefix_ok={ok_pref}")
        ids[(split, side)] = df.entity_id.values
        for m in EMB_MODELS:
            e = load_emb(m, split, side)
            r = rep[f"{split}_{side}"].setdefault("emb", {})
            r[m] = {"rows": len(e), "dim": e.dim}
            if len(e) != n:
                problems.append(f"{m} {split}_{side}: {len(e)} emb rows vs {n} data rows")
            eids = emb_ids(m, split, side)
            if eids is not None and not np.array_equal(np.asarray(eids, dtype=object), df.entity_id.values):
                problems.append(f"{m} {split}_{side}: id order differs from parquet")
            samp = e.take(np.random.default_rng(0).choice(len(e), 2000, replace=False)).astype(np.float32)
            nr = np.linalg.norm(samp, axis=1)
            r[m]["norm_min"], r[m]["norm_max"] = float(nr.min()), float(nr.max())
            if abs(nr.mean() - 1) > 0.01:
                problems.append(f"{m} {split}_{side}: vectors not unit-norm (mean {nr.mean():.3f})")
        log.info(f"checked {split}_{side}: {rep[f'{split}_{side}']}")
    # ground truth
    gt = pd.read_csv(PATHS["ground_truth"], sep="\t", dtype=str, keep_default_na=False)
    s1ids = set(ids[("train", "s1")])
    pool_ids = set(ids[("train", "pool")])
    matched = [x for m in gt.matched_entity_ids if m for x in m.split(",")]
    cnt = Counter(matched)
    rep["ground_truth"] = {"rows": len(gt), "s1_covered": int(gt.source1_entity_id.isin(s1ids).sum()),
                           "n_matched": len(matched), "unknown_ids": int(sum(1 for x in cnt if x not in pool_ids)),
                           "multi_owner": int(sum(1 for v in cnt.values() if v > 1))}
    if rep["ground_truth"]["unknown_ids"]:
        problems.append("ground truth references ids missing from train pool")
    if rep["ground_truth"]["multi_owner"]:
        problems.append("S2/S3 ids owned by >1 S1: one-owner arbitration invalid, folds need label-graph components")
    # embedding semantic check on train: true pairs must score far above random pairs
    s1_index = {x: i for i, x in enumerate(ids[("train", "s1")])}
    pool_index = {x: i for i, x in enumerate(ids[("train", "pool")])}
    sub = gt[gt.matched_entity_ids != ""].sample(3000, random_state=0)
    a = np.array([s1_index[s] for s in sub.source1_entity_id])
    b = np.array([pool_index[m.split(",")[0]] for m in sub.matched_entity_ids])
    for m in EMB_MODELS:
        es, ep = load_emb(m, "train", "s1"), load_emb(m, "train", "pool")
        qa, pb = es.take(a).astype(np.float32), ep.take(b).astype(np.float32)
        pr = ep.take(np.random.default_rng(1).choice(len(ep), len(b))).astype(np.float32)
        tpos, trnd = float((qa * pb).sum(1).mean()), float((qa * pr).sum(1).mean())
        rep["ground_truth"][f"{m}_cos_true_vs_random"] = [tpos, trnd]
        if tpos - trnd < 0.1:
            problems.append(f"{m}: true pairs not separated from random (cos {tpos:.3f} vs {trnd:.3f})")
    rep["problems"] = problems
    save_json(rep, WORK / "reports" / "manifest.json")
    if problems:
        raise RuntimeError("integrity manifest failed: " + "; ".join(problems))
    log.info("manifest OK")


# ----------------------------------------------------------------------------------------------- prep

_FORMS = None


def _init(forms):
    global _FORMS
    _FORMS = forms


def _tok_chunk(names):
    return [tu.name_tokens(x) for x in names]


def _proc_chunk(args):
    names, addrs, countries = args
    rows = []
    for n, a, c in zip(names, addrs, countries):
        forms, long_single = _FORMS.get(c, ([], set()))
        r = tu.process_name(n, forms, long_single)
        r.update(tu.parse_address(a))
        rows.append(r)
    return pd.DataFrame(rows)


def norm_country(s):
    return s.strip().lower() or "unknown"


def run_prep():
    PREP.mkdir(parents=True, exist_ok=True)
    forms_path = PREP / "legal_forms.json"
    raws = {}
    with Pool(n_workers()) as pool:
        if not forms_path.exists():
            by_c = {}
            for split, side in SIDES:
                df = read_raw(split, side)
                df["cn"] = df.country.map(norm_country)
                raws[(split, side)] = df
                for c, g in df.groupby("cn"):
                    by_c.setdefault(c, []).extend(g.business_name.tolist())
            rng = random.Random(CFG["seed"])
            forms = {}
            for c, names in by_c.items():
                if len(names) > CFG["legal_sample"]:
                    names = rng.sample(names, CFG["legal_sample"])
                chunks = [names[i:i + 50_000] for i in range(0, len(names), 50_000)]
                toks = [t for part in pool.map(_tok_chunk, chunks) for t in part]
                f = tu.discover_legal_forms(toks, CFG["legal_min_frac"], CFG["legal_boundary_ratio"],
                                            CFG["legal_max_cooc"], CFG["legal_bigram_tight"])
                forms[c] = [" ".join(x) for x in f]
                log.info(f"legal forms [{c}] from {len(names)} names: {forms[c]}")
            save_json(forms, forms_path)
    forms = load_json(forms_path)
    forms_t = {c: ([tuple(x.split()) for x in v], {x for x in v if " " not in x and len(x) >= 3})
               for c, v in forms.items()}
    with Pool(n_workers(), initializer=_init, initargs=(forms_t,)) as pool:
        for split, side in SIDES:
            out = prep_path(split, side)
            if out.exists():
                continue
            df = raws.pop((split, side), None)
            if df is None:
                df = read_raw(split, side)
                df["cn"] = df.country.map(norm_country)
            step = 100_000
            args = [(df.business_name.values[i:i + step], df.business_address.values[i:i + step],
                     df.cn.values[i:i + step]) for i in range(0, len(df), step)]
            prog = Progress(log, len(args), f"prep {split}_{side}")
            parts = []
            for part in pool.imap(_proc_chunk, args):
                parts.append(part)
                prog.step()
            feat = pd.concat(parts, ignore_index=True)
            res = pd.DataFrame({"id": df.entity_id.values, "country": df.cn.values,
                                "src": df.entity_id.str[1].astype(np.int8).values,
                                "raw_name": df.business_name.values, "raw_addr": df.business_address.values})
            res = pd.concat([res, feat], axis=1)
            save_parquet(res, out)
            log.info(f"wrote {out} rows={len(res)}")


# ----------------------------------------------------------------------------------------------- splits

def run_splits():
    s1 = load_prep("train", "s1", ["id", "country"])
    pool_ids = load_prep("train", "pool", ["id"]).id.values
    rng = np.random.default_rng(CFG["seed"])
    perm = rng.permutation(len(s1))
    a, b, c = CFG["n_ce"], CFG["n_stack"], CFG["n_val"]
    groups = {"ce": np.sort(perm[:a]), "stack": np.sort(perm[a:a + b]), "val": np.sort(perm[a + b:a + b + c])}
    q_train = np.sort(np.concatenate(list(groups.values())))
    n_test = len(load_prep("test", "s1", ["id"]))
    q_test = np.arange(n_test)
    if CFG["test_limit"]:
        q_test = np.sort(rng.choice(n_test, CFG["test_limit"], replace=False))
    # ground truth as index pairs
    gt = pd.read_csv(PATHS["ground_truth"], sep="\t", dtype=str, keep_default_na=False)
    s1_index = pd.Series(np.arange(len(s1)), index=s1.id.values)
    pool_index = pd.Series(np.arange(len(pool_ids)), index=pool_ids)
    gt = gt[gt.matched_entity_ids != ""]
    ex = gt.assign(m=gt.matched_entity_ids.str.split(",")).explode("m")
    gq = s1_index.loc[ex.source1_entity_id.values].values.astype(np.int64)
    gc = pool_index.loc[ex.m.values].values.astype(np.int64)
    n_true = np.bincount(gq, minlength=len(s1)).astype(np.int16)
    save_npz(WORK / "splits.npz", q_train=q_train, q_test=q_test, n_true=n_true, gt_q=gq, gt_c=gc,
             **{f"g_{k}": v for k, v in groups.items()})
    rep = {k: {"n": int(len(v)), "countries": pd.Series(s1.country.values[v]).value_counts().to_dict(),
               "singleton_rate": float((n_true[v] == 0).mean())} for k, v in groups.items()}
    rep["q_test"] = int(len(q_test))
    save_json(rep, WORK / "reports" / "splits.json")
    log.info(f"splits: {rep}")


def load_splits():
    z = np.load(WORK / "splits.npz")
    return {k: z[k] for k in z.files}