src/data_pbdb.py

"""
Paleobiology Database (PBDB) real-data client.

VERIFICATION STATUS (be precise about what's actually confirmed vs. not —
this matters because paleobiodb.org's robots.txt blocked automated
fetching in the development sandbox, so the exact live JSON field names
below were NOT confirmed against a real response in that environment):

  CONFIRMED (via web search of PBDB's own documentation, dev.paleobiodb.org
  and multiple independent third-party docs, cross-checked):
    - Base endpoint: https://paleobiodb.org/data1.2/occs/list.{json,csv,txt}
    - No authentication required; data is CC0.
    - `limit`: positive integer, 0, or the literal string "all" (NOT a
      fixed default of 500 — that claim in an earlier draft plan was
      unverified and is not repeated here).
    - `offset`: skips N records at the start of the result set (used for
      pagination together with `limit`).
    - `base_name`: taxon + all children (e.g. "Dinosauria").
    - `interval`: named geologic interval (e.g. "Miocene").
    - `lngmin`/`lngmax`/`latmin`/`latmax`: bounding-box filter.
    - `show`: comma-separated list of extra output blocks (e.g. "coords",
      "classext", "full").
    - JSON responses wrap records in a top-level "records" array.

  NOT CONFIRMED as of the previous session, NOW PARTIALLY CONFIRMED:
    A real live-data example (not just prose docs) surfaced this session
    in the `paleobioDB` R package documentation, showing an actual
    response row. It confirms:
      - Numeric age fields are named `max_ma` / `min_ma`.
        CORRECTION: an earlier version of this module assumed
        `early_age`/`late_age` — that was never verified and was WRONG.
        Fixed below; REQUIRED_FIELDS and occurrences_to_field() now use
        max_ma/min_ma.
      - `show=classext` returns `phylum`, `class`, `order`, `family`,
        `genus` as direct columns — used by data_pbdb_taxonomy.py to
        build the taxonomic tree from this SAME verified endpoint,
        rather than introducing a separate, unverified `/taxa/list` API.
      - `occurrence_no` (confirmed this session via multiple independent
        real-example sources: ropensci paleobioDB docs, CRAN README,
        GitHub) is the primary-key field present on every occurrence
        record — used as data_pbdb_taxonomy.py's TAXONOMY_REQUIRED_FIELDS
        instead of any specific taxonomic rank, since ranks are
        legitimately sparse per-record on real fossil data.
      - PBDB omits keys entirely for empty/null values rather than
        including them as null — confirmed via the same sources (a field
        can be genuinely present in the dataset but absent from any
        individual record, including record 0). _validate_record_schema()
        therefore checks a union of keys over a sample of up to 50
        records, not just the first one.
    Field names beyond these remain unconfirmed for other `show`/`vocab`
    combinations. This module still does NOT hardcode a fixed set of
    trusted field names and silently proceed if they're missing.
      Instead: discover_schema() must be run first, by a human, to see the
      ACTUAL fields returned; fetch_occurrences() then validates every
      declared REQUIRED_FIELD is present and raises SchemaValidationError
      (naming the fields that were actually present) if not. This is the
      same "validate the declared contract, never guess" pattern as
      LocalWellHDF5 in data_real.py.

Run discovery FIRST, before trusting anything else in this module:
    python -m src.data_pbdb --discover --base-name Dinosauria --limit 5
"""
from __future__ import annotations
import hashlib
import json
import os
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence

import torch
from torch.utils.data import Dataset

from .provenance import DataLoadError, SchemaValidationError

BASE_URL = "https://paleobiodb.org/data1.2/occs/list.json"
USER_AGENT = "well_poincare_rl-research-client/0.3 (contact: set via PBDB_CONTACT env var)"
POLITE_DELAY_SECONDS = 0.5  # spacing between paginated requests; no documented
                             # rate limit was found, so this is precautionary
                             # good-citizen behavior, not a confirmed requirement.
MAX_RETRIES = 3


REQUIRED_FIELDS = ("lng", "lat", "max_ma", "min_ma", "genus")


class PBDBFetchError(DataLoadError):
    outcome_code = "PBDB_FETCH_FAILED"


def _cache_key(params: Dict[str, Any]) -> str:
    blob = json.dumps(params, sort_keys=True).encode()
    return hashlib.sha256(blob).hexdigest()[:24]


