from __future__ import annotations

import json
import multiprocessing as mp
import os
import shutil
import struct
import subprocess
from collections import deque
from concurrent.futures import ProcessPoolExecutor
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import h5py
import numpy as np
try:
    from tqdm import tqdm
except ImportError:
    def tqdm(iterable=None, **kwargs):
        if iterable is None:
            class _NoopTqdm:
                def __init__(self, total=None, **_):
                    self.total = total

                def update(self, _=1):
                    pass

                def close(self):
                    pass

            return _NoopTqdm(**kwargs)
        return iterable

from src.common.config import Config
from src.common.experiment_paths import normalize_tag
from src.data.curve_processing import clean_curve_for_dataset, is_valid_curve, resample_curve_to_features_with_time
from src.data.params import Params, Schedule, generate_params_dataset, write_params_bin
from src.data.schedule_encoding import encode_schedule_to_timegrid

_WORKER_CFG_PATH: Optional[str] = None
_WORKDIR: Optional[Path] = None
_PARAMS_BIN: Optional[Path] = None
_RESULT_BIN: Optional[Path] = None
_SUBPROC_ENV: Optional[dict] = None


"""生成正演代理模型训练使用的 HDF5 样本。

每个样本由“物理参数采样 + 流量制度采样 + C++ 数值求解器正演”生成，
随后将求解器输出的双对数曲线转换为固定长度的模型输入和监督目标。
"""


def _seed32(x: int) -> int:
    return int(x) & 0xFFFFFFFF


def _read_result_bin(result_bin_path: Path) -> Optional[Dict[str, Any]]:
    """读取 runner.exe 输出的紧凑二进制结果文件。"""
    try:
        with open(result_bin_path, "rb") as f:
            magic, version = struct.unpack("<II", f.read(8))
            expected_magic = ord("R") | (ord("S") << 8) | (ord("B") << 16) | (ord("1") << 24)
            if magic != expected_magic or version != 1:
                return None

            nWells, nSteps = struct.unpack("<II", f.read(8))
            t = np.fromfile(f, dtype="<f8", count=nSteps).astype(np.float64).tolist()
            pw = [np.fromfile(f, dtype="<f8", count=nSteps).astype(np.float64).tolist() for _ in range(nWells)]

            loglog = []
            for _ in range(nWells):
                (nLogLog,) = struct.unpack("<I", f.read(4))
                loglog_t = np.fromfile(f, dtype="<f8", count=nLogLog).astype(np.float64).tolist()
                loglog_p = np.fromfile(f, dtype="<f8", count=nLogLog).astype(np.float64).tolist()
                loglog_deriv = np.fromfile(f, dtype="<f8", count=nLogLog).astype(np.float64).tolist()
                loglog.append({"t": loglog_t, "p": loglog_p, "deriv": loglog_deriv})

            return {"nWells": nWells, "nSteps": nSteps, "t": t, "pw": pw, "loglog": loglog}
    except Exception:
        return None


def _pick_mixture(rng, items):
    probs = np.array([float(it.get("prob", 0.0)) for it in items], dtype=np.float64)
    s = float(np.sum(probs))
    if s <= 0:
        return items[int(rng.randint(0, len(items)))]
    probs = probs / s

    u = rng.rand()
    c = 0.0
    for it, p in zip(items, probs):
        c += float(p)
        if u <= c:
            return it
    return items[-1]


def _normalize_durations_to_total(dt: np.ndarray, total: float, min_dt: float) -> np.ndarray:
    dt = np.maximum(dt.astype(np.float64), float(min_dt))
    s = float(np.sum(dt))
    total = max(float(total), float(min_dt) * float(len(dt)))
    if s <= 0:
        return np.full_like(dt, total / len(dt))
    dt = dt * (total / s)
    dt = np.maximum(dt, float(min_dt))
    diff = total - float(np.sum(dt))
    dt[-1] += diff
    return np.maximum(dt, float(min_dt))


SCHEDULE_META_NAMES = [
    "family_id",
    "section_index",
    "n_sections",
    "n_prod",
    "prod_total_time",
    "shutin_dt",
    "q_first_prod",
    "q_last_prod",
    "q_mean_prod",
    "q_std_prod",
    "q_min_prod",
    "q_max_prod",
    "q_end_start_ratio",
    "q_range_ratio",
    "max_step_ratio",
    "mean_step_ratio",
    "trend_slope",
    "is_nonincreasing",
    "is_strong_decline",
]


