#!/usr/bin/env catnip
# Checkpoint Redis d'un pré-agrégat OpenAlex via freeze/thaw
# - premier run: parsing séquentiel + pré-agrégation ND process + checkpoint Redis
# - second run: thaw depuis Redis + fusion finale uniquement
# Sur un vrai export OpenAlex, on augmente batch_size et nd_workers.
#
# DEPS: redis orjson
# REQUIS: serveur Redis local (localhost:6379)

pragma('nd_mode', 'process')
pragma('nd_workers', 4)

pathlib = import('pathlib')
redis = import('redis')
orjson = import('orjson')
time = import('time')
math = import('math')

openalex_path = pathlib.Path(META.file).parent / "data" / "openalex-works-sample.jsonl"
batch_size = 20
checkpoint_key = "catnip:openalex:works:checkpoint:v1"
checkpoint_ttl = 3600

r = redis.Redis(host="localhost", port=6379, db=0)

# Wrapper de batch : une liste de rows serait broadcastée récursivement
# par le ND-map, un struct est une feuille
struct Chunk { rows; }

# I/O: lecture JSONL en batches de bytes bruts

read_jsonl_batches = (path, size: int) => {
    raw = pathlib.Path(path).read_bytes()
    lines = list()
    for line in raw.split(b"\n") {
        stripped = line.strip()
        if stripped { lines = lines + list(stripped) }
    }
    batches = list()
    i = 0
    while i < len(lines) {
        batches = batches + list(lines.[i:i + size])
        i = i + size
    }
    batches
}

# Parsing : désérialise un batch de lignes JSON brutes

parse_batch = (lines) => {
    rows = list()
    for line in lines {
        obj = orjson.loads(line)

        concepts = list()
        seen = 0
        for c in (obj['concepts'] ?? list()) {
            name = c['display_name'] ?? None
            if name and seen < 3 {
                concepts = concepts + list(name)
                seen = seen + 1
            }
        }

        rows = rows + list(dict(
                year    =obj['publication_year'] ?? None,
                type    =obj['type'] ?? "unknown",
                cited   =obj['cited_by_count'] ?? 0,
                concepts=concepts,
            ))
    }
    rows
}

# Résumé : pré-agrégat pur par batch (CPU-bound)
# Coût artificiel via cpu_cost pour rendre le checkpoint visible.
# Import math dans la fonction: en mode process, chaque worker est
# un processus séparé sans accès au scope parent.

cpu_cost = (n: int): float => {
    math = import('math')  # noqa: W204 -- réimport worker ND
    acc = 0.0
    for i in range(n) {
        x = i + 1
        acc = acc + math.sqrt(x) / (x + 0.5)
    }
    acc
}

summarize_batch = (rows) => {
    total = 0
    cited_sum = 0
    by_year = dict()
    by_type = dict()
    by_concept = dict()

    for row in rows {
        year = row['year']
        if year == None { continue }

        cpu_cost(150)

        total = total + 1
        cited_sum = cited_sum + row['cited']

        by_year[year] = by_year.get(year, 0) + 1
        by_type[row['type']] = by_type.get(row['type'], 0) + 1

        for concept in row['concepts'] {
            by_concept[concept] = by_concept.get(concept, 0) + 1
        }
    }

    dict(
        total=total,
        cited_sum=cited_sum,
        by_year=by_year,
        by_type=by_type,
        by_concept=by_concept,
    )
}

# Merge : fusion des pré-agrégats partiels

merge_counts = (dst, src) => {
    for k in src.keys() {
        dst[k] = dst.get(k, 0) + src[k]
    }
    dst
}

merge_summaries = (partials) => {
    total = 0
    cited_sum = 0
    by_year = dict()
    by_type = dict()
    by_concept = dict()

    for p in partials {
        total = total + p['total']
        cited_sum = cited_sum + p['cited_sum']
        by_year = merge_counts(by_year, p['by_year'])
        by_type = merge_counts(by_type, p['by_type'])
        by_concept = merge_counts(by_concept, p['by_concept'])
    }

    dict(
        total=total,
        cited_sum=cited_sum,
        avg_citations=if total > 0 { round(cited_sum / total, 2) } else { 0 },
        by_year=by_year,
        by_type=by_type,
        by_concept=by_concept,
    )
}

# Pipeline principal

started = time.perf_counter()

blob = r.get(checkpoint_key)

if blob {
    print("checkpoint hit:", checkpoint_key)
    partials = thaw(blob)
    print("reprise depuis Redis -> fusion finale uniquement")
} else {
    print("checkpoint miss:", checkpoint_key)
    print("lecture source:", openalex_path)

    batches = read_jsonl_batches(openalex_path, batch_size)
    print("batches:", len(batches))

    # Parsing séquentiel (I/O + orjson dans le scope parent)
    parsed = list()
    for batch in batches {
        parsed = parsed + list(Chunk(parse_batch(batch)))
    }

    # Summarize ND process (CPU-bound, un worker par batch)
    partials = parsed.[~>(chunk) => { summarize_batch(chunk.rows) }]

    r.setex(checkpoint_key, checkpoint_ttl, freeze(partials))
    print("checkpoint écrit dans Redis (TTL=", checkpoint_ttl, "s)")
}

final = merge_summaries(partials)
elapsed = round(time.perf_counter() - started, 3)

print("\nrésumé global")
print("works:", final['total'])
print("citations total:", final['cited_sum'])
print("citations moyennes:", final['avg_citations'])
print("années:", final['by_year'])
print("types:", final['by_type'])
print("concepts:", final['by_concept'])
print("\ntemps total:", elapsed, "s")