# pendulum_transformer.py
import argparse, math, os
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt

# ============ 工具函数 ============
def moving_average(a, w=5):
    if w <= 1: return a
    pad = np.pad(a, (w//2, w-1-w//2), mode='edge')
    k = np.ones(w)/w
    return np.convolve(pad, k, mode='valid')

def fit_circle_lstsq(x, y):
    # 解: x^2 + y^2 = 2 a x + 2 b y + c
    A = np.c_[2*x, 2*y, np.ones_like(x)]
    b = x**2 + y**2
    sol, *_ = np.linalg.lstsq(A, b, rcond=None)
    a, b_, c = sol
    cx, cy = a, b_
    r = math.sqrt(max(cx*cx + cy*cy + c, 0.0))
    return cx, cy, r

def wrap_to_pi(a):
    return (a + np.pi) % (2*np.pi) - np.pi

def build_windows(arr, s_in=120, s_out=30, stride=10):
    # arr: [T, F]
    T, F = arr.shape
    X, Y, t_idx = [], [], []
    for t in range(0, T - s_in - s_out + 1, stride):
        X.append(arr[t:t+s_in])
        Y.append(arr[t+s_in:t+s_in+s_out])
        t_idx.append(t)
    return np.array(X), np.array(Y), np.array(t_idx)  # [N, S, F], [N, S, F], [N]

def sincos_to_angle(s, c):
    return np.arctan2(s, c)

def angle_rmse(pred_sc, true_sc):
    # pred_sc, true_sc: [..., 2] -> (sin, cos)
    ps, pc = pred_sc[...,0], pred_sc[...,1]
    ts, tc = true_sc[...,0], true_sc[...,1]
    pa = np.arctan2(ps, pc)
    ta = np.arctan2(ts, tc)
    diff = wrap_to_pi(pa - ta)
    return np.sqrt(np.mean(diff**2))

def split_even_indices(n, ratios=(0.7, 0.15, 0.15), block=50, seed=123):
    """将 [0..n-1] 均匀地分成三部分（train/val/test）"""
    if n <= 0:
        return np.array([], dtype=int), np.array([], dtype=int), np.array([], dtype=int)
    probs = np.array(ratios, dtype=np.float64)
    s = probs.sum()
    probs = probs / (s if s > 0 else 1.0)
    rng = np.random.default_rng(seed)
    idxs = np.arange(n)
    tr, va, te = [], [], []
    for start in range(0, n, block):
        end = min(start + block, n)
        m = end - start
        target = probs * m
        base = np.floor(target).astype(int)
        rem = m - int(base.sum())
        if rem > 0:
            frac = target - base
            order = np.argsort(-(frac + 1e-9 * rng.random(frac.shape)))
            for k in order[:rem]:
                base[k] += 1
        local = idxs[start:end].copy()
        rng.shuffle(local)
        a, b, c = int(base[0]), int(base[1]), int(base[2])
        tr.extend(local[:a]); va.extend(local[a:a+b]); te.extend(local[a+b:a+b+c])
    return np.array(tr, dtype=int), np.array(va, dtype=int), np.array(te, dtype=int)

# ============ 数据集 ============
class SeqDataset(Dataset):
    def __init__(self, X, Y, mu=None, std=None):
        self.X = X.astype(np.float32)
        self.Y = Y.astype(np.float32)
        if mu is None or std is None:
            mu = self.X.reshape(-1, self.X.shape[-1]).mean(0, keepdims=True)
            std = self.X.reshape(-1, self.X.shape[-1]).std(0, keepdims=True) + 1e-8
        self.mu  = mu.astype(np.float32)
        self.std = std.astype(np.float32)

    def __len__(self): return self.X.shape[0]
    def __getitem__(self, i):
        x = (self.X[i] - self.mu) / self.std
        y = (self.Y[i] - self.mu) / self.std
        return torch.from_numpy(x), torch.from_numpy(y)

# ============ 位置编码 ============
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=4096):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
        div = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float32) * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer('pe', pe.unsqueeze(0))  # [1, L, D]
    def forward(self, x):
        return x + self.pe[:, :x.size(1), :]