def discover_schema(base_name: str = "Dinosauria", limit: int = 5,
                     show: str = "full") -> Dict[str, Any]:
    import requests
    params = {"base_name": base_name, "limit": limit, "show": show, "vocab": "pbdb"}
    resp = requests.get(BASE_URL, params=params,
                         headers={"User-Agent": USER_AGENT}, timeout=30)
    if resp.status_code != 200:
        raise PBDBFetchError(
            f"discovery request failed: HTTP {resp.status_code} — {resp.text[:300]}",
            outcome_code="PBDB_HTTP_ERROR",
        )
    data = resp.json()
    records = data.get("records", [])
    if not records:
        raise PBDBFetchError(
            f"discovery request returned zero records for base_name={base_name!r} "
            f"— check the taxon name, or the API contract may have changed.",
            outcome_code="PBDB_EMPTY_RESPONSE",
        )
    observed_fields = sorted(set().union(*(r.keys() for r in records)))
    return {
        "n_records": len(records),
        "observed_fields": observed_fields,
        "required_fields_present": [f for f in REQUIRED_FIELDS if f in observed_fields],
        "required_fields_missing": [f for f in REQUIRED_FIELDS if f not in observed_fields],
        "sample_record": records[0],
    }


def _validate_record_schema(records: List[Dict[str, Any]], required_fields: Sequence[str] = REQUIRED_FIELDS) -> None:
    if not records:
        raise SchemaValidationError(
            "fetch returned zero occurrence records", outcome_code="EMPTY_PBDB_RESULT",
        )

    sample = records[: min(50, len(records))]
    observed = set().union(*(r.keys() for r in sample))
    missing = [f for f in required_fields if f not in observed]
    if missing:
        raise SchemaValidationError(
            f"PBDB response is missing required field(s) {missing}. "
            f"Fields actually present (union of first {len(sample)} records): "
            f"{sorted(observed)}. "
            f"This means either the request params (show=/vocab=) need "
            f"adjusting, or PBDB's schema has changed since REQUIRED_FIELDS "
            f"was declared. Run discover_schema() to see the live response "
            f"before changing REQUIRED_FIELDS — never guess a fix.",
            outcome_code="PBDB_SCHEMA_MISMATCH",
        )


def fetch_occurrences(
    base_name: str,
    interval: Optional[str] = None,
    max_records: int = 5000,
    page_size: int = 500,
    show: str = "full",
    cache_dir: Optional[str] = "data/pbdb_cache",
    required_fields: Sequence[str] = REQUIRED_FIELDS,
) -> List[Dict[str, Any]]:

    import requests

    params_base = {"base_name": base_name, "show": show, "vocab": "pbdb"}
    if interval:
        params_base["interval"] = interval

    cache_path = None
    if cache_dir:
        os.makedirs(cache_dir, exist_ok=True)
        cache_key_params = dict(params_base, _required_fields=list(required_fields))
        cache_path = Path(cache_dir) / f"{_cache_key(cache_key_params)}_{max_records}.json"
        if cache_path.exists():
            records = json.loads(cache_path.read_text())
            _validate_record_schema(records, required_fields=required_fields)
            return records

    all_records: List[Dict[str, Any]] = []
    offset = 0
    while len(all_records) < max_records:
        params = dict(params_base)
        params["limit"] = min(page_size, max_records - len(all_records))
        params["offset"] = offset

        last_err = None
        for attempt in range(MAX_RETRIES):
            try:
                resp = requests.get(BASE_URL, params=params,
                                     headers={"User-Agent": USER_AGENT}, timeout=30)
                if resp.status_code == 200:
                    break
                last_err = f"HTTP {resp.status_code}: {resp.text[:200]}"
            except requests.RequestException as e:
                last_err = str(e)
            time.sleep(1.5 * (attempt + 1))
        else:
            raise PBDBFetchError(
                f"failed after {MAX_RETRIES} attempts at offset={offset}: {last_err}",
                outcome_code="PBDB_HTTP_ERROR",
            )

        page = resp.json().get("records", [])
        if not page:
            break  # exhausted the result set before hitting max_records
        all_records.extend(page)
        offset += len(page)
        if len(page) < params["limit"]:
            break  # short page => last page
        time.sleep(POLITE_DELAY_SECONDS)

    _validate_record_schema(all_records, required_fields=required_fields)

    if cache_path:
        cache_path.write_text(json.dumps(all_records))

    return all_records

