"""
ml/pretrain_bc.py - Behavioral Cloning from Kinematics Teacher
"""
import sys
import time
import argparse
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from stable_baselines3 import PPO
from tqdm import tqdm  # <--- Progress Bar Support

# Ensure project root is in sys.path
sys.path.append(str(Path(__file__).resolve().parent.parent))
from ml.env import JackBotEnv, CurriculumPhase


def collect_kinematics_dataset(num_samples: int = 100_000, use_gui: bool = False):
    """Collects (Observation, Action) pairs directly from Kinematics Teacher."""
    print(f"\n[Pretrain] Collecting {num_samples} samples from Kinematics Teacher (GUI={use_gui})...")
    
    env = JackBotEnv(use_gui=use_gui, robot_mode="kinematics")
    env.curriculum_phase = CurriculumPhase.FULL_COMMAND

    observations = []
    actions = []

    obs, _ = env.reset()
    
    # --- PROGRESS BAR: Data Collection ---
    pbar = tqdm(range(num_samples), desc=" Collecting Data", unit="step")
    for i in pbar:
        # 1. Sample random movement command
        if i % 120 == 0:
            env.command = env.sample_command()

        # 2. Let Robot compute Kinematics target angles
        cmd_vx, cmd_vy, cmd_omega = env.command
        env.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
        env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
        
        # Advance internal IK tick to calculate joint angles
        env.robot.tick()
        
        # Extract .data and flatten (6, 3) matrix to 18-dim 1D array
        target_ik_rad = env.robot.current_rad.data.flatten().copy()

        # 3. Convert target angles back to normalized [-1, 1] action space
        normalized_action = np.where(
            target_ik_rad >= env.default_joint_angles,
            (target_ik_rad - env.default_joint_angles) / np.maximum(1e-5, env.max_joint_limits - env.default_joint_angles),
            (target_ik_rad - env.default_joint_angles) / np.maximum(1e-5, env.default_joint_angles - env.min_joint_limits)
        )
        normalized_action = np.clip(normalized_action, -1.0, 1.0)

        # 4. Save sample
        observations.append(obs.copy())
        actions.append(normalized_action.copy())

        # Step simulation environment
        obs, _, terminated, truncated, _ = env.step(normalized_action)

        if use_gui:
            time.sleep(1.0 / 60.0)

        if terminated or truncated:
            obs, _ = env.reset()

    env.close()
    print("[Pretrain] Data collection complete!\n")
    return np.array(observations, dtype=np.float32), np.array(actions, dtype=np.float32)


def pretrain_policy(
    save_path: str = "ml/checkpoints/jackbot_kinematics_base.zip",
    epochs: int = 15,
    batch_size: int = 256,
    num_samples: int = 100_000,
    use_gui: bool = False
):
    # Collect dataset from Kinematics teacher
    obs_data, action_data = collect_kinematics_dataset(num_samples=num_samples, use_gui=use_gui)

    # Initialize Dummy Env & Fresh SB3 PPO Model
    dummy_env = JackBotEnv(use_gui=False, robot_mode="direct")
    model = PPO("MlpPolicy", dummy_env, learning_rate=5e-4, verbose=0, device="cpu")

    # Extract PyTorch Policy Network & Optimizer
    policy = model.policy
    optimizer = torch.optim.Adam(policy.parameters(), lr=5e-4)
    loss_fn = nn.MSELoss()

    # Convert to PyTorch Dataloader
    dataset = TensorDataset(torch.tensor(obs_data), torch.tensor(action_data))
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    print(f"[Pretrain] Pre-training Policy Network ({epochs} Epochs)...")
    policy.train()

    # --- PROGRESS BAR: Epoch Training ---
    for epoch in range(epochs):
        epoch_loss = 0.0
        
        batch_pbar = tqdm(loader, desc=f" Epoch {epoch + 1:02d}/{epochs:02d}", leave=True, unit="batch")
        for batch_obs, batch_actions in batch_pbar:
            optimizer.zero_grad()
            
            distribution = policy.get_distribution(batch_obs)
            predicted_actions = distribution.distribution.mean

            loss = loss_fn(predicted_actions, batch_actions)
            loss.backward()
            optimizer.step()

            current_loss = loss.item()
            epoch_loss += current_loss * len(batch_obs)
            
            # Dynamic loss update in progress bar tail
            batch_pbar.set_postfix({"loss": f"{current_loss:.6f}"})

        avg_loss = epoch_loss / len(dataset)
        tqdm.write(f"  └─ Epoch {epoch + 1:02d}/{epochs:02d} Complete | Mean MSE Loss: {avg_loss:.6f}")

    # Save SB3 Model Checkpoint
    out_file = Path(save_path)
    out_file.parent.mkdir(parents=True, exist_ok=True)
    model.save(out_file)
    dummy_env.close()
    print(f"\n[Pretrain] Successfully saved pre-trained base model to: {out_file.resolve()}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="JackBot Behavioral Cloning Pre-trainer")
    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering during collection")
    parser.add_argument("--num-samples", type=int, default=100_000, help="Number of dataset samples to collect")
    parser.add_argument("--epochs", type=int, default=15, help="Number of BC training epochs")
    parser.add_argument("--save-path", type=str, default="ml/checkpoints/jackbot_kinematics_base.zip", help="Output path for pre-trained model .zip")
    args = parser.parse_args()

    pretrain_policy(
        save_path=args.save_path,
        epochs=args.epochs,
        num_samples=args.num_samples,
        use_gui=args.gui
    )