def _family_id_map(cfg: Config) -> dict[str, int]:
    mode = str(cfg.raw["schedule"]["generation_mode"]).lower()
    if mode == "family_random":
        families = cfg.raw["schedule"]["family_random"]["families"]
        mapping = {str(item.get("name", f"family_{i}")).lower(): i for i, item in enumerate(families)}
        mapping.setdefault("unknown", -1)
        return mapping
    return {
        "fixed_case": 0,
        "case_neighborhood": 1,
        "unknown": -1,
    }


def build_schedule_metadata(
    cfg: Config,
    timeQ: List[float],
    q: List[float],
    family_name: str,
    section_index: int,
) -> tuple[np.ndarray, str]:
    """构造描述流量制度形态的分析元数据。"""
    dt = np.asarray(timeQ, dtype=np.float64).reshape(-1)
    qq = np.asarray(q, dtype=np.float64).reshape(-1)
    n_sections = int(len(dt))
    q_thr = float((cfg.raw["schedule"].get("canonicalize_for_model", {}) or {}).get("q_thr", 1e-6))

    has_shutin = bool(n_sections > 0 and qq[-1] <= q_thr)
    n_prod = int(max(n_sections - 1, 0) if has_shutin else n_sections)
    prod_dt = dt[:n_prod] if n_prod > 0 else np.zeros((0,), dtype=np.float64)
    prod_q = qq[:n_prod] if n_prod > 0 else np.zeros((0,), dtype=np.float64)
    shutin_dt = float(dt[-1]) if has_shutin else 0.0
    q_eps = 1.0e-12

    if n_prod > 0:
        q_first = float(prod_q[0])
        q_last = float(prod_q[-1])
        q_min = float(np.min(prod_q))
        q_max = float(np.max(prod_q))
        q_end_start_ratio = q_last / max(abs(q_first), q_eps)
        q_range_ratio = (q_max - q_min) / max(abs(q_max), q_eps)
    else:
        q_first = q_last = q_min = q_max = 0.0
        q_end_start_ratio = 0.0
        q_range_ratio = 0.0

    if n_prod > 1:
        prev_q = prod_q[:-1]
        next_q = prod_q[1:]
        step_ratio = np.maximum(prev_q, next_q) / np.maximum(np.minimum(prev_q, next_q), q_eps)
        max_step_ratio = float(np.max(step_ratio))
        mean_step_ratio = float(np.mean(step_ratio))
        is_nonincreasing = bool(np.all(next_q <= prev_q * (1.0 + 1.0e-9)))
        has_drop = bool(np.any(next_q < prev_q * (1.0 - 1.0e-9)))
        x = np.cumsum(prod_dt) - 0.5 * prod_dt
        trend_slope = float(np.polyfit(x, prod_q, deg=1)[0]) if np.ptp(x) > q_eps else 0.0
    else:
        max_step_ratio = 1.0
        mean_step_ratio = 1.0
        is_nonincreasing = False
        has_drop = False
        trend_slope = 0.0

    is_strong_decline = bool(is_nonincreasing and has_drop and q_end_start_ratio < 0.70)

    family_name = str(family_name).lower()
    family_id = int(_family_id_map(cfg).get(family_name, -1))

    meta_vec = np.asarray(
        [
            float(family_id),
            float(section_index),
            float(n_sections),
            float(n_prod),
            float(np.sum(prod_dt)) if n_prod > 0 else 0.0,
            float(shutin_dt),
            q_first,
            q_last,
            float(np.mean(prod_q)) if n_prod > 0 else 0.0,
            float(np.std(prod_q)) if n_prod > 0 else 0.0,
            q_min,
            q_max,
            q_end_start_ratio,
            q_range_ratio,
            max_step_ratio,
            mean_step_ratio,
            trend_slope,
            float(is_nonincreasing),
            float(is_strong_decline),
        ],
        dtype=np.float32,
    )
    return meta_vec, family_name


def _sample_schedule_fixed_case(cfg: Config, rng) -> Tuple[List[float], List[float], Dict[str, Any]]:
    sc = cfg.raw["schedule"]["case_schedule"]
    return (
        list(map(float, sc["timeQ"])),
        list(map(float, sc["q"])),
        {"family_name": "fixed_case"},
    )


