"""Per-cell h3 calls over millions of cells, split across worker processes (h3 is a C library called once per cell
from Python: the loop, not the math, is the cost). Same values as the serial loops, in the same order."""
import os

import h3.api.basic_int as h3
import numpy as np

MIN_PARALLEL = 200_000         # fewer cells: the serial loop beats starting workers


def rings(ids):
    """Neighbours of each cell (h3.grid_ring order), flat, with per-cell counts (pentagons have five)."""
    ids = np.asarray(ids, dtype=np.uint64)
    counts = np.empty(len(ids), np.int64)

    def walk():                                                    # streamed: no list of millions of Python ints
        for k, c in enumerate(ids.tolist()):
            r = h3.grid_ring(c, 1)
            counts[k] = len(r)
            yield from r
    flat = np.fromiter(walk(), dtype=np.uint64)
    return flat, counts


def centres_areas(ids):
    """(lat, lon) degrees (n, 2) and areas in rad² (n,) of the cells."""
    ids = np.asarray(ids, dtype=np.uint64)
    ll = np.fromiter((v for c in ids.tolist() for v in h3.cell_to_latlng(c)), dtype=np.float64, count=2 * len(ids))
    area = np.fromiter((h3.cell_area(c, unit="rads^2") for c in ids.tolist()), dtype=np.float64, count=len(ids))
    return ll.reshape(-1, 2), area


def workers() -> int:
    v = os.environ.get("MAPVIEW_H3_WORKERS")
    if v is not None:
        return max(0, int(v))
    return max(0, min(6, (os.cpu_count() or 2) - 2))


def run(fn, ids):
    """fn(ids) (a function of this module returning arrays per cell, or (flat, counts)), split over workers."""
    ids = np.asarray(ids, dtype=np.uint64)
    w = workers()
    if w < 2 or len(ids) < MIN_PARALLEL:
        return fn(ids)
    import multiprocessing as mp
    chunks = np.array_split(ids, w * 2)
    with mp.get_context("fork").Pool(w) as pool:             # fork, as tiles' and export's pools
        parts = pool.map(fn, chunks)
    return tuple(np.concatenate([p[i] for p in parts]) for i in range(len(parts[0])))