#!/usr/bin/env python3
"""离线强化学习：在训练期窗口、训练期种子上采日志，然后 Fitted Q Iteration。

离线的含义（严格）：学习过程中一次都不与环境交互。日志由行为策略一次性采完后冻结，
FQI 只在这个固定数据集上做回归迭代。训练日志一条都不来自留出期（2016-01 之后）。
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

import numpy as np
from sklearn.ensemble import ExtraTreesRegressor

sys.path.insert(0, str(Path(__file__).resolve().parent))
import env_depolarization as E

PKG = Path(__file__).resolve().parents[2]
TRAIN_SEEDS = [1, 2, 3, 4, 5, 6, 7, 8]
ROLLOUTS_PER_SEED = 3
GAMMA = 0.9
FQI_ITERS = 5
TREES = dict(n_estimators=300, min_samples_leaf=20, n_jobs=3, random_state=0)


def collect(graph, log):
    """行为策略：每窗以 0.5 概率出手；出手时 50% 随机挑、50% 极端度贪心。"""
    X, A, R, XN, DONE = [], [], [], [], []
    n_train = graph.n_train
    budget_total = E.BUDGET_PER_WINDOW * n_train // 2
    for seed in TRAIN_SEEDS:
        for rep in range(ROLLOUTS_PER_SEED):
            x, s, noise = E.initial_state(graph, seed)
            rng = np.random.default_rng(10_000 + seed * 100 + rep)
            budget_left = budget_total
            phi_prev = a_prev = r_prev = None
            for t in range(n_train):
                cache = graph.window_stats(t - 1)
                phi = E.features(graph, t, x, s, budget_left, budget_total, n_train, cache)
                act = np.zeros(graph.N, dtype=np.int8)
                if budget_left >= E.BUDGET_PER_WINDOW and rng.random() < 0.5:
                    if rng.random() < 0.5:
                        chosen = rng.choice(graph.N, size=E.BUDGET_PER_WINDOW, replace=False)
                    else:
                        chosen = np.argsort(-np.abs(x))[:E.BUDGET_PER_WINDOW]
                    act[chosen] = 1
                    budget_left -= E.BUDGET_PER_WINDOW
                x_before = x.copy()
                x = E.apply_perturbation(x, np.flatnonzero(act))
                x = E.step_window(graph, t, x, s, noise[t])
                reward = (x_before ** 2 - x ** 2) / graph.N
                if phi_prev is not None:
                    X.append(phi_prev); A.append(a_prev); R.append(r_prev)
                    XN.append(phi); DONE.append(0)
                phi_prev, a_prev, r_prev = phi, act, reward
            X.append(phi_prev); A.append(a_prev); R.append(r_prev)
            XN.append(phi_prev); DONE.append(1)
            log.write(f"collect seed={seed} rep={rep} spent={budget_total-budget_left}/{budget_total}\n")
    return (np.vstack(X), np.concatenate(A), np.concatenate(R),
            np.vstack(XN), np.repeat(DONE, graph.N))


def main() -> int:
    z = np.load(PKG / "work/data/temporal_graph.npz", allow_pickle=True)
    graph = E.Graph(z)
    logp = PKG / "work/logs/offline_rl.log"
    with logp.open("w", encoding="utf-8") as log:
        t0 = time.time()
        log.write(f"start {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}\n")
        X, A, R, XN, DONE = collect(graph, log)
        log.write(f"logged transitions: {len(X)} rows, {X.shape[1]} features, "
                  f"action=1 fraction {A.mean():.4f}, collected in {time.time()-t0:.1f}s\n")
        log.flush()

        SA = np.column_stack([X, A])
        SA0 = np.column_stack([XN, np.zeros(len(XN))])
        SA1 = np.column_stack([XN, np.ones(len(XN))])
        model, history = None, []
        for k in range(FQI_ITERS):
            tk = time.time()
            if model is None:
                target = R.copy()
            else:
                q0, q1 = model.predict(SA0), model.predict(SA1)
                target = R + GAMMA * (1 - DONE) * np.maximum(q0, q1)
            model = ExtraTreesRegressor(**TREES).fit(SA, target)
            fit = model.predict(SA)
            rmse = float(np.sqrt(np.mean((fit - target) ** 2)))
            adv = model.predict(np.column_stack([X, np.ones(len(X))])) - \
                  model.predict(np.column_stack([X, np.zeros(len(X))]))
            history.append({"iter": k + 1, "train_rmse": rmse,
                            "mean_advantage": float(adv.mean()),
                            "frac_positive_advantage": float((adv > 0).mean()),
                            "seconds": round(time.time() - tk, 1)})
            log.write(f"FQI iter {k+1}: rmse={rmse:.3e} mean_adv={adv.mean():.3e} "
                      f"pos_adv_frac={(adv>0).mean():.3f} {time.time()-tk:.1f}s\n")
            log.flush()

        import pickle
        (PKG / "work/results/fqi_model.pkl").write_bytes(pickle.dumps(model))
        meta = {
            "algorithm": "Fitted Q Iteration (Ernst et al. 2005) with ExtraTrees function approximation",
            "offline": "学习过程中零环境交互；日志一次采完后冻结",
            "train_seeds": TRAIN_SEEDS, "rollouts_per_seed": ROLLOUTS_PER_SEED,
            "train_windows": [graph.months[0], graph.months[graph.n_train - 1]],
            "n_transitions": int(len(X)), "n_features": int(X.shape[1]),
            "feature_names": E.FEATURE_NAMES,
            "action_fraction_in_log": float(A.mean()),
            "gamma": GAMMA, "fqi_iterations": FQI_ITERS, "tree_params": TREES,
            "per_node_reward": "(x_i(t)^2 - x_i(t+1)^2)/N —— 极化方差均方项的逐节点精确分解（已声明的信用分配近似）",
            "history": history,
            "total_seconds": round(time.time() - t0, 1),
        }
        (PKG / "work/results/fqi_training.json").write_text(
            json.dumps(meta, ensure_ascii=False, indent=1), encoding="utf-8")
        log.write(f"done in {time.time()-t0:.1f}s\n")
    print(json.dumps({k: v for k, v in meta.items() if k != "history"}, ensure_ascii=False, indent=1))
    for h in history:
        print(h)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())