def occurrences_to_field(
    records: List[Dict[str, Any]],
    n_time_bins: int = 12,
    height: int = 24,
    width: int = 24,
    age_min: Optional[float] = None,
    age_max: Optional[float] = None,
) -> torch.Tensor:
    ages, lats, lngs, genera = [], [], [], []
    for r in records:
        try:
            ea, la = float(r["max_ma"]), float(r["min_ma"])
            mid_age = (ea + la) / 2.0
            lng, lat = float(r["lng"]), float(r["lat"])
        except (KeyError, TypeError, ValueError):
            continue  # malformed individual record; skipped, not fabricated
        ages.append(mid_age)
        lats.append(lat)
        lngs.append(lng)
        genera.append(r.get("genus") or r.get("accepted_name") or "")

    if not ages:
        raise SchemaValidationError(
            "no occurrence records had usable age/coordinate fields after parsing",
            outcome_code="NO_USABLE_RECORDS",
        )

    age_max = age_max if age_max is not None else max(ages)
    age_min = age_min if age_min is not None else min(ages)
    if age_max <= age_min:
        raise SchemaValidationError(
            f"degenerate age range [{age_min}, {age_max}] — cannot bin into "
            f"{n_time_bins} time steps",
            outcome_code="DEGENERATE_AGE_RANGE",
        )

    field = torch.zeros(n_time_bins, 2, height, width)
    for age, lat, lng, genus in zip(ages, lats, lngs, genera):
        t = int((age_max - age) / (age_max - age_min) * (n_time_bins - 1e-6))
        t = min(max(t, 0), n_time_bins - 1)
        h = int((lat + 90.0) / 180.0 * (height - 1e-6))
        h = min(max(h, 0), height - 1)
        w = int((lng + 180.0) / 360.0 * (width - 1e-6))
        w = min(max(w, 0), width - 1)
        field[t, 0, h, w] += 1.0

    genus_sets = [[[set() for _ in range(width)] for _ in range(height)] for _ in range(n_time_bins)]
    for age, lat, lng, genus in zip(ages, lats, lngs, genera):
        if not genus:
            continue
        t = min(max(int((age_max - age) / (age_max - age_min) * (n_time_bins - 1e-6)), 0), n_time_bins - 1)
        h = min(max(int((lat + 90.0) / 180.0 * (height - 1e-6)), 0), height - 1)
        w = min(max(int((lng + 180.0) / 360.0 * (width - 1e-6)), 0), width - 1)
        genus_sets[t][h][w].add(genus)
    for t in range(n_time_bins):
        for h in range(height):
            for w in range(width):
                field[t, 1, h, w] = float(len(genus_sets[t][h][w]))

    return field


class PBDBFieldDataset(Dataset):
    def __init__(
        self,
        taxon_groups: Sequence[str],
        n_time_bins: int = 12,
        height: int = 24,
        width: int = 24,
        max_records_per_group: int = 3000,
        show: str = "full",
        cache_dir: Optional[str] = "data/pbdb_cache",
    ):
        self.taxon_groups = list(taxon_groups)
        self.fields: List[torch.Tensor] = []
        self.group_names: List[str] = []
        for group in self.taxon_groups:
            records = fetch_occurrences(
                base_name=group, max_records=max_records_per_group,
                show=show, cache_dir=cache_dir,
            )
            field = occurrences_to_field(records, n_time_bins=n_time_bins, height=height, width=width)
            self.fields.append(field)
            self.group_names.append(group)

        if not self.fields:
            raise DataLoadError(
                "no taxon groups produced usable fields", outcome_code="NO_PBDB_FIELDS",
            )

    def __len__(self):
        return len(self.fields)

    def __getitem__(self, idx):
        return {"fields": self.fields[idx], "idx": idx, "taxon_group": self.group_names[idx]}


DEFAULT_TAXON_GROUPS = (
    "Trilobita", "Ammonoidea", "Bivalvia", "Dinosauria",
    "Mammalia", "Brachiopoda", "Gastropoda", "Crinoidea",
)


def get_pbdb_dataset(
    taxon_groups: Sequence[str] = DEFAULT_TAXON_GROUPS,
    n_time_bins: int = 12,
    height: int = 24,
    width: int = 24,
    max_records_per_group: int = 3000,
    cache_dir: Optional[str] = "data/pbdb_cache",
):
    ds = PBDBFieldDataset(
        taxon_groups=taxon_groups, n_time_bins=n_time_bins, height=height,
        width=width, max_records_per_group=max_records_per_group, cache_dir=cache_dir,
    )
    print(f"[data] using REAL_PBDB data: {len(ds)} taxon-group trajectories "
          f"({', '.join(ds.group_names)})")
    return ds, "REAL_PBDB"


if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--discover", action="store_true",
                    help="Fetch a small sample and print the RAW schema for "
                         "human inspection. Run this first.")
    p.add_argument("--base-name", default="Dinosauria")
    p.add_argument("--limit", type=int, default=5)
    args = p.parse_args()

    if args.discover:
        result = discover_schema(base_name=args.base_name, limit=args.limit)
        print(json.dumps(result, indent=2, default=str))
        if result["required_fields_missing"]:
            print(f"\nWARNING: REQUIRED_FIELDS {result['required_fields_missing']} "
                  f"are NOT present in the live response. Update REQUIRED_FIELDS "
                  f"in src/data_pbdb.py to match reality before running fetch_occurrences().")
        else:
            print("\nAll REQUIRED_FIELDS confirmed present in the live response.")
    else:
        print("Pass --discover to inspect the live PBDB schema before trusting this module.")