import random
from collections import deque
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class QNetwork(nn.Module):
    def __init__(self, state_dim, n_actions, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, n_actions),
        )

    def forward(self, x):
        return self.net(x)


class ReplayBuffer:
    def __init__(self, capacity=10000):
        self.buffer = deque(maxlen=capacity)

    def push(self, s, a, r, s_next, done):
        self.buffer.append((s, a, r, s_next, done))

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        s, a, r, s_next, done = zip(*batch)
        return (
            torch.tensor(np.array(s), dtype=torch.float32),
            torch.tensor(a, dtype=torch.int64),
            torch.tensor(r, dtype=torch.float32),
            torch.tensor(np.array(s_next), dtype=torch.float32),
            torch.tensor(done, dtype=torch.float32),
        )

    def __len__(self):
        return len(self.buffer)


def train_dqn(
    env, state_dim, n_actions, episodes=300, gamma=0.99,
    lr=1e-3, buffer_size=10000, batch_size=64,
    target_update_every=200, double_dqn=True, seed=0,
):
    torch.manual_seed(seed)
    online = QNetwork(state_dim, n_actions)
    target = QNetwork(state_dim, n_actions)
    target.load_state_dict(online.state_dict())
    optimizer = torch.optim.Adam(online.parameters(), lr=lr)
    buffer = ReplayBuffer(buffer_size)

    epsilon, eps_min, eps_decay = 1.0, 0.05, 0.995
    step_count = 0
    episode_returns = []

    for ep in range(episodes):
        state = env.reset()
        total_reward = 0.0
        done = False
        while not done:
            if random.random() < epsilon:
                action = random.randrange(n_actions)
            else:
                with torch.no_grad():
                    q = online(torch.tensor(state, dtype=torch.float32))
                    action = int(q.argmax())

            next_state, reward, done = env.step(action)
            buffer.push(state, action, reward, next_state, done)
            state = next_state
            total_reward += reward
            step_count += 1

            if len(buffer) >= batch_size:
                s, a, r, s_next, d = buffer.sample(batch_size)
                q_values = online(s).gather(1, a.unsqueeze(1)).squeeze(1)

                with torch.no_grad():
                    if double_dqn:
                        next_actions = online(s_next).argmax(dim=1)
                        next_q = target(s_next).gather(
                            1, next_actions.unsqueeze(1)
                        ).squeeze(1)
                    else:
                        next_q = target(s_next).max(dim=1).values
                    td_target = r + gamma * next_q * (1 - d)

                loss = F.mse_loss(q_values, td_target)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            if step_count % target_update_every == 0:
                target.load_state_dict(online.state_dict())

        epsilon = max(eps_min, epsilon * eps_decay)
        episode_returns.append(total_reward)

    return online, episode_returns
