02 — Pipeline idioms

from geotoolz import Operator, Sequential


class Add(Operator):
    """Tiny test operator: x → x + n."""

    def __init__(self, n: int) -> None:
        self.n = n

    def _apply(self, x: int) -> int:
        return x + self.n

    def get_config(self) -> dict:
        return {"n": self.n}


class _NamedStep(Operator):
    """Named identity-with-tag stand-in for any domain op (NDVI, MaskClouds, ...)."""

    def __init__(self, name: str, *, factor: float = 1.0) -> None:
        self.name = name
        self.factor = factor

    def _apply(self, x):
        return x * self.factor

    def get_config(self) -> dict:
        return {"name": self.name, "factor": self.factor}


# A 'pipeline' that stands in for a real preprocessing chain
demo_pipe = Sequential(
    [
        _NamedStep("MaskClouds", factor=1.0),
        _NamedStep("PercentileClip", factor=0.98),
        _NamedStep("NDVI", factor=0.5),
    ]
)
print(demo_pipe)
print("demo_pipe(10) =", demo_pipe(10))

from geotoolz import Tap


seen = []
pipe = Sequential(
    [
        Add(1),
        Tap(lambda x: seen.append(("after Add(1)", x))),
        Add(10),
        Tap(lambda x: seen.append(("after Add(10)", x))),
    ]
)
out = pipe(0)
print("result :", out)
print("log    :", seen)

from geotoolz import Snapshot


snap = Snapshot()
pipe = Sequential(
    [
        Add(1),
        snap.at("after_first"),
        Add(10),
        snap.at("after_second"),
        Add(100),
    ]
)
pipe(0)
print("captured:", snap.captures)

from pipekit import Profile


prof = Profile()
pipe = Sequential([prof.wrap(Add(1)), prof.wrap(Add(10)), prof.wrap(Add(100))])
for _ in range(100):
    pipe(0)
print("Per-step timings:")
for name, stats in prof.report().items():
    print(
        f"  {name:18s}  mean={stats['mean'] * 1e6:7.1f} µs  "
        f"calls={int(stats['calls'])}  total={stats['total'] * 1e3:.2f} ms"
    )

import numpy as np
from pipekit import Histogram


hist = Histogram(bins=5, to_array=lambda a: np.asarray(a, dtype=float).ravel().tolist())
pipe = Sequential(
    [
        hist.at("input"),
        _NamedStep("PercentileClip", factor=0.5),
        hist.at("after_clip"),
    ]
)
pipe(np.arange(20))
print("input    counts:", hist.results["input"]["counts"])
print("clipped  counts:", hist.results["after_clip"]["counts"])

from geotoolz import Lambda, ShapeTrace


trace = ShapeTrace(mode="diff_only")
pipe = Sequential(
    [
        trace,
        Lambda(
            lambda x: np.asarray([x, x * 2, x * 3], dtype=np.int16), name="to_array"
        ),
        trace,
        Lambda(lambda a: a.astype(np.float32), name="to_float"),
        trace,
    ]
)
pipe(7)

import contextvars
from contextlib import contextmanager


_active_spies: contextvars.ContextVar[tuple] = contextvars.ContextVar(
    "geotoolz_spy_stack", default=()
)


class Spy:
    def __init__(self) -> None:
        self._hooks: dict[type, list] = {}

    @classmethod
    @contextmanager
    def scoped(cls):
        spy = cls()
        token = _active_spies.set((*_active_spies.get(), spy))
        try:
            yield spy
        finally:
            _active_spies.reset(token)

    def on(self, op_type: type, fn) -> None:
        self._hooks.setdefault(op_type, []).append(fn)


def _fire(op, x_in, x_out):
    for spy in _active_spies.get():
        for fn in spy._hooks.get(type(op), ()):
            fn(op, x_in, x_out)


# Wrap call sites manually for the demo (in v0.2, Operator.__call__ does this).
class _TrackedAdd(Add):
    def __call__(self, x):
        out = super().__call__(x)
        _fire(self, x, out)
        return out


pipe = Sequential([_TrackedAdd(1), _TrackedAdd(10), _TrackedAdd(100)])

with Spy.scoped() as spy:
    spy.on(_TrackedAdd, lambda op, xi, xo: print(f"  {op}: {xi} → {xo}"))
    pipe(0)

print()
print("outside the scope — silent:")
pipe(0)

class DiffError(Exception):
    pass


class Diff(Operator):
    """Compare input against a stored reference; raise on drift beyond atol."""

    forbid_in_yaml = True

    def __init__(self, *, reference, atol: float = 1e-6) -> None:
        self.reference = np.asarray(reference)
        self.atol = atol

    def _apply(self, x):
        arr = np.asarray(x)
        diff = float(np.max(np.abs(arr - self.reference)))
        if diff > self.atol:
            raise DiffError(f"max abs diff {diff:.4g} exceeds atol={self.atol}")
        return x

    def get_config(self) -> dict:
        return {"atol": self.atol, "ref_shape": self.reference.shape}


