"""Optional navigation-only behavior-cloning initialization, followed by PPO.

The teacher is used ONLY to collect training data. Export contains the trained MLP,
not a hidden scripted navigation fallback. Later stages never use this initializer.
"""

import math

import mujoco
import numpy as np
import torch

from .env import MobileManipulatorEnv
from .kernel import TASK, navigation_error


def navigation_teacher(state):
    _, yaw = navigation_error(state)
    x = state[37] - TASK["navigationOffset"][0] - state[0]
    y = state[38] - TASK["navigationOffset"][1] - state[1]
    action = np.zeros(TASK["actionSize"], dtype=np.float32)
    action[0] = 8 * (math.cos(yaw) * x + math.sin(yaw) * y)
    action[1] = 8 * (-math.sin(yaw) * x + math.cos(yaw) * y)
    action[2] = -4 * yaw
    return np.clip(action, -1, 1)


def bootstrap_navigation(agent, package, params, reset_options, seed):
    steps = params["navigationBootstrapSteps"]
    if not steps:
        return
    print(f"Navigation BC initialization: {steps} simulated teacher steps, then PPO", flush=True)
    env = MobileManipulatorEnv(
        package,
        stage="navigate",
        reset_options=reset_options,
        position_jitter=params["positionJitter"],
    )
    rng = np.random.default_rng(seed)
    observations, actions = [], []

    def reset():
        env.reset()
        # Cover heading errors rather than cloning only a straight, zero-yaw path.
        yaw = rng.uniform(-0.6, 0.6)
        env.data.qpos[env.base_q + 3 : env.base_q + 7] = [
            math.cos(yaw / 2),
            0,
            0,
            math.sin(yaw / 2),
        ]
        mujoco.mj_forward(env.model, env.data)
        env.hold()
        return env.observe().copy()

    try:
        env.reset(seed=seed)
        observation = reset()
        for _ in range(steps):
            action = navigation_teacher(env.state)
            observations.append(observation.copy())
            actions.append(action)
            # Cover nearby off-teacher states without teleportation or privileged deployment inputs.
            noisy = action.copy()
            noisy[:3] += rng.normal(0, 0.2, 3)
            observation, _, terminated, truncated, _ = env.step(noisy)
            if terminated or truncated:
                observation = reset()
    finally:
        env.close()
    x = torch.as_tensor(np.asarray(observations), device=agent.device)
    y = torch.as_tensor(np.asarray(actions), device=agent.device)
    policy = agent.policy
    optimizer = torch.optim.Adam(policy.parameters(), lr=1e-3)
    for _ in range(30):
        for batch in torch.randperm(steps, device=agent.device).split(256):
            features = policy.extract_features(x[batch])
            predicted = policy.action_net(policy.mlp_extractor.forward_actor(features))
            loss = torch.nn.functional.mse_loss(predicted, y[batch])
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(policy.parameters(), 1)
            optimizer.step()
    print(f"Navigation BC initialization finished; final batch MSE={loss.item():.6g}", flush=True)