import os
import random
import time

import cv2
import gymnasium as gym
import imageio
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from stable_baselines3.common.buffers import ReplayBuffer
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

import wandb


# ===== CONFIGURATION =====
class Config:
    # Experiment settings
    exp_name = "DQN-CartPole"
    seed = 42
    env_id = "CartPole-v1"

    # Training parameters
    total_timesteps = 20000
    learning_rate = 2.5e-4
    buffer_size = 10000
    gamma = 0.99
    tau = 1.0
    target_network_frequency = 50
    batch_size = 128
    start_e = 1.0
    end_e = 0.05
    exploration_fraction = 0.5
    learning_starts = 1000
    train_frequency = 10

    # Logging & saving
    capture_video = True
    save_model = True
    upload_model = True
    hf_entity = ""  # Your Hugging Face username

    # WandB settings
    use_wandb = True
    wandb_project = "cleanRL"
    wandb_entity = ""  # Your WandB username/team


class QNet(nn.Module):
    def __init__(self, state_space, action_space):
        super().__init__()
        self.fc1 = nn.Linear(state_space, 256)
        self.fc2 = nn.Linear(256, 512)
        self.q_value = nn.Linear(512, action_space)

    def forward(self, x):
        return self.q_value(torch.relu(self.fc2(torch.relu(self.fc1(x)))))


class LinearEpsilonDecay(nn.Module):
    def __init__(self, initial_eps, end_eps, total_timesteps):
        super().__init__()
        self.initial_eps = initial_eps
        # self.decay_factor = decay_factor
        self.total_timesteps = total_timesteps
        self.end_eps = end_eps

    def forward(self, current_timestep, decay_factor):
        slope = (self.end_eps - self.initial_eps) / (
            self.total_timesteps * decay_factor
        )
        return max(slope * current_timestep + self.initial_eps, self.end_eps)


def make_env(env_id, seed, capture_video, run_name, eval_mode=False):
    """Create environment with video recording"""
    env = gym.make(env_id, render_mode="rgb_array")
    env = gym.wrappers.RecordEpisodeStatistics(env)

    # Video recording setup
    if capture_video:
        if eval_mode:
            # Evaluation videos
            video_prefix = f"videos/{run_name}/eval"
        else:
            # Training videos
            video_prefix = f"videos/{run_name}/train"
            env = gym.wrappers.RecordVideo(
                env,
                video_prefix,
                episode_trigger=lambda x: x % 100 == 0,  # Record every 100 episodes
            )

    env.action_space.seed(seed)

    return env


def evaluate(model, device, run_name, num_eval_eps=10, record=False):
    eval_env = make_env(
        env_id=Config.env_id,
        seed=Config.seed,
        capture_video=True,
        run_name=run_name,
        eval_mode=True,
    )
    eval_env.action_space.seed(Config.seed)

    model = model.to(device)
    model = model.eval()
    returns = []
    frames = []

    for _ in range(num_eval_eps):
        obs, _ = eval_env.reset()
        done = False
        episode_reward = 0.0
        # episode_frames = []

        while not done:
            if record:
                if episode_reward > 500:
                    print("Hooray! Episode reward exceeded 500, stopping early.")
                    break
                frame = eval_env.render()
                frames.append(frame)  # Capture all frames

            # with torch.no_grad():
            action = (
                model(torch.tensor(obs, device=device).unsqueeze(0)).argmax().item()
            )
            obs, reward, terminated, truncated, _ = eval_env.step(action)
            done = terminated or truncated
            episode_reward += reward

        returns.append(episode_reward)
        # if eps == 0:  # Save frames only for the first episode (optional)
        #     frames = episode_frames.copy()  # Avoid memory issues

    eval_env.close()

    # Save video
    if frames:
        os.makedirs(f"videos/{run_name}/eval", exist_ok=True)
        imageio.mimsave(f"videos/{run_name}/eval/eval_video.mp4", frames, fps=30)

    return returns, frames