def _sample_schedule_case_neighborhood(cfg: Config, rng) -> Tuple[List[float], List[float], Dict[str, Any]]:
    base = cfg.raw["schedule"]["case_schedule"]
    ncfg = cfg.raw["schedule"]["case_neighborhood"]

    base_timeQ = np.asarray(base["timeQ"], dtype=np.float64)
    base_q = np.asarray(base["q"], dtype=np.float64)
    out_t = base_timeQ.copy()
    out_q = base_q.copy()
    prod_n = len(base_timeQ) - 1

    noise_t = rng.uniform(-float(ncfg["dt_jitter_rel"]), float(ncfg["dt_jitter_rel"]), size=prod_n)
    out_t[:prod_n] = np.maximum(base_timeQ[:prod_n] * (1.0 + noise_t), float(ncfg["min_dt"]))
    tail_noise = rng.uniform(-float(ncfg["shutin_dt_jitter_rel"]), float(ncfg["shutin_dt_jitter_rel"]))
    out_t[-1] = max(base_timeQ[-1] * (1.0 + tail_noise), float(ncfg["min_dt"]))

    noise_q = rng.uniform(-float(ncfg["q_jitter_rel"]), float(ncfg["q_jitter_rel"]), size=prod_n)
    out_q[:prod_n] = np.clip(base_q[:prod_n] * (1.0 + noise_q), float(ncfg["q_min"]), float(ncfg["q_max"]))

    if bool(ncfg["keep_monotonic_prod"]) and prod_n > 1:
        out_q[:prod_n] = np.maximum.accumulate(out_q[:prod_n])
    if bool(ncfg["keep_last_q_zero"]):
        out_q[-1] = 0.0

    return out_t.tolist(), out_q.tolist(), {"family_name": "case_neighborhood"}


def _sample_schedule_family_random(cfg: Config, rng) -> Tuple[List[float], List[float], Dict[str, Any]]:
    """随机抽取一种流量制度族，并生成对应的流动段时长和流量。"""
    fcfg = cfg.raw["schedule"]["family_random"]
    fam = _pick_mixture(rng, fcfg["families"])
    fam_name = str(fam.get("name", "inc_tail_shutin")).lower()

    n_lo, n_hi = fcfg["n_prod_sections_range"]
    n_prod = int(rng.randint(int(n_lo), int(n_hi) + 1))

    prod_total_lo, prod_total_hi = fcfg["prod_total_time_range"]
    prod_total = float(prod_total_lo + rng.rand() * (prod_total_hi - prod_total_lo))

    mu = float(fcfg["duration_lognormal_mu"])
    sigma = float(fcfg["duration_lognormal_sigma"])
    dt_prod = rng.lognormal(mean=mu, sigma=sigma, size=n_prod).astype(np.float64)
    dt_prod = _normalize_durations_to_total(dt_prod, total=prod_total, min_dt=0.05)

    q_lo, q_hi = fcfg["q_range"]
    q0 = float(q_lo + rng.rand() * (q_hi - q_lo))
    max_rel_step = float(fcfg["max_rel_step"])
    mult_noise_sigma = float(fcfg["mult_noise_sigma"])
    step_jump_lo, step_jump_hi = fcfg["step_jump_rel_range"]

    q_prod = np.zeros((n_prod,), dtype=np.float64)
    q_prod[0] = q0

    if fam_name == "inc_tail_shutin":
        for i in range(1, n_prod):
            q_prod[i] = q_prod[i - 1] * rng.uniform(1.02, max_rel_step)
        q_prod = np.maximum.accumulate(q_prod)
    elif fam_name == "dec_tail_shutin":
        for i in range(1, n_prod):
            q_prod[i] = max(q_prod[i - 1] * rng.uniform(1.0 / max_rel_step, 0.98), q_lo)
        q_prod = np.minimum.accumulate(q_prod)
    elif fam_name == "strong_dec_tail_shutin":
        end_ratio_lo, end_ratio_hi = fcfg.get("strong_dec_end_start_ratio_range", [0.25, 0.65])
        q_end = max(q_lo, q_prod[0] * rng.uniform(float(end_ratio_lo), float(end_ratio_hi)))
        shape = np.linspace(0.0, 1.0, n_prod)
        q_prod = q_prod[0] * (q_end / max(q_prod[0], 1.0e-12)) ** shape
        q_prod = np.minimum.accumulate(q_prod)
    elif fam_name == "mild_step_tail_shutin":
        q_prod[:] = q0
        jump_idx = int(rng.randint(1, max(2, n_prod)))
        q_prod[jump_idx:] *= float(rng.uniform(step_jump_lo, step_jump_hi))
    else:
        for i in range(1, n_prod):
            rel = rng.uniform(1.0 / max_rel_step, max_rel_step)
            rel = 1.0 + 0.15 * (rel - 1.0)
            q_prod[i] = q_prod[i - 1] * rel

    if mult_noise_sigma > 0:
        q_prod *= np.exp(rng.normal(loc=0.0, scale=mult_noise_sigma, size=n_prod))

    q_prod = np.clip(q_prod, q_lo, q_hi)

    shut_lo, shut_hi = fcfg["shutin_dt_range"]
    dt_shut = float(shut_lo + rng.rand() * (shut_hi - shut_lo))

    return dt_prod.tolist() + [dt_shut], q_prod.tolist() + [0.0], {"family_name": fam_name}


