#!/usr/bin/env python3
"""
Merge volumes and editions of the same work into one entry, and link translations.

Usage:
    python group.py        # last step: after build_overview, years, geolocate and keywords

Within one library and language, entries are grouped when they have the same first author's
surname and
- the same title once volume markers are removed ("— Volume 2", ", Vol. II", "Bd. 3", "Part 4"),
- or nearly the same title (at least 80% shared content words, titles of 4+ words),
- or Wikidata says they are editions of the same work (field "wk" from years.py).
A group becomes one entry: the complete edition or first volume leads; the others are listed
under "ed" with their own reading links. Categories, places, keywords and libraries are
merged; the year is the earliest found. Entries in different languages that Wikidata ties to
the same work are linked as translations ("tr").
"""
import argparse
import json
import re
import sys
from collections import defaultdict

from years import VOLUME, norm, surname, words

ROMAN = {"i": 1, "v": 5, "x": 10, "l": 50, "c": 100}
ORDINALS = ["erste", "zweite", "dritte", "vierte", "fünfte", "sechste", "siebente", "achte", "neunte", "zehnte"]
ORDINALS_EN = ["first", "second", "third", "fourth", "fifth", "sixth", "seventh", "eighth", "ninth", "tenth"]
ORDINAL = re.compile(r"[\s.,;:—–-]*\b(%s)[rs]?\s+(band|theil|teil|abt(h)?eilung|th\.|bd\.)\b.*$|[\s.,;:—–-]*\b(%s)\s+(volume|part|series)\b.*$"
                     % ("|".join(ORDINALS), "|".join(ORDINALS_EN)), re.I)
OF = re.compile(r"\s*\((\d+) of \d+\).*$", re.I)


def vol_no(t):
    """Volume number named in a title, or None."""
    m = ORDINAL.search(t)
    if m:
        w = (m.group(1) or m.group(4)).lower()
        return (ORDINALS.index(w) if w in ORDINALS else ORDINALS_EN.index(w)) + 1
    m = OF.search(t)
    if m:
        return int(m.group(1))
    m = VOLUME.search(t)
    if not m or not m.group(5):
        return None
    v = m.group(5).lower()
    if v.isdigit():
        return int(v)
    n = 0
    for a, b in zip(v, v[1:] + " "):
        n += -ROMAN[a] if ROMAN.get(b, 0) > ROMAN[a] else ROMAN[a]
    return n


def base(t):
    return VOLUME.sub("", OF.sub("", ORDINAL.sub("", t))).strip(" .,;:—–-")


class Groups:
    def __init__(self, n):
        self.parent = list(range(n))

    def find(self, i):
        while self.parent[i] != i:
            self.parent[i] = self.parent[self.parent[i]]
            i = self.parent[i]
        return i

    def join(self, i, j):
        self.parent[self.find(i)] = self.find(j)


def merge(members):
    vols = [vol_no(r["t"]) for r in members]
    order = sorted(range(len(members)), key=lambda i: (vols[i] is not None, vols[i] or 0, -(members[i].get("s") or 0),
                                                       int(members[i]["id"]) if str(members[i]["id"]).isdigit() else 0))
    members = [members[i] for i in order]
    vols = [vols[i] for i in order]
    lead = members[0]
    out = dict(lead)
    ids = [int(r["id"]) for r in members if str(r["id"]).isdigit()]
    consecutive = len(ids) == len(members) and max(ids) - min(ids) < len(ids) + 2
    if all(v is not None for v in vols):
        out["t"], kind = base(lead["t"]) or lead["t"], "volumes"  # "Fünfter Theil der ..." has nothing left
    elif consecutive and len({r["t"] for r in members}) == 1:
        kind = "volumes"
    else:
        kind = "versions"  # editions, reprints, or a complete edition beside its parts
    out["ek"] = kind
    out["ed"] = [{k: r[k] for k in ("id", "t", "r", "u", "y") if r.get(k) is not None} for r in members]
    uniq = lambda xs: list(dict.fromkeys(x for x in xs if x))
    out["c"] = [c for c in ("natural_history", "travel") if any(c in r["c"] for r in members)]
    out["p"] = uniq(p for r in members for p in r.get("p") or [])
    out["kw"] = uniq(k for r in members for k in r.get("kw") or [])[:6]
    out["also"] = uniq(a for r in members for a in r.get("also") or [])
    out["sub"] = " | ".join(uniq(s for r in members for s in (r.get("sub") or "").split(" | ")))
    out["s"] = max((r.get("s") or 0) for r in members) if lead.get("s") is not None else lead.get("s")
    dated = [r for r in members if r.get("y")]
    if dated:
        first = min(dated, key=lambda r: r["y"])
        out["y"] = first["y"]
        if first.get("ys"):
            out["ys"] = first["ys"]
    for k in ("p", "kw", "also", "sub"):
        if not out[k]:
            out.pop(k)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="overview/data.json")
    ap.add_argument("--show", action="store_true", help="print the groups")
    args = ap.parse_args()
    rows = json.load(open(args.data, encoding="utf-8"))
    if any(r.get("ek") in ("volumes", "versions") for r in rows):
        sys.exit("data.json is already grouped; rebuild it with build_overview.py first")

    g = Groups(len(rows))
    side = lambda r: (r["src"], tuple(r["l"]), norm(surname(r)))
    by_title, by_side, by_work = {}, defaultdict(list), {}
    for i, r in enumerate(rows):
        if "ed" in r:
            continue  # grouped at the source (Humboldt's correspondences)
        key = side(r) + (" ".join(words(base(r["t"]))),)
        if key in by_title:
            g.join(i, by_title[key])
        else:
            by_title[key] = i
        by_side[side(r)].append(i)
        if r.get("wk"):
            wk = (r["src"], tuple(r["l"]), r["wk"])
            if wk in by_work:
                g.join(i, by_work[wk])
            else:
                by_work[wk] = i
    for idx in by_side.values():
        if not side(rows[idx[0]])[2]:
            continue  # anonymous: titles alone are too weak
        ws = {i: set(words(base(rows[i]["t"]))) for i in idx}
        for a in range(len(idx)):
            for b in range(a + 1, len(idx)):
                wa, wb = ws[idx[a]], ws[idx[b]]
                if len(wa) >= 4 and len(wb) >= 4 and len(wa & wb) / len(wa | wb) >= 0.8:
                    g.join(idx[a], idx[b])

    clusters = defaultdict(list)
    for i, r in enumerate(rows):
        clusters[g.find(i)].append(r)
    out = []
    for members in clusters.values():
        if len(members) == 1:
            out.append(members[0])
            continue
        m = merge(members)
        out.append(m)
        if args.show:
            print(f"{m['src']} {m['ek']:8} {m['t'][:60]!r} <- " + " | ".join(r["t"][:40] for r in members))

    # translations: same Wikidata work, other language
    work = defaultdict(list)
    for r in out:
        if r.get("wk"):
            work[(r["src"], r["wk"])].append(r)
    linked = 0
    for rs in work.values():
        if len({tuple(r["l"]) for r in rs}) < 2:
            continue
        for r in rs:
            r["tr"] = [{"id": o["id"], "t": o["t"], "l": o["l"], "r": o["r"]} for o in rs if o is not r and o["l"] != r["l"]]
            linked += 1

    json.dump(out, open(args.data, "w", encoding="utf-8"), ensure_ascii=False, separators=(",", ":"))
    merged = sum(1 for r in out if "ed" in r)
    print(f"{len(rows)} entries -> {len(out)} works; {merged} merged groups; {linked} linked as translations", file=sys.stderr)


if __name__ == "__main__":
    main()