def train_dqn(
    env_id="CartPole-v1",
    total_timesteps=20000,
    seed=42,
    learning_rate=2.5e-4,
    buffer_size=10000,
    gamma=0.99,
    tau=1.0,
    target_network_frequency=50,
    batch_size=128,
    start_e=1.0,
    end_e=0.05,
    exploration_fraction=0.5,
    learning_starts=1000,
    train_frequency=10,
    capture_video=False,
    use_wandb=False,
    wandb_project="cleanRL",
    wandb_entity="",
    exp_name="DQN",
):
    """
    Train a DQN agent on a Gymnasium environment.

    Args:
        env_id: Gymnasium environment ID
        total_timesteps: Total training timesteps
        seed: Random seed
        learning_rate: Learning rate for optimizer
        buffer_size: Replay buffer size
        gamma: Discount factor
        tau: Target network update rate
        target_network_frequency: How often to update target network
        batch_size: Batch size for training
        start_e: Initial epsilon for exploration
        end_e: Final epsilon for exploration
        exploration_fraction: Fraction of timesteps for epsilon decay
        learning_starts: When to start learning
        train_frequency: How often to train
        capture_video: Whether to capture training videos
        use_wandb: Whether to use Weights & Biases logging
        wandb_project: W&B project name
        wandb_entity: W&B entity/username
        exp_name: Experiment name

    Returns:
        Trained Q-network model
    """
    run_name = f"{env_id}__{exp_name}__{seed}__{int(time.time())}"

    # Initialize WandB
    if use_wandb:
        wandb.init(
            project=wandb_project,
            entity=wandb_entity,
            sync_tensorboard=True,
            config=locals(),
            name=run_name,
            monitor_gym=True,
            save_code=True,
        )

    if capture_video:
        os.makedirs(f"videos/{run_name}/train", exist_ok=True)
        os.makedirs(f"videos/{run_name}/eval", exist_ok=True)
    os.makedirs(f"runs/{run_name}", exist_ok=True)
    writer = SummaryWriter(f"runs/{run_name}")

    # Set seeds
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    env = make_env(env_id, seed, capture_video, run_name)
    q_network = QNet(env.observation_space.shape[0], env.action_space.n).to(device)
    target_net = QNet(env.observation_space.shape[0], env.action_space.n).to(device)
    target_net.load_state_dict(q_network.state_dict())
    optimizer = optim.Adam(q_network.parameters(), lr=learning_rate)
    eps_decay = LinearEpsilonDecay(start_e, end_e, total_timesteps)

    q_network.train()
    target_net.train()

    replay_buffer = ReplayBuffer(
        buffer_size,
        env.observation_space,
        env.action_space,
        device=device,
        handle_timeout_termination=False,
    )

    obs, _ = env.reset()

    for step in tqdm(range(total_timesteps)):
        eps = eps_decay(step, exploration_fraction)
        rnd = random.random()

        if rnd < eps:
            action = env.action_space.sample()
        else:
            action = (
                q_network(torch.tensor(obs, device=device).unsqueeze(0)).argmax().item()
            )

        new_obs, reward, terminated, truncated, info = env.step(action)
        done = terminated or truncated
        replay_buffer.add(
            obs, new_obs, np.array(action), np.array(reward), np.array(done), [info]
        )

        # Log episode returns
        if "episode" in info:
            print(
                f"Step={step}, Return={info['episode']['r']:.2f}, Length={info['episode']['l']}"
            )

            # WandB logging
            if use_wandb:
                wandb.log(
                    {
                        "episodic_return": info["episode"]["r"],
                        "episodic_length": info["episode"]["l"],
                        "epsilon": eps,
                        "global_step": step,
                    }
                )

        if step > learning_starts and step % train_frequency == 0:
            data = replay_buffer.sample(batch_size)

            target_max = target_net(data.next_observations).max(1)[0]
            td_target = data.rewards.flatten() + gamma * target_max * (
                1 - data.dones.flatten()
            )
            old_val = q_network(data.observations).gather(1, data.actions).squeeze()

            optimizer.zero_grad()
            loss = nn.functional.mse_loss(old_val, td_target)
            loss.backward()
            optimizer.step()

            # Log loss and metrics every 100 steps
            if step % 100 == 0:
                if use_wandb:
                    wandb.log(
                        {
                            "losses/td_loss": loss.item(),
                        }
                    )

        # Update target network
        if step % target_network_frequency == 0:
            for q_params, target_params in zip(
                q_network.parameters(), target_net.parameters()
            ):
                target_params.data.copy_(
                    tau * q_params.data + (1.0 - tau) * target_params.data
                )

        # Model evaluation & saving
        if step % 1000 == 0:
            episodic_returns, _ = evaluate(q_network, device, run_name)
            avg_return = np.mean(episodic_returns)

            if use_wandb:
                wandb.log({"val_avg_return": avg_return, "val_step": step})
            print(f"Evaluation returns: {episodic_returns}, Average: {avg_return:.2f}")

        if done:
            obs, _ = env.reset()
        else:
            obs = new_obs

    # Save final video to WandB
    if use_wandb:
        train_video_path = "videos/final.mp4"
        _, frames = evaluate(q_network, device, run_name, record=True)
        imageio.mimsave(train_video_path, frames, fps=30)
        print(f"Final training video saved to {train_video_path}")
        wandb.finish()

    if capture_video:
        cv2.destroyAllWindows()

    env.close()
    writer.close()

    return q_network
