#!/usr/bin/env python3
"""01_process.py — NOMAD 钙钛矿 JSONL → 清洗后的条目级与聚合级 CSV
数据源: NOMAD (nomad-lab.eu) 公开 DFT 数据库, structure_name=perovskite
许可: CC BY 4.0 (NOMAD 公开数据)
"""
import json
import re
from collections import Counter

import numpy as np
import pandas as pd

from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
RAW = ROOT / "data_raw" / "nomad_perovskite.jsonl"
OUT_ENTRIES = ROOT / "data" / "nomad_perovskite_entries.csv"
OUT_AGG = ROOT / "data" / "nomad_perovskite_aggregated.csv"
OUT_DIR = ROOT / "output"


def parse_band_gap(bg_list):
    """band_gap 为数组(每个 spin channel / 来源一项), 取最大值作为该条目带隙
    注意: NOMAD 存储单位为 J(焦耳), 此处换算为 eV"""
    if not bg_list:
        return np.nan, None
    vals, types = [], []
    for g in bg_list:
        v = g.get("value")
        if v is None:
            continue
        vals.append(float(v))
        types.append(g.get("type"))
    if not vals:
        return np.nan, None
    imax = int(np.argmax(vals))
    return vals[imax] / 1.602176634e-19, types[imax]  # J -> eV


def classify_functional(names):
    """将 xc_functional_names 数组归入大类"""
    if not names:
        return "unknown"
    s = " ".join(str(x) for x in names).lower()
    if "hf" in s or "hse" in s or "hybrid" in s:
        return "hybrid"
    if "gga" in s:
        if "u" in s and "-u" in s.replace("gga-u", "gga_u"):
            return "gga_u"
        return "gga"
    if "lda" in s:
        return "lda"
    if "meta" in s or "scan" in s or "tpss" in s:
        return "meta_gga"
    return "other"


FORMULA_RE = re.compile(r"([A-Z][a-z]?)(\d*\.?\d*)")


def parse_formula(formula):
    """解析化学式为 {元素: 化学计量} 字典"""
    if not formula or not isinstance(formula, str):
        return {}
    out = {}
    for el, num in FORMULA_RE.findall(formula):
        if not el:
            continue
        n = float(num) if num else 1.0
        out[el] = out.get(el, 0.0) + n
    return out


def main():
    if not Path(RAW).exists():
        if Path(OUT_ENTRIES).exists() and Path(OUT_AGG).exists():
            print(f"SKIP 01_process: raw JSONL missing ({RAW}), but cleaned CSVs already present in data/")
            return
        raise FileNotFoundError(f"Missing raw JSONL and cleaned CSVs: {RAW}")
    rows = []
    with open(RAW) as f:
        for line in f:
            try:
                e = json.loads(line)
            except json.JSONDecodeError:
                continue
            r = e.get("results", {}) or {}
            mat = r.get("material", {}) or {}
            props = (r.get("properties", {}) or {}).get("electronic", {}) or {}
            sym = mat.get("symmetry", {}) or {}
            meth = (r.get("method", {}) or {}).get("simulation", {}) or {}
            dft = meth.get("dft", {}) or {}

            bg_val, bg_type = parse_band_gap(props.get("band_gap"))
            rows.append(
                {
                    "entry_id": e.get("entry_id"),
                    "formula": mat.get("chemical_formula_reduced"),
                    "formula_hill": mat.get("chemical_formula_hill"),
                    "space_group": sym.get("space_group_number"),
                    "crystal_system": sym.get("crystal_system"),
                    "band_gap_eV": bg_val,
                    "band_gap_type": bg_type,
                    "functional_names": dft.get("xc_functional_names"),
                    "functional_class": classify_functional(dft.get("xc_functional_names")),
                }
            )
    df = pd.DataFrame(rows)
    print(f"原始条目: {len(df)}")
    print(f"含带隙数值条目: {df.band_gap_eV.notna().sum()}")
    print(f"唯一化学式(reduced): {df.formula.nunique()}")

    # ---- 条目级清洗 ----
    df = df.dropna(subset=["formula", "band_gap_eV"])
    df = df[(df.band_gap_eV >= 0) & (df.band_gap_eV <= 15)]  # 物理范围
    df["n_elements"] = df.formula.map(lambda s: len(parse_formula(s)))
    print(f"清洗后条目: {len(df)}")

    print("\n泛函分布:")
    print(df.functional_class.value_counts().to_string())
    print("\n带隙类型分布:")
    print(df.band_gap_type.value_counts(dropna=False).to_string())
    print("\n晶体系统分布:")
    print(df.crystal_system.value_counts().to_string())

    df.to_csv(OUT_ENTRIES, index=False)

    # ---- 聚合: (formula, functional_class) 中位数 ----
    agg = (
        df.groupby(["formula", "functional_class"])
        .agg(
            band_gap_eV=("band_gap_eV", "median"),
            band_gap_std=("band_gap_eV", "std"),
            n_entries=("band_gap_eV", "size"),
            space_group=("space_group", lambda s: s.mode().iloc[0] if len(s.mode()) else np.nan),
            crystal_system=("crystal_system", lambda s: s.mode().iloc[0] if len(s.mode()) else None),
            band_gap_type=("band_gap_type", lambda s: s.mode().iloc[0] if len(s.mode()) else None),
        )
        .reset_index()
    )
    agg["band_gap_std"] = agg["band_gap_std"].fillna(0)
    print(f"\n聚合后 (formula × functional_class) 组合: {len(agg)}")
    print(f"唯一化学式: {agg.formula.nunique()}")
    agg.to_csv(OUT_AGG, index=False)

    # 顶层泛函子集规模(用于选择训练集)
    top = agg.groupby("functional_class").size().sort_values(ascending=False)
    print("\n各泛函类聚合组数:")
    print(top.to_string())

    # 元素统计
    elem_counter = Counter()
    for s in agg.formula:
        for el in parse_formula(s):
            elem_counter[el] += 1
    print(f"\n涉及元素数: {len(elem_counter)}")
    print("Top-15 元素:", elem_counter.most_common(15))


if __name__ == "__main__":
    main()