# ============ 极简 Transformer ============
class TimeSeriesTransformer(nn.Module):
    def __init__(self, in_dim=2, d_model=256, nhead=4, enc_layers=3, dec_layers=3, d_ff=512, dropout=0.1, out_dim=2):
        super().__init__()
        self.in_proj  = nn.Linear(in_dim, d_model)
        self.out_proj = nn.Linear(d_model, out_dim)
        self.tgt_proj = nn.Linear(out_dim, d_model)
        self.pos_enc  = PositionalEncoding(d_model)
        self.tf = nn.Transformer(
            d_model=d_model, nhead=nhead,
            num_encoder_layers=enc_layers, num_decoder_layers=dec_layers,
            dim_feedforward=d_ff, dropout=dropout, batch_first=True
        )

    def forward(self, src, tgt_in, tgt_mask=None):
        src = self.pos_enc(self.in_proj(src))
        tgt = self.pos_enc(self.tgt_proj(tgt_in))
        out = self.tf(src, tgt, tgt_mask=tgt_mask)  # [B, S_out, D]
        return self.out_proj(out)

# ============ 主流程 ============
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", type=str, nargs='+', required=True, help="CSV files")
    ap.add_argument("--split", type=float, nargs=3, default=[0.7, 0.15, 0.15])
    ap.add_argument("--s_in", type=int, default=180)
    ap.add_argument("--s_out", type=int, default=30)
    ap.add_argument("--stride", type=int, default=10)
    ap.add_argument("--epochs", type=int, default=10)
    ap.add_argument("--batch", type=int, default=64)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--ckpt", type=str, default="ckpt/pendulum_tf_base.pt")
    ap.add_argument("--eval_only", action="store_true")
    # 新增：评估相关
    ap.add_argument("--eval_all", action="store_true", help="Evaluate on the entire test set")
    ap.add_argument("--eval_samples", type=int, default=6, help="How many random samples to plot from test set")
    ap.add_argument("--seed", type=int, default=123, help="Random seed for eval sampling")
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 规范化路径
    def _norm(p: str) -> str:
        return os.path.abspath(os.path.expanduser(p))

    csv_list = sorted(set(_norm(p) for p in (args.csv or [])))
    if not csv_list:
        raise ValueError("No CSV files provided.")

    ratios = np.array(args.split, dtype=np.float64)
    s = ratios.sum()
    ratios = ratios / (s if s > 0 else 1.0)
    print(f"[Split] auto per-CSV even mode: files={len(csv_list)} | ratios train/val/test = {ratios.tolist()}")

    Xtr_list, Ytr_list, Xva_list, Yva_list, Xte_list, Yte_list = [], [], [], [], [], []

    for i, csv_path in enumerate(csv_list):
        df = pd.read_csv(csv_path)
        x = df['x'].to_numpy(dtype=np.float64)
        y = df['y'].to_numpy(dtype=np.float64)

        # 圆拟合 => 悬点
        cx, cy, r = fit_circle_lstsq(x, y)
        print(f"[{i+1}/{len(csv_list)}] {os.path.basename(csv_path)} pivot=({cx:.2f},{cy:.2f}), r~{r:.2f}")

        # 像素 -> 角度；竖直向下为 0
        theta = np.arctan2(x - cx, cy - y)
        theta = np.unwrap(theta)
        sin_t, cos_t = np.sin(theta), np.cos(theta)
        feats = np.stack([sin_t, cos_t], axis=1)  # [T, 2]

        # 滑窗
        Xi, Yi, _ = build_windows(feats, s_in=args.s_in, s_out=args.s_out, stride=args.stride)
        if Xi.size == 0:
            continue
        tr_idx, va_idx, te_idx = split_even_indices(Xi.shape[0], ratios)
        if tr_idx.size: Xtr_list.append(Xi[tr_idx]); Ytr_list.append(Yi[tr_idx])
        if va_idx.size: Xva_list.append(Xi[va_idx]); Yva_list.append(Yi[va_idx])
        if te_idx.size: Xte_list.append(Xi[te_idx]); Yte_list.append(Yi[te_idx])

    def _concat(lst, steps):
        if not lst:
            return np.zeros((0, steps, 2), dtype=np.float32)
        return np.concatenate(lst, axis=0).astype(np.float32)

    Xtr = _concat(Xtr_list, args.s_in)
    Ytr = _concat(Ytr_list, args.s_out)
    Xva = _concat(Xva_list, args.s_in)
    Yva = _concat(Yva_list, args.s_out)
    Xte = _concat(Xte_list, args.s_in)
    Yte = _concat(Yte_list, args.s_out)

    if Xtr.shape[0] == 0: raise ValueError("No training samples produced.")
    if Xva.shape[0] == 0: raise ValueError("No validation samples produced.")
    if Xte.shape[0] == 0: raise ValueError("No test samples produced.")

    # 标准化参数（仅基于训练集）
    base = Xtr
    mu = base.reshape(-1, base.shape[-1]).mean(0, keepdims=True).astype(np.float32)
    std = base.reshape(-1, base.shape[-1]).std(0, keepdims=True).astype(np.float32) + 1e-8

    dtr = SeqDataset(Xtr, Ytr, mu, std)
    dva = SeqDataset(Xva, Yva, mu, std)
    dte = SeqDataset(Xte, Yte, mu, std)

    nw = 2
    Ltr = DataLoader(dtr, batch_size=args.batch, shuffle=True, drop_last=True, pin_memory=True, num_workers=nw)
    Lva = DataLoader(dva, batch_size=args.batch, shuffle=False, pin_memory=True, num_workers=nw)
    Lte = DataLoader(dte, batch_size=args.batch, shuffle=False, pin_memory=True, num_workers=nw)

    cfg = dict(d_model=256, nhead=4, enc_layers=3, dec_layers=3, d_ff=512)
    net = TimeSeriesTransformer(in_dim=2, out_dim=2, **cfg).to(device)

    mse = nn.MSELoss()
    opt = torch.optim.AdamW(net.parameters(), lr=args.lr, weight_decay=1e-4)

    # ---- warmup + cosine 调度（按 step）----
    total_steps = max(1, args.epochs * len(Ltr))
    warmup = min(500, max(1, total_steps // 10))
    def lr_lambda(step):
        if step < warmup:
            return float(step + 1) / float(warmup)
        progress = (step - warmup) / max(1, total_steps - warmup)
        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
        min_factor = 0.01
        return min_factor + (1 - min_factor) * cosine
    scheduler = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda=lr_lambda)

    # === 标准化 <-> 原始空间 转换 & 单位化 ===
    mu_t = torch.from_numpy(mu).to(device)    # [1, 2]
    std_t = torch.from_numpy(std).to(device)  # [1, 2]

    def to_raw(z):  return z * std_t + mu_t
    def to_std(z):  return (z - mu_t) / std_t
    def unitize(z, eps=1e-6):
        n = torch.clamp(torch.linalg.norm(z, dim=-1, keepdim=True), min=eps)
        return z / n

    # === 预测一个 batch 的自回归函数（返回 raw 空间预测） ===
    def rollout_raw(xb_std, steps):
        # xb_std: [B, S_in, 2] 标准化空间
        tgt_tokens = xb_std[:, -1:, :2].clone()  # 起始 token（标准化空间）
        for _ in range(steps):
            mask = nn.Transformer.generate_square_subsequent_mask(tgt_tokens.size(1)).to(device)
            out = net(xb_std, tgt_tokens, tgt_mask=mask)    # 标准化
            next_std = out[:, -1:, :]
            next_raw = to_raw(next_std)                      # -> 原始
            next_raw = unitize(next_raw)                    # 单位圆约束
            next_std = to_std(next_raw)                     # -> 标准化（与训练分布一致）
            tgt_tokens = torch.cat([tgt_tokens, next_std], dim=1)
        pred_std = tgt_tokens[:, 1:, :]                     # [B, steps, 2] 标准化
        return to_raw(pred_std)                             # -> 原始

    # === 改动一：闭环 rollout 损失所需的角度/方向损失 ===
    def cosine_ang_loss(pred_sc_raw, true_sc_raw, eps=1e-6):
        """
        pred_sc_raw/true_sc_raw: [B,S,2]，分别是 RAW 空间下的 (sin,cos)
        计算 1 - cos(Δθ) 的平均值，忽略向量长度，只关注方向（角度）差。
        """
        pu = pred_sc_raw / torch.clamp(torch.linalg.norm(pred_sc_raw, dim=-1, keepdim=True), min=eps)
        tu = true_sc_raw / torch.clamp(torch.linalg.norm(true_sc_raw, dim=-1, keepdim=True), min=eps)
        return (1.0 - (pu * tu).sum(dim=-1)).mean()

    # === 训练 or 加载 ===
    os.makedirs(os.path.dirname(args.ckpt) or ".", exist_ok=True)
    if args.eval_only:
        if not os.path.exists(args.ckpt):
            raise FileNotFoundError(f"Checkpoint not found: {args.ckpt}. 请先训练一次或指定正确路径。")
        # 关键修复：显式关闭 weights_only 的安全反序列化限制（你信任本地 ckpt）
        payload = torch.load(args.ckpt, map_location=device, weights_only=False)
        # 若 ckpt 中包含 cfg，用 ckpt 的配置重建网络，避免结构不一致
        ckpt_cfg = payload.get("cfg", None)
        if ckpt_cfg is not None:
            # 如果与当前构造的不同，则重建
            if ckpt_cfg != cfg:
                net = TimeSeriesTransformer(in_dim=2, out_dim=2, **ckpt_cfg).to(device)
        net.load_state_dict(payload["model"], strict=True)

        mu = payload["mu"]; std = payload["std"]
        mu_t = torch.from_numpy(mu).to(device); std_t = torch.from_numpy(std).to(device)
        print(f"Loaded checkpoint from {args.ckpt}")
    else:
        # Sanity：原始空间下 ||yb|| 应 ~1
        with torch.no_grad():
            xb0, yb0 = next(iter(Ltr))
            m_norm_raw = torch.linalg.norm(to_raw(yb0.to(device)), dim=-1).mean().item()
            print(f"[Sanity] mean ||yb|| in RAW space ~= {m_norm_raw:.4f}")

        best_val = float("inf")
        global_step = 0
        for ep in range(1, args.epochs+1):
            net.train()
            total = 0.0
            total_tf = 0.0
            total_roll = 0.0

            # ρ：闭环损失权重（线性升到 1.0；如需更稳可改为升到 0.5）
            rho = min(1.0, ep / max(1, args.epochs // 2))

            for xb, yb in Ltr:
                xb = xb.to(device, non_blocking=True); yb = yb.to(device, non_blocking=True)
                B, S, F = yb.shape

                # ---------- (A) Teacher Forcing 分支（与原来一致） ----------
                x_last = xb[:, -1:, :F]
                tgt_in = torch.cat([x_last, yb[:, :-1, :]], dim=1)
                mask = nn.Transformer.generate_square_subsequent_mask(S).to(device)

                pred_tf = net(xb, tgt_in, tgt_mask=mask)     # 标准化空间
                pred_tf_u = unitize(to_raw(pred_tf))         # 原始空间单位圆（只对方向约束）
                yb_u      = unitize(to_raw(yb))
                loss_tf   = mse(pred_tf_u, yb_u)

                # ---------- (B) 闭环 rollout 分支（改动一的核心） ----------
                with torch.no_grad():
                    y_true_raw = to_raw(yb)                 # RAW，用于对齐真值方向
                y_pred_raw = rollout_raw(xb, steps=S)       # RAW，模型自回归滚动预测
                loss_roll  = cosine_ang_loss(y_pred_raw, y_true_raw)

                # ---------- 合成总损失 ----------
                loss = (1.0 - rho) * loss_tf + rho * loss_roll

                opt.zero_grad(set_to_none=True)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
                opt.step()
                scheduler.step()

                total += loss.item() * xb.size(0)
                total_tf += loss_tf.item() * xb.size(0)
                total_roll += loss_roll.item() * xb.size(0)
                global_step += 1

            tr_loss = total / len(dtr)
            tr_tf   = total_tf / len(dtr)
            tr_rl   = total_roll / len(dtr)

            # 验证：保持原先的 TF-式度量（也可扩展加闭环验证，这里先不改动）
            net.eval()
            with torch.no_grad():
                total_v = 0.0
                for xb, yb in Lva:
                    xb = xb.to(device, non_blocking=True); yb = yb.to(device, non_blocking=True)
                    x_last = xb[:, -1:, :2]
                    tgt_in = torch.cat([x_last, yb[:, :-1, :]], dim=1)
                    mask = nn.Transformer.generate_square_subsequent_mask(yb.size(1)).to(device)
                    pred = net(xb, tgt_in, tgt_mask=mask)
                    pred_u = unitize(to_raw(pred))
                    yb_u   = unitize(to_raw(yb))
                    total_v += mse(pred_u, yb_u).item() * xb.size(0)
                va_loss = total_v / len(dva)

            cur_lr = scheduler.get_last_lr()[0]
            print(f"[Epoch {ep:02d}] ρ={rho:.2f} | train TOTAL={tr_loss:.6f} (TF={tr_tf:.6f}, ROLL={tr_rl:.6f}) | val MSE(u_raw)={va_loss:.6f} | lr={cur_lr:.6e}")

            if not np.isnan(va_loss) and va_loss < best_val:
                best_val = va_loss
                payload = {"model": net.state_dict(), "mu": mu, "std": std, "cfg": cfg}
                torch.save(payload, args.ckpt)
                print(f"Saved checkpoint -> {args.ckpt}")

    # ===== 评估 =====
    net.eval()
    rng = np.random.default_rng(args.seed)

    with torch.no_grad():
        # --- 1) 全量评估（可选） ---
        if args.eval_all:
            steps = args.s_out
            sum_sq = 0.0
            n_total = 0
            per_step_sq = np.zeros(steps, dtype=np.float64)
            total_samples = 0

            for xb, yb in Lte:
                xb = xb.to(device, non_blocking=True); yb = yb.to(device, non_blocking=True)
                y_pred_raw = rollout_raw(xb, steps)                  # [B, steps, 2] raw
                y_true_raw = to_raw(yb)                              # [B, steps, 2] raw
                y_pred_np = y_pred_raw.cpu().numpy()
                y_true_np = y_true_raw.cpu().numpy()

                # 单位圆 -> 角度误差
                n1 = np.clip(np.linalg.norm(y_pred_np, axis=-1, keepdims=True), 1e-6, None)
                n2 = np.clip(np.linalg.norm(y_true_np, axis=-1, keepdims=True), 1e-6, None)
                y_pred_u = y_pred_np / n1
                y_true_u = y_true_np / n2

                pa = np.arctan2(y_pred_u[...,0], y_pred_u[...,1])   # 注意顺序 (sin,cos)
                ta = np.arctan2(y_true_u[...,0], y_true_u[...,1])
                diff = wrap_to_pi(pa - ta)                           # [B, steps]

                sum_sq += np.sum(diff**2)
                per_step_sq += np.sum(diff**2, axis=0)
                n_total += diff.size
                total_samples += diff.shape[0]

            overall_rmse = math.sqrt(sum_sq / n_total)
            per_step_rmse = np.sqrt(per_step_sq / max(1, total_samples))

            print(f"[Test:ALL] overall angle RMSE = {overall_rmse:.6f} rad ({overall_rmse*180/np.pi:.3f} deg)")
            # 保存每步 RMSE 曲线
            ts = np.arange(1, steps+1)
            plt.figure(figsize=(6,4))
            plt.plot(ts, per_step_rmse * 180/np.pi, marker='o')
            plt.xlabel("horizon (step)"); plt.ylabel("RMSE (deg)")
            plt.title("Per-step angle RMSE on test set")
            plt.grid(True, linestyle="--", alpha=0.5)
            plt.tight_layout()
            plt.savefig("rmse_per_step.png", dpi=150)
            print("Saved plot -> rmse_per_step.png")

        # --- 2) 随机抽样 K 个样本可视化 ---
        X_list, Y_list = [], []
        for xb, yb in Lte:
            X_list.append(xb); Y_list.append(yb)
        X_all = torch.cat(X_list, dim=0).to(device, non_blocking=True)
        Y_all = torch.cat(Y_list, dim=0).to(device, non_blocking=True)

        N = X_all.size(0)
        K = min(max(0, args.eval_samples), N)
        if K > 0:
            idxs = rng.choice(N, size=K, replace=False)
            steps = Y_all.size(1)
            y_pred_raw_all = rollout_raw(X_all[idxs], steps)     # [K, steps, 2]
            y_true_raw_all = to_raw(Y_all[idxs])                 # [K, steps, 2]

            def unit_np(z):
                n = np.linalg.norm(z, axis=-1, keepdims=True)
                n = np.clip(n, 1e-6, None)
                return z / n

            y_pred_np = y_pred_raw_all.cpu().numpy()
            y_true_np = y_true_raw_all.cpu().numpy()
            y_pred_u = unit_np(y_pred_np)
            y_true_u = unit_np(y_true_np)

            for j in range(K):
                ya = np.arctan2(y_true_u[j,:,0], y_true_u[j,:,1])
                pa = np.arctan2(y_pred_u[j,:,0], y_pred_u[j,:,1])
                pa_aligned = pa + 2*np.pi * np.round((ya - pa)/(2*np.pi))

                plt.figure(figsize=(8,4))
                plt.plot(np.arange(steps), ya, label="true θ")
                plt.plot(np.arange(steps), pa_aligned, '--', label="pred θ (aligned)")
                plt.xlabel("future step"); plt.ylabel("angle (rad)")
                plt.title(f"{steps}-step forecast (test sample {j})")
                plt.legend(); plt.tight_layout()
                out_png = f"pred_plot_sample_{j}.png"
                plt.savefig(out_png, dpi=150)
                plt.close()
                print(f"Saved plot -> {out_png}")

        # 如果没开 --eval_all，也输出一个“单 batch”对齐的整体 RMSE 以免落差太大
        if not args.eval_all:
            steps = args.s_out
            xb, yb = X_all[:args.batch], Y_all[:args.batch]
            y_pred_raw = rollout_raw(xb, steps)
            y_true_raw = to_raw(yb)
            y_pred_np = y_pred_raw.cpu().numpy()
            y_true_np = y_true_raw.cpu().numpy()
            n1 = np.clip(np.linalg.norm(y_pred_np, axis=-1, keepdims=True), 1e-6, None)
            n2 = np.clip(np.linalg.norm(y_true_np, axis=-1, keepdims=True), 1e-6, None)
            y_pred_u = y_pred_np / n1; y_true_u = y_true_np / n2
            pa = np.arctan2(y_pred_u[...,0], y_pred_u[...,1])
            ta = np.arctan2(y_true_u[...,0], y_true_u[...,1])
            diff = wrap_to_pi(pa - ta)
            rmse = math.sqrt(np.mean(diff**2))
            print(f"[Test:Mini] angle RMSE = {rmse:.6f} rad ({rmse*180/np.pi:.3f} deg)")

if __name__ == "__main__":
    main()