"""H3 hexagonal grid as a CSR cell graph. Stage `grid`."""
from __future__ import annotations

from dataclasses import dataclass
from functools import cached_property

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

from .sphere import latlon_to_xyz, tangent_dir


EDGE_BLOCK = 1 << 20


def edge_blocks(n: int, block: int | None = None):
    """Slices over n edges, EDGE_BLOCK at a time: per-edge (row-wise) formulas give the same values block by block,
    with (block, 3) float64 temporaries instead of (edges, 3) ones (290 MB each at r5, 2 GB at r6)."""
    block = block or EDGE_BLOCK
    return [slice(a, min(a + block, n)) for a in range(0, n, block)]


def rowdot_at(A, idx, t) -> np.ndarray:
    """np.sum(A[idx] * t, axis=1) — per edge, A's row at an end dotted with the edge's vector — in edge blocks."""
    out = np.empty(len(idx))
    for s in edge_blocks(len(idx)):
        out[s] = np.sum(A[idx[s]] * t[s], axis=1)
    return out


def cell_parents(cells, res: int) -> np.ndarray:
    """h3.cell_to_parent for an array of cells (all of resolution ≥ res), by the index bits: the resolution field
    set to res, the digits below it set to 7 (unused). Same ids as h3, without a Python call per cell."""
    cells = np.asarray(cells, dtype=np.uint64)
    if len(cells) and int(((cells >> np.uint64(52)) & np.uint64(15)).min()) < res:
        raise ValueError(f"cell_parents: a cell is coarser than resolution {res}")
    unused = 0
    for r in range(res + 1, 16):
        unused |= 7 << ((15 - r) * 3)
    out = (cells & np.uint64(~(15 << 52) & (2**64 - 1))) | np.uint64(res << 52)
    return out | np.uint64(unused)


@dataclass
class Grid:
    res: int
    radius_km: float
    ids: np.ndarray
    lat: np.ndarray
    lon: np.ndarray
    xyz: np.ndarray
    area_km2: np.ndarray
    nbr_ptr: np.ndarray
    nbr_idx: np.ndarray

    @property
    def n(self) -> int:
        return len(self.ids)

    @cached_property
    def counts(self):
        return np.diff(self.nbr_ptr)

    @cached_property
    def src(self):
        return np.repeat(np.arange(self.n), self.counts)

    @property
    def dst(self):
        return self.nbr_idx

    @cached_property
    def edge_km(self):
        out = np.empty(len(self.dst))
        for s in edge_blocks(len(self.dst)):
            d = np.sum(self.xyz[self.src[s]] * self.xyz[self.dst[s]], axis=1)
            out[s] = self.radius_km * np.arccos(np.clip(d, -1.0, 1.0))
        return out

    @cached_property
    def edge_tangents(self):
        out = np.empty((len(self.dst), 3))
        for s in edge_blocks(len(self.dst)):
            out[s] = tangent_dir(self.xyz[self.src[s]], self.xyz[self.dst[s]])
        return out

    @cached_property
    def spacing_km(self) -> float:
        return float(self.edge_km.mean())

    def cell_index(self, lat: float, lon: float) -> int:
        c = np.uint64(h3.latlng_to_cell(float(lat), float(lon), self.res))
        return int(np.searchsorted(self.ids, c))

    def to_arrays(self) -> dict:
        return {"g_ids": self.ids, "g_lat": self.lat, "g_lon": self.lon, "g_xyz": self.xyz,
                "g_area_km2": self.area_km2, "g_nbr_ptr": self.nbr_ptr, "g_nbr_idx": self.nbr_idx}

    @classmethod
    def from_arrays(cls, d: dict, res: int, radius_km: float) -> "Grid":
        return cls(res, radius_km, d["g_ids"], d["g_lat"], d["g_lon"], d["g_xyz"], d["g_area_km2"],
                   d["g_nbr_ptr"], d["g_nbr_idx"])


def build_grid(res: int, radius_km: float) -> Grid:
    cells: list[int] = []
    for r0 in h3.get_res0_cells():
        cells.extend(h3.cell_to_children(r0, res))
    ids = np.array(sorted(cells), dtype=np.uint64)
    ll = np.array([h3.cell_to_latlng(int(c)) for c in ids], dtype=np.float64)
    area = np.array([h3.cell_area(int(c), unit="rads^2") for c in ids]) * radius_km**2
    rings = [h3.grid_ring(int(c), 1) for c in ids]
    counts = np.array([len(r) for r in rings], dtype=np.int64)
    ptr = np.concatenate([[0], np.cumsum(counts)]).astype(np.int64)
    flat = np.array([c for r in rings for c in r], dtype=np.uint64)
    idx = np.searchsorted(ids, flat).astype(np.int64)
    return Grid(res, radius_km, ids, ll[:, 0], ll[:, 1], latlon_to_xyz(ll[:, 0], ll[:, 1]), area, ptr, idx)


def run(ctx) -> dict:
    return build_grid(ctx.res, float(ctx.cfg["planet"]["radius_km"])).to_arrays()