def sample_schedule_by_mode(cfg: Config, rng) -> Tuple[List[float], List[float], Dict[str, Any]]:
    """根据配置中的 generation_mode 分发到对应的流量制度采样方法。"""
    mode = str(cfg.raw["schedule"]["generation_mode"]).lower()
    if mode == "fixed_case":
        return _sample_schedule_fixed_case(cfg, rng)
    if mode == "case_neighborhood":
        return _sample_schedule_case_neighborhood(cfg, rng)
    if mode == "family_random":
        return _sample_schedule_family_random(cfg, rng)
    raise ValueError(f"Unknown schedule generation_mode: {mode}")


def _resolve_section_indices(cfg: Config, timeQ, q, rng):
    policy = cfg.raw["schedule"]["section_policy"]
    mode = str(policy["mode"]).lower()
    n = int(len(timeQ))
    if mode == "fixed_last":
        return [n]
    if mode == "fixed_value":
        return [int(np.clip(int(policy["fixed_value"]), 1, n))]
    if mode == "all_sections":
        return list(range(1, n + 1))
    if mode == "uniform_one":
        return [int(rng.randint(1, n + 1))]
    raise ValueError(f"Unknown section_policy.mode: {mode}")


def _init_process_worker(cfg_path: str):
    global _WORKER_CFG_PATH, _WORKDIR, _PARAMS_BIN, _RESULT_BIN, _SUBPROC_ENV

    _WORKER_CFG_PATH = cfg_path
    cfg = Config(cfg_path)
    cfg.ensure_dirs()

    base_worker_dir = cfg.paths.temp_dir / "parallel_workers"
    pid = os.getpid()
    _WORKDIR = base_worker_dir / f"proc_{pid}"
    _WORKDIR.mkdir(parents=True, exist_ok=True)

    _PARAMS_BIN = _WORKDIR / "params.bin"
    _RESULT_BIN = _WORKDIR / "result.bin"

    env = os.environ.copy()
    env.setdefault("OMP_NUM_THREADS", "1")
    env.setdefault("OPENBLAS_NUM_THREADS", "1")
    env.setdefault("MKL_NUM_THREADS", "1")
    env.setdefault("NUMEXPR_NUM_THREADS", "1")
    _SUBPROC_ENV = env