# Bless once when you trust the output:
golden = np.array([1, 11, 111])

# Then catches drift on subsequent runs:
arr_in = np.array([0, 10, 110])
pipe = Sequential(
    [
        Lambda(lambda a: a + 1, name="add1"),
        Diff(reference=golden, atol=1e-6),
    ]
)
print("matches golden:", pipe(arr_in))

try:
    pipe(np.array([0, 10, 111]))  # off by one in last element
except DiffError as e:
    print("caught regression:", e)

from geotoolz import Branch, Identity


guarded = Branch(
    predicate=lambda x: x > 0,
    if_true=Add(100),
    if_false=Identity(),  # default — pass-through for non-positive
)
print("guarded(5) =", guarded(5))
print("guarded(-5)=", guarded(-5))

from geotoolz import Switch


router = Switch(
    key=lambda x: "even" if x % 2 == 0 else "odd",
    cases={
        "even": Add(1),
        "odd": Sequential([Add(100), Add(100)]),
    },
)
print("router(4) =", router(4))
print("router(3) =", router(3))

from pipekit import Try


class Boom(Operator):
    def _apply(self, x):
        raise ConnectionError(f"network unreachable on input {x}")


robust = Try(primary=Boom(), fallback=Add(0), on=(ConnectionError,))
print("robust(7) =", robust(7))

from pipekit import Coalesce


def is_positive(out):
    return out > 0


coalesced = Coalesce(
    sources=[
        Add(-100),  # produces negative — rejected
        Add(-10),  # still negative — rejected
        Add(50),  # positive — wins
    ],
    is_ok=is_positive,
)
print("coalesced(5) =", coalesced(5))

from pipekit import Retry


_attempts = {"count": 0}


class FlakeyOp(Operator):
    def _apply(self, x):
        _attempts["count"] += 1
        if _attempts["count"] < 3:
            raise TimeoutError(f"transient failure (attempt {_attempts['count']})")
        return x + 1000


robust_call = Retry(op=FlakeyOp(), attempts=5, base_delay=0.0, on=(TimeoutError,))
print("retrying...")
out = robust_call(7)
print(f"succeeded after {_attempts['count']} attempts, out = {out}")

from geotoolz import Fanout


products = Fanout(
    {
        "doubled": Lambda(lambda x: x * 2, name="double"),
        "squared": Lambda(lambda x: x * x, name="square"),
        "labelled": _NamedStep("classify", factor=10),
    }
)
print(products(7))

class ApplyToBands(Operator):
    """Apply ``inner`` to each slice along ``axis``, stack results."""

    def __init__(self, inner: Operator, *, axis: int = 0) -> None:
        self.inner = inner
        self.axis = axis

    def _apply(self, arr):
        arr = np.asarray(arr)
        slices = [
            self.inner(np.take(arr, i, axis=self.axis))
            for i in range(arr.shape[self.axis])
        ]
        return np.stack(slices, axis=self.axis)

    def get_config(self) -> dict:
        return {
            "inner": type(self.inner).__name__,
            "inner_config": self.inner.get_config(),
            "axis": self.axis,
        }


# stand-in: scale each "band" by a factor of 10
op = ApplyToBands(Lambda(lambda x: x * 10, name="scale"), axis=0)
arr = np.arange(12).reshape(3, 4)
print("input  :")
print(arr)
print("per-band scaled:")
print(op(arr))

import hashlib
import json

from pipekit import Cache


cached = Cache(Sequential([Add(1), Add(10)]))
for _ in range(5):
    cached(0)
print(f"5 calls → {cached.hits} hits, {cached.misses} miss")
print("result :", cached(0))

_current_mode: contextvars.ContextVar = contextvars.ContextVar(
    "geotoolz_pipeline_mode", default=None
)


@contextmanager
def pipeline_mode(name: str):
    """Enter a named mode for the duration of the with block."""
    if name not in ("train", "eval"):
        raise ValueError(f"mode must be 'train' or 'eval', got {name!r}")
    token = _current_mode.set(name)
    try:
        yield name
    finally:
        _current_mode.reset(token)


class ModeGated(Operator):
    """Wraps an op that only fires under a specific mode."""

    forbid_in_yaml = True

    def __init__(self, inner: Operator, *, mode: str) -> None:
        self.inner, self.required_mode = inner, mode

    def _apply(self, x):
        if _current_mode.get() == self.required_mode:
            return self.inner(x)
        return x

    def get_config(self) -> dict:
        return {"mode": self.required_mode, "inner": type(self.inner).__name__}


# Augmentation only fires in train mode; deterministic preprocess always runs.
pipe = Sequential(
    [
        Add(1),  # preprocess (always)
        ModeGated(Add(1000), mode="train"),  # augmentation (train only)
    ]
)

with pipeline_mode("train"):
    print("train mode:", pipe(0))

with pipeline_mode("eval"):
    print("eval mode :", pipe(0))

# Outside any scope: mode is None, so ModeGated is silent
print("no mode   :", pipe(0))

