# Part of the Hex9 (H9) Project
# Copyright ©2025, Ben Griffin
# Licensed under the Apache License, Version 2.0

"""
GCDBary: parallel-aware direct hop g_gcd -> b_oct.

Outside PostgreSQL: uses ProcessPoolExecutor so each worker runs its own Python
interpreter and GIL, giving true CPU parallelism for the compute-heavy
c_ell->c_oct->b_raw chain

Inside PostgreSQL (plpython3u): process spawning is not permitted, so the large-
batch path falls back to in-process serial execution.  Hex-fill queries are
bounded in size so the performance cost is acceptable.
"""

import numpy as np
from hhg9 import Points
from hhg9.base.projection import Projection
from concurrent.futures import ProcessPoolExecutor
from functools import partial

# Detect PostgreSQL plpython3u environment once at import time.
try:
    import plpy as _plpy        # noqa: F401
    _IN_POSTGRES = True
except ImportError:
    _IN_POSTGRES = False


# --- WORKER (top-level for pickling) ---

def _worker_project_gcd_to_braw(chunk_ll: np.ndarray, accuracy: float = 1e-9):
    """
    Process worker: rebuilds the minimal projection stack,
    and projects one chunk g_gcd -> b_oct.
    Each worker runs in a separate process with its own GIL.
    """
    from hhg9 import Registrar, Points
    reg = Registrar()
    _ = reg.domain('g_gcd')
    _ = reg.domain('c_ell')
    _ = reg.domain('c_oct')
    b_raw = reg.domain('b_raw')
    reg.projection('oct_ell').accuracy = accuracy
    pts = Points(np.ascontiguousarray(chunk_ll, dtype=np.float64), 'g_gcd')
    bc = reg.project(pts, ['g_gcd', 'c_ell', 'c_oct', b_raw])
    return (np.ascontiguousarray(bc.coords, dtype=np.float64),
            np.ascontiguousarray(bc.oid))


def _project_gcd_to_braw_parallel(ll_array: np.ndarray, accuracy, workers=0, chunk=8_000):
    """Orchestrate process-parallel projection over chunks of ll_array."""
    if workers <= 0:
        import os
        cpu_count = os.cpu_count() or 2
        needed_workers = int(np.ceil(len(ll_array) / chunk))
        workers = min(cpu_count, max(1, needed_workers))

    chunks = [np.ascontiguousarray(ll_array[i:i + chunk])
              for i in range(0, len(ll_array), chunk)]

    worker_func = partial(_worker_project_gcd_to_braw, accuracy=accuracy)

    out_coords, out_oids = [], []
    with ProcessPoolExecutor(max_workers=workers) as ex:
        try:
            for xy, oid in ex.map(worker_func, chunks):
                out_coords.append(xy)
                out_oids.append(oid)
        finally:
            ex.shutdown(cancel_futures=True)

    if not out_coords:
        return np.empty((0, 2)), np.empty(0, dtype=np.uint8)

    return np.vstack(out_coords), np.concatenate(out_oids)


# --- THE PROJECTION CLASS ---

class GCDBraw(Projection):
    """
    Parallel-Aware Direct Hop: g_gcd -> b_raw

    Switches between serial execution (small batches or PostgreSQL) and
    ProcessPoolExecutor (large batches outside PostgreSQL).
    """

    def __init__(self, registrar, *,
                 workers: int | None = None,
                 chunk: int = 8_000,
                 threshold: int = 8_000):

        super().__init__(registrar, 'gcd_brw', 'g_gcd', 'b_raw')
        self.reg = registrar
        self.workers = workers
        self.chunk = int(chunk)
        self.threshold = int(threshold)
        self.b_raw = registrar.domain('b_raw')
        # Cached for process workers (they can't share the parent's objects)
        self.accuracy = registrar.projection('oct_ell').accuracy

    def set_parallel(self, *, workers=None, chunk=None, threshold=None):
        """Configure parallel execution."""
        if workers is not None: self.workers = workers
        if chunk is not None: self.chunk = int(chunk)
        if threshold is not None: self.threshold = int(threshold)
        return self

    def forward(self, arr: Points) -> Points:
        coords = arr.coords
        n = 0 if coords is None else len(coords)

        # 1. Edge Case: Empty
        if n == 0:
            return Points(np.empty((0, 2)), self.fwd_cs,
                          oid=np.empty(0, dtype=np.uint8),
                          samples=arr.samples)

        # 2. Serial: small batch, or always serial inside PostgreSQL
        if n < self.threshold or _IN_POSTGRES:
            return self.reg.project(arr, ['g_gcd', 'c_ell', 'c_oct', 'b_raw'])

        # 3. Large batch outside PostgreSQL: process pool
        xy, oid = _project_gcd_to_braw_parallel(
            coords,
            accuracy=self.accuracy,
            workers=self.workers or 0,
            chunk=self.chunk,
        )
        return Points(xy, self.fwd_cs, oid=oid, samples=arr.samples)

    def backward(self, arr: Points) -> Points:
        return self.reg.project(arr, ['b_raw', 'c_oct', 'c_ell', 'g_gcd'])