def _worker_simulate_parallel(args):
    """子进程侧仿真：调用求解器并返回一个编码后的有效样本。"""
    task_idx, params_dict = args
    cfg = Config(_WORKER_CFG_PATH)

    try:
        p = Params(
            k=float(params_dict["k"]),
            skin=float(params_dict["skin"]),
            wellboreC=float(params_dict["wellboreC"]),
            phi=float(params_dict["phi"]),
            h=float(params_dict["h"]),
            Cf=float(params_dict["Cf"]),
            schedule=None,
        )

        sc = params_dict["schedule"]
        sch = Schedule(
            sectionIndex=int(sc["sectionIndex"]),
            timeQ=list(map(float, sc["timeQ"])),
            q=list(map(float, sc["q"])),
        )
        p.schedule = sch

        write_params_bin(str(_PARAMS_BIN), p, cfg=cfg, include_schedule=True)

        if _RESULT_BIN.exists():
            _RESULT_BIN.unlink()

        cmd = [str(cfg.runner_exe), str(cfg.dataset_bin), str(_PARAMS_BIN), str(_RESULT_BIN)]
        result = subprocess.run(
            cmd,
            cwd=str(cfg.runner_exe.parent),
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
            timeout=int(cfg.get("dataset_runtime", "runner_timeout_sec", default=240)),
            encoding="utf-8",
            errors="ignore",
            env=_SUBPROC_ENV,
        )

        if result.returncode != 0 or not _RESULT_BIN.exists():
            return task_idx, None, "runner_failed", None

        result_dict = _read_result_bin(_RESULT_BIN)
        if not result_dict or not result_dict["loglog"] or not result_dict["loglog"][0]:
            return task_idx, None, "empty_loglog", None

        t = np.asarray(result_dict["loglog"][0]["t"], dtype=np.float64)
        p_curve = np.asarray(result_dict["loglog"][0]["p"], dtype=np.float64)
        d_curve = np.asarray(result_dict["loglog"][0]["deriv"], dtype=np.float64)

        cleaned = clean_curve_for_dataset(cfg, t, p_curve, d_curve)
        if cleaned is None:
            return task_idx, None, "too_short_after_clean", None

        t, p_curve, d_curve = cleaned
        ok, reason = is_valid_curve(cfg, t, p_curve, d_curve)
        if not ok:
            return task_idx, None, reason, None

        curve_feat, curve_time = resample_curve_to_features_with_time(cfg, t, p_curve, d_curve)

        sec = int(np.clip(int(sch.sectionIndex), 1, max(len(sch.timeQ), 1)))
        enc = encode_schedule_to_timegrid(
            cfg,
            sectionIndex=sec,
            timeQ=sch.timeQ,
            q=sch.q,
            n_sections=len(sch.timeQ),
        )

        params_vec = np.asarray(
            [p.k, p.skin, p.wellboreC, p.phi, p.h, p.Cf],
            dtype=np.float32,
        )

        schedule_parts = [
            np.asarray(enc.x_sched, dtype=np.float32).reshape(-1),
            np.asarray(enc.x_sec, dtype=np.float32).reshape(-1),
        ]

        if bool(cfg.raw.get("schedule", {}).get("use_metadata_features_for_model", False)):
            names = cfg.raw.get("schedule", {}).get("metadata_features_for_model", []) or []
            name_to_idx = {name: i for i, name in enumerate(SCHEDULE_META_NAMES)}
            meta_src = np.asarray(params_dict["schedule_meta"]["meta_vec"], dtype=np.float32).reshape(-1)
            meta_values = []
            for name in names:
                idx = name_to_idx.get(str(name))
                meta_values.append(float(meta_src[idx]) if idx is not None and idx < meta_src.size else 0.0)
            schedule_parts.append(np.asarray(meta_values, dtype=np.float32).reshape(-1))

        schedule_vec = np.concatenate(
            [
                *schedule_parts,
            ],
            axis=0,
        ).astype(np.float32)

        sample = {
            "group_id": int(params_dict.get("schedule_id", -1)),
            "params": params_vec,
            "schedule": schedule_vec,
            "curve": np.asarray(curve_feat, dtype=np.float32),
            "curve_time": np.asarray(curve_time, dtype=np.float32),
            "schedule_meta": np.asarray(params_dict["schedule_meta"]["meta_vec"], dtype=np.float32),
            "family_name": str(params_dict["schedule_meta"]["family_name"]),
        }
        return task_idx, sample, "", None

    except subprocess.TimeoutExpired:
        return task_idx, None, "timeout", None
    except Exception as e:
        return task_idx, None, f"exception_{type(e).__name__}", None


