#!/usr/bin/env python3
"""从 NOMAD API 分页拉取钙钛矿结构的 DFT 带隙数据"""
import json
import time
import urllib.request

API = "https://nomad-lab.eu/prod/v1/api/v1/entries/query"
from pathlib import Path
OUT = str(Path(__file__).resolve().parents[1] / "data_raw" / "nomad_perovskite.jsonl")

BODY = {
    "query": {"results.material.symmetry.structure_name": "perovskite"},
    "pagination": {"page_size": 1000},
    "required": {
        "include": [
            "results.material.chemical_formula_reduced",
            "results.material.chemical_formula_hill",
            "results.material.symmetry.space_group_number",
            "results.material.symmetry.crystal_system",
            "results.properties.electronic.band_gap.value",
            "results.properties.electronic.band_gap.type",
            "results.method.simulation.dft.xc_functional_names",
            "entry_id",
        ],
        "statistics": False,
    },
}

n = 0
after = None
t0 = time.time()
with open(OUT, "w") as f:
    for page in range(60):  # 最多 6 万条
        body = json.loads(json.dumps(BODY))
        if after:
            body["pagination"]["page_after_value"] = after
        req = urllib.request.Request(
            API, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}
        )
        for attempt in range(3):
            try:
                resp = urllib.request.urlopen(req, timeout=120)
                data = json.loads(resp.read())
                break
            except Exception as e:
                print(f"page {page} attempt {attempt} error: {e}", flush=True)
                time.sleep(5)
        else:
            print("连续失败，停止", flush=True)
            break
        entries = data.get("data", [])
        for e in entries:
            f.write(json.dumps(e) + "\n")
        n += len(entries)
        pag = data.get("pagination", {})
        after = pag.get("next_page_after_value")
        print(f"page {page}: +{len(entries)} (total {n}) elapsed {time.time()-t0:.0f}s", flush=True)
        if not after or not entries:
            break
print(f"DONE: {n} entries -> {OUT}", flush=True)