class Provenance(Operator):
    """Wraps a pipeline; attaches a metadata dict on its way out.

    Real pipelines write this to the output COG's tags via georeader.
    """

    forbid_in_yaml = True

    def __init__(self, inner: Operator) -> None:
        self.inner = inner

    def _apply(self, x):
        out = self.inner(x)
        cfg = self.inner.get_config()
        cfg_canon = json.dumps(cfg, sort_keys=True)
        provenance = {
            "pipeline_hash": hashlib.sha256(cfg_canon.encode()).hexdigest()[:16],
            "operators": _op_names(self.inner),
            "geotoolz_version": "0.0.1",
        }
        # In real code: out.metadata["provenance"] = provenance
        return {"data": out, "_provenance": provenance}

    def get_config(self) -> dict:
        return {"inner": type(self.inner).__name__}


def _op_names(op: Operator) -> list:
    if isinstance(op, Sequential):
        return [type(o).__name__ for o in op.operators]
    return [type(op).__name__]


tracked = Provenance(Sequential([Add(1), Add(10)]))
result = tracked(0)
print("data        :", result["data"])
print("provenance  :")
for k, v in result["_provenance"].items():
    print(f"  {k}: {v}")

from pipekit import AssertDType, AssertHasAttribute, AssertShape, QCError


class FakeCarrier:
    """Minimal stand-in with the attributes the Assert* family inspects."""

    def __init__(self, shape, dtype, crs="EPSG:4326"):
        self.shape = shape
        self.dtype = dtype
        self.crs = crs


# Three contracts in one Sequential — input must be (4, H, W) float32 with a CRS.
contract = (
    AssertShape((4, None, None)) | AssertDType("float32") | AssertHasAttribute("crs")
)
print("good carrier passes through:", contract(FakeCarrier((4, 256, 256), "float32")))

print("\nbad shape raises:")
try:
    contract(FakeCarrier((3, 256, 256), "float32"))
except QCError as e:
    print(" ", e)


# Domain-specific range check, written in the same pass-through shape:
class AssertValueRange(Operator):
    """Pass-through; raise QCError on values outside [min_val, max_val]."""

    def __init__(self, *, min_val: float, max_val: float) -> None:
        self.min, self.max = min_val, max_val

    def _apply(self, x):
        arr = np.asarray(x)
        lo, hi = float(np.nanmin(arr)), float(np.nanmax(arr))
        if lo < self.min or hi > self.max:
            raise QCError(
                f"range [{lo:.3g}, {hi:.3g}] outside [{self.min}, {self.max}]"
            )
        return x

    def get_config(self) -> dict:
        return {"min_val": self.min, "max_val": self.max}


print("\nrange check on raw values:")
range_check = AssertValueRange(min_val=0, max_val=100)
try:
    range_check(np.array([5, 25, 150]))
except QCError as e:
    print(" caught:", e)

from pipekit import Quarantine


log: list = []

pipe = Sequential(
    [
        Quarantine(
            check=lambda arr: float(np.nanmax(arr)) <= 100,
            sentinel=np.array([-1]),  # marker for "quarantined"
            on_quarantine=lambda arr: log.append(
                f"quarantined max={float(np.nanmax(arr))}"
            ),
        ),
    ]
)

print("good:", pipe(np.array([5, 25])))
print("bad: ", pipe(np.array([5, 999])))
print("log :", log)

from geotoolz import Const, Sink


# Branch with explicit Identity in if_false — self-documenting no-op
opt_clean = Branch(
    predicate=lambda x: x > 0,
    if_true=Sequential([Add(100), Add(-50)]),
    if_false=Identity(),
)
print("opt_clean(10) =", opt_clean(10))
print("opt_clean(-5) =", opt_clean(-5))


# Sink: write side effect + pass through
written: list = []
checkpoint_pipe = Sequential(
    [
        Add(1),
        Sink(written.append, name="checkpoint_after_first"),
        Add(10),
    ]
)
print("checkpoint_pipe(0) =", checkpoint_pipe(0))
print("checkpoint stored  :", written)


# Const: fixed value, useful for synthetic sources / golden tests
golden_input = Const(np.arange(5))
synth_test = Sequential(
    [
        golden_input,  # ignores actual input
        Lambda(lambda a: a.sum(), name="sum"),
    ]
)
print("synth test sum:", synth_test("real_input_ignored"))

class Subsample(Operator):
    """Stride-decimate the last two axes. Standalone transform."""

    def __init__(self, *, stride: int = 10) -> None:
        if stride <= 0:
            raise ValueError("stride must be positive")
        self.stride = stride

    def _apply(self, arr):
        arr = np.asarray(arr)
        return arr[..., :: self.stride, :: self.stride]

    def get_config(self) -> dict:
        return {"stride": self.stride}


big = np.arange(100).reshape(10, 10)
small = Subsample(stride=3)(big)
print("input shape   :", big.shape)
print("decimated shape:", small.shape)
print(small)