class HDF5Appender:
    """将生成样本追加写入 HDF5，避免把完整数据集放入内存。"""

    def __init__(
        self,
        filepath: Path,
        param_dim: int,
        schedule_dim: int,
        curve_dim: int,
        compression=None,
        chunk_rows: int = 2048,
    ):
        self.filepath = Path(filepath)
        self.filepath.parent.mkdir(parents=True, exist_ok=True)

        self.param_dim = int(param_dim)
        self.schedule_dim = int(schedule_dim)
        self.curve_dim = int(curve_dim)
        self.time_dim = int(curve_dim) // 3

        self.f = h5py.File(self.filepath, "w")
        self._n = 0
        self.f.attrs["schedule_meta_names"] = np.asarray(SCHEDULE_META_NAMES, dtype="S")

        def ds(name, shape_tail, dtype, fillvalue=None):
            # 所有数据集都在第一维支持扩展，用于流式写入样本。
            return self.f.create_dataset(
                name,
                shape=(0, *shape_tail),
                maxshape=(None, *shape_tail),
                chunks=(chunk_rows, *shape_tail),
                dtype=dtype,
                compression=compression,
                fillvalue=fillvalue,
            )

        self.d_group_id = ds("group_id", (), np.int64, fillvalue=-1)
        self.d_params = ds("params", (self.param_dim,), np.float32, fillvalue=np.nan)
        self.d_schedule = ds("schedule", (self.schedule_dim,), np.float32, fillvalue=np.nan)
        self.d_curve = ds("curve", (self.curve_dim,), np.float32, fillvalue=np.nan)
        self.d_curve_time = ds("curve_time", (self.time_dim,), np.float32, fillvalue=np.nan)
        self.d_schedule_meta = ds("schedule_meta", (len(SCHEDULE_META_NAMES),), np.float32, fillvalue=np.nan)
        self.d_family_name = self.f.create_dataset(
            "family_name",
            shape=(0,),
            maxshape=(None,),
            chunks=(chunk_rows,),
            dtype=h5py.string_dtype(encoding="utf-8"),
        )

    @property
    def n_samples(self) -> int:
        return int(self._n)

    def append_samples(self, samples: List[Dict[str, Any]]):
        """追加一批有效样本，并更新样本数量属性。"""
        if not samples:
            return

        B = len(samples)
        start, end = self._n, self._n + B

        self.d_group_id.resize((end,))
        self.d_params.resize((end, self.param_dim))
        self.d_schedule.resize((end, self.schedule_dim))
        self.d_curve.resize((end, self.curve_dim))
        self.d_curve_time.resize((end, self.time_dim))
        self.d_schedule_meta.resize((end, len(SCHEDULE_META_NAMES)))
        self.d_family_name.resize((end,))

        gid = np.full((B,), -1, dtype=np.int64)
        params = np.full((B, self.param_dim), np.nan, dtype=np.float32)
        schedule = np.full((B, self.schedule_dim), np.nan, dtype=np.float32)
        curve = np.full((B, self.curve_dim), np.nan, dtype=np.float32)
        curve_time = np.full((B, self.time_dim), np.nan, dtype=np.float32)
        schedule_meta = np.full((B, len(SCHEDULE_META_NAMES)), np.nan, dtype=np.float32)
        family_name: list[str] = ["" for _ in range(B)]

        for i, s in enumerate(samples):
            gid[i] = int(s.get("group_id", -1))
            params[i] = np.asarray(s["params"], dtype=np.float32).reshape(-1)[: self.param_dim]
            schedule[i] = np.asarray(s["schedule"], dtype=np.float32).reshape(-1)[: self.schedule_dim]
            curve[i] = np.asarray(s["curve"], dtype=np.float32).reshape(-1)[: self.curve_dim]
            if "curve_time" in s:
                curve_time[i] = np.asarray(s["curve_time"], dtype=np.float32).reshape(-1)[: self.time_dim]
            schedule_meta[i] = np.asarray(s["schedule_meta"], dtype=np.float32).reshape(-1)[: len(SCHEDULE_META_NAMES)]
            family_name[i] = str(s.get("family_name", ""))

        self.d_group_id[start:end] = gid
        self.d_params[start:end] = params
        self.d_schedule[start:end] = schedule
        self.d_curve[start:end] = curve
        self.d_curve_time[start:end] = curve_time
        self.d_schedule_meta[start:end] = schedule_meta
        self.d_family_name[start:end] = family_name

        self._n = end
        self.f.attrs["n_samples"] = int(self._n)

    def flush(self):
        self.f.flush()

    def close(self):
        self.flush()
        self.f.close()


