#!/usr/bin/env python3
"""Warm-start TARGET_DIM weights from old 44-dim trained weights."""
import numpy as np
from pathlib import Path

SRC = Path("tools/training_runner/snapshots/best_post_maint_r38003")
DST = Path("PPO_Bot/weights/latest")
OLD_DIM = 44
TARGET_DIM = 57  # change this to expand to a different input dimension

# Pad w1 [64, OLD_DIM] -> [64, TARGET_DIM] with zeros (weights and Adam moments)
for prefix in ("actor", "critic"):
    # Weight
    old = np.load(SRC / f"{prefix}_w1.npy")
    assert old.shape == (64, OLD_DIM), f"unexpected shape {old.shape}"
    new = np.zeros((64, TARGET_DIM), dtype=old.dtype)
    new[:, :OLD_DIM] = old
    np.save(DST / f"{prefix}_w1.npy", new)
    print(f"  {prefix}_w1: {old.shape} -> {new.shape}")

    # Adam moments for w1: pad same way
    adam_base = f"adam_{'a' if prefix == 'actor' else 'c'}w1"
    for moment in ("_m", "_v"):
        old_m = np.load(SRC / f"{adam_base}{moment}.npy")
        new_m = np.zeros((64, TARGET_DIM), dtype=old_m.dtype)
        new_m[:, :OLD_DIM] = old_m
        np.save(DST / f"{adam_base}{moment}.npy", new_m)
        print(f"  {adam_base}{moment}: {old_m.shape} -> {new_m.shape}")

# Copy unchanged weight files as-is
unchanged = [
    "actor_w2", "actor_w3", "actor_b1", "actor_b2", "actor_b3",
    "critic_w2", "critic_w3", "critic_b1", "critic_b2", "critic_b3",
]
for name in unchanged:
    data = np.load(SRC / f"{name}.npy")
    np.save(DST / f"{name}.npy", data)
    print(f"  {name}: {data.shape} copied")

# Initialize log_std to -1.0 (std ≈ 0.37) — snapshot values (2.27–4.68) are too high for fine-tuning
log_std = np.full(6, -1.0, dtype=np.float32)
np.save(DST / "log_std.npy", log_std)
print(f"  log_std: initialized to -1.0 (std≈0.37), shape={log_std.shape}")

# Copy unchanged Adam moments (all except w1, which were handled above)
unchanged_adam = [
    "adam_aw2", "adam_cw2",
    "adam_aw3", "adam_cw3",
    "adam_ab1", "adam_cb1",
    "adam_ab2", "adam_cb2",
    "adam_ab3", "adam_cb3",
    "adam_logstd",
]
for base in unchanged_adam:
    for moment in ("_m", "_v"):
        data = np.load(SRC / f"{base}{moment}.npy")
        np.save(DST / f"{base}{moment}.npy", data)
        print(f"  {base}{moment}: {data.shape} copied")

(DST / "adam_t.txt").write_text("1\n")
(DST.parent / "round_counter.txt").write_text("0\n")
print("  adam_t.txt -> 1, round_counter.txt -> 0")

# Verify
w1 = np.load(DST / "actor_w1.npy")
old_w1 = np.load(SRC / "actor_w1.npy")
assert w1.shape == (64, TARGET_DIM), f"bad shape {w1.shape}"
assert np.allclose(w1[:, :OLD_DIM], old_w1), "old columns don't match"
assert np.all(w1[:, OLD_DIM:] == 0), "new columns not zero"
print(f"\nOK: actor_w1 shape={w1.shape}, cols 0-{OLD_DIM-1} match old, cols {OLD_DIM}-{TARGET_DIM-1} are zero")