class ParallelDatasetGenerator:
    """大规模求解器生成数据集的并行调度层。"""

    def __init__(self, cfg: Config, n_workers: int | None = None):
        self.cfg = cfg
        self.cfg.ensure_dirs()
        self.n_workers = int(
            n_workers or cfg.get("parallel", "n_workers", default=max(1, int(mp.cpu_count() * 0.75)))
        )
        self.output_dir = cfg.paths.samples_dir

    def _generate_dataset(self) -> bool:
        """缺少 dataset.bin 时调用 C++ training.exe 自动生成。"""
        try:
            result = subprocess.run(
                [str(self.cfg.training_exe)],
                cwd=str(self.cfg.training_exe.parent),
                timeout=60,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
                encoding="utf-8",
                errors="ignore",
            )
            if result.returncode == 0 and self.cfg.dataset_bin.exists():
                return True

            fallback = self.cfg.training_exe.parent / "dataset.bin"
            if result.returncode == 0 and fallback.exists():
                self.cfg.dataset_bin.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(fallback, self.cfg.dataset_bin)
                return True

            return False
        except Exception:
            return False

    def _cleanup_workers(self):
        base_worker_dir = self.cfg.paths.temp_dir / "parallel_workers"
        if base_worker_dir.exists():
            shutil.rmtree(base_worker_dir, ignore_errors=True)

    def generate(self, n_samples=None, method=None, random_seed=None, dataset_tag: str | None = None):
        """生成一份完整的正演代理模型训练 HDF5 数据集。"""
        cfg = self.cfg
        n_samples = int(n_samples or cfg.get("generation", "n_samples", default=50000))
        method = method or cfg.raw["params"].get("sampling_method", "sobol")
        random_seed = int(random_seed or cfg.get("generation", "random_seed", default=42))

        if not cfg.dataset_bin.exists():
            if bool(cfg.get("dataset_runtime", "auto_build_dataset_bin", default=True)):
                if not self._generate_dataset():
                    raise RuntimeError("无法生成 dataset.bin")
            else:
                raise RuntimeError(f"dataset.bin 不存在: {cfg.dataset_bin}")

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        mode = str(cfg.raw["schedule"]["generation_mode"]).lower()
        sec_mode = str(cfg.raw["schedule"]["section_policy"]["mode"]).lower()
        tag = normalize_tag(dataset_tag)
        if tag:
            filepath = self.output_dir / f"dataset_{tag}_target{n_samples}_{timestamp}.h5"
        else:
            filepath = self.output_dir / f"dataset_{mode}_{sec_mode}_target{n_samples}_{timestamp}.h5"

        curve_dim = cfg.curve_dim
        Nu, Cu = cfg.schedule_grid_shape
        sched_dim = Nu * Cu + cfg.sec_feat_dim
        if bool(cfg.raw.get("schedule", {}).get("use_metadata_features_for_model", False)):
            sched_dim += len(cfg.raw.get("schedule", {}).get("metadata_features_for_model", []) or [])
        param_dim = 6

        # HDF5 布局与正演模型输入输出一致：
        # params/schedule 是输入，curve 是数值求解器生成的监督目标。
        write_batch_size = int(cfg.get("streaming_hdf5", "write_batch_size", default=2000))
        compression = cfg.get("streaming_hdf5", "compression", default=None)
        chunk_rows = int(cfg.get("streaming_hdf5", "chunk_rows", default=2048))

        app = HDF5Appender(
            filepath=filepath,
            param_dim=param_dim,
            schedule_dim=sched_dim,
            curve_dim=curve_dim,
            compression=compression,
            chunk_rows=chunk_rows,
        )

        successful = 0
        total_requested = 0
        fail_reasons = {}
        fail_examples = {}
        max_fail_examples_per_reason = int(cfg.get("generation", "max_fail_examples_per_reason", default=50))
        max_rounds = int(cfg.get("generation", "max_rounds", default=999999))
        max_total_requested = int(cfg.get("generation", "max_total_requested", default=int(n_samples * 12.0)))

        checkpoint_every_n = int(cfg.get("parallel", "checkpoint_every_n", default=5000))
        checkpoint_every_sec = float(cfg.get("parallel", "checkpoint_every_sec", default=180.0))
        last_ckpt_n = 0
        last_ckpt_ts = datetime.now().timestamp()

        seed_base = int(random_seed)
        round_id = 0
        param_keys = cfg.raw["params"]["all_physical_param_names"]
        expand_ma = 1.0
        alpha = 0.2
        expand_stats = []
        buffer_samples = []

        def flush_buffer(tag: str):
            nonlocal buffer_samples, last_ckpt_n, last_ckpt_ts
            if not buffer_samples:
                return

            # 写入进度元数据，便于长时间生成任务中断后检查当前状态。
            app.append_samples(buffer_samples)
            buffer_samples = []
            app.flush()

            meta = {
                "is_partial": True,
                "tag": tag,
                "n_saved": int(app.n_samples),
                "round_id": int(round_id),
                "total_requested": int(total_requested),
                "generated_at": datetime.now().isoformat(),
                "generation_mode": mode,
                "section_policy": sec_mode,
            }
            with open(filepath.with_suffix(".progress.json"), "w", encoding="utf-8") as f:
                json.dump(meta, f, indent=2, ensure_ascii=False)

            last_ckpt_n = int(app.n_samples)
            last_ckpt_ts = datetime.now().timestamp()

        try:
            pbar = tqdm(total=n_samples, desc="valid samples", unit="sample")

            while successful < n_samples:
                round_id += 1
                if round_id > max_rounds or total_requested >= max_total_requested:
                    break

                need = n_samples - successful
                request_params = int(np.ceil(need / max(expand_ma, 1.0) * 1.15))
                request_params = max(1, min(request_params, 8000))

                params_list = generate_params_dataset(
                    cfg,
                    request_params,
                    method=method,
                    random_seed=seed_base + round_id * 1000,
                )

                tasks = []
                base_seed = seed_base + round_id * 1000
                task_counter = 0

                for pidx, p in enumerate(params_list):
                    pd_base = {k: float(getattr(p, k)) for k in param_keys}
                    schedule_id = int((round_id << 32) ^ (pidx << 1) ^ (base_seed & 0xFFFFFFFF))
                    rng_i = np.random.RandomState(_seed32(base_seed ^ (pidx + 12345)))
                    timeQ, q, sched_info = sample_schedule_by_mode(cfg, rng_i)
                    section_indices = _resolve_section_indices(cfg, timeQ, q, rng_i)

                    for sec in section_indices:
                        pd = dict(pd_base)
                        pd["schedule_id"] = schedule_id
                        pd["schedule"] = {
                            "sectionIndex": int(sec),
                            "timeQ": timeQ,
                            "q": q,
                        }
                        meta_vec, family_name = build_schedule_metadata(
                            cfg=cfg,
                            timeQ=timeQ,
                            q=q,
                            family_name=str(sched_info.get("family_name", "unknown")),
                            section_index=int(sec),
                        )
                        pd["schedule_meta"] = {
                            "meta_vec": meta_vec.tolist(),
                            "family_name": family_name,
                        }
                        tasks.append((task_counter, pd))
                        task_counter += 1

                total_requested += len(tasks)
                expand_now = len(tasks) / max(len(params_list), 1)
                expand_ma = (1 - alpha) * expand_ma + alpha * expand_now
                expand_stats.append(
                    {
                        "round": round_id,
                        "expand_now": float(expand_now),
                        "expand_ma": float(expand_ma),
                    }
                )

                reached = False
                max_in_flight = int(cfg.get("parallel", "max_in_flight", default=self.n_workers * 4))

                with ProcessPoolExecutor(
                    max_workers=self.n_workers,
                    initializer=_init_process_worker,
                    initargs=(str(cfg.path),),
                ) as executor:
                    task_iter = iter(tasks)
                    in_flight = deque()

                    for _ in range(min(max_in_flight, len(tasks))):
                        try:
                            task = next(task_iter)
                            in_flight.append(executor.submit(_worker_simulate_parallel, task))
                        except StopIteration:
                            break

                    while in_flight:
                        fut = in_flight.popleft()
                        try:
                            task_idx, sample, fail_reason, fail_ctx = fut.result()
                        except Exception:
                            fail_reasons["future_exception"] = fail_reasons.get("future_exception", 0) + 1
                        else:
                            if (not reached) and (sample is not None):
                                buffer_samples.append(sample)
                                successful += 1
                                pbar.update(1)

                                if len(buffer_samples) >= write_batch_size:
                                    flush_buffer(f"batch_R{round_id}")

                                if successful >= n_samples:
                                    reached = True

                            elif (not reached) and (sample is None) and fail_reason:
                                fail_reasons[fail_reason] = fail_reasons.get(fail_reason, 0) + 1
                                if len(fail_examples.get(fail_reason, [])) < max_fail_examples_per_reason:
                                    fail_examples.setdefault(fail_reason, []).append(
                                        {
                                            "task_idx": int(task_idx),
                                            "reason": str(fail_reason),
                                        }
                                    )

                        now = datetime.now().timestamp()
                        if ((app.n_samples + len(buffer_samples)) - last_ckpt_n >= checkpoint_every_n) or (
                            now - last_ckpt_ts >= checkpoint_every_sec
                        ):
                            flush_buffer(f"ckpt_R{round_id}")

                        if not reached:
                            try:
                                task = next(task_iter)
                                in_flight.append(executor.submit(_worker_simulate_parallel, task))
                            except StopIteration:
                                pass

            flush_buffer("final")
            pbar.close()

            with open(filepath.with_suffix(".fail_stats.json"), "w", encoding="utf-8") as f:
                json.dump(
                    {
                        "fail_reasons": fail_reasons,
                        "fail_examples": fail_examples,
                        "expand_stats": expand_stats,
                        "total_requested": total_requested,
                        "n_success_target": int(n_samples),
                        "n_success_written": int(app.n_samples),
                        "generated_at": datetime.now().isoformat(),
                    },
                    f,
                    indent=2,
                    ensure_ascii=False,
                )

            self._cleanup_workers()
            return str(filepath)

        finally:
            try:
                if buffer_samples:
                    app.append_samples(buffer_samples)
                app.close()
            except Exception:
                pass