"""
Implicit Q-Learning (IQL) for offline RL.
Kostrikov et al., 2022 - "Offline Reinforcement Learning with Implicit Q-Learning"

Three networks, three losses:
  1. V-loss: expectile regression of Q-values
  2. Q-loss: MSE with V-targets (no max over actions)
  3. Policy loss: advantage-weighted regression (AWR)
"""

import os
import csv
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal

# Joint limits for clamping policy output
JOINT_LOWER = torch.tensor([-1.606, -1.221, -3.142, -2.251, -3.142, -2.16, -3.142, 0.0, 0.0])
JOINT_UPPER = torch.tensor([1.606, 1.518, 3.142, 2.251, 3.142, 3.142, 3.142, 0.05, 0.05])


# ============================================================
# Networks
# ============================================================

class QNetwork(nn.Module):
    """Q(s, a) -> scalar"""
    def __init__(self, state_dim=9, action_dim=9, hidden_dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, state, action):
        x = torch.cat([state, action], dim=-1)
        return self.net(x).squeeze(-1)


class VNetwork(nn.Module):
    """V(s) -> scalar"""
    def __init__(self, state_dim=9, hidden_dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, state):
        return self.net(state).squeeze(-1)


class GaussianPolicy(nn.Module):
    """Diagonal Gaussian policy: pi(a|s)"""
    def __init__(self, state_dim=9, action_dim=9, hidden_dim=256):
        super().__init__()
        self.trunk = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )
        self.mean_head = nn.Linear(hidden_dim, action_dim)
        self.log_std_head = nn.Linear(hidden_dim, action_dim)

    def forward(self, state):
        h = self.trunk(state)
        mean = self.mean_head(h)
        log_std = self.log_std_head(h).clamp(-5.0, 2.0)
        return mean, log_std

    def log_prob(self, state, action):
        """Log probability of action under the policy (no clamping on action)."""
        mean, log_std = self.forward(state)
        std = log_std.exp()
        dist = Normal(mean, std)
        return dist.log_prob(action).sum(dim=-1)

    def sample(self, state):
        """Sample an action and return (clamped_action, log_prob)."""
        mean, log_std = self.forward(state)
        std = log_std.exp()
        dist = Normal(mean, std)
        raw_action = dist.rsample()
        log_p = dist.log_prob(raw_action).sum(dim=-1)
        clamped = torch.clamp(raw_action, JOINT_LOWER.to(raw_action.device),
                              JOINT_UPPER.to(raw_action.device))
        return clamped, log_p


# ============================================================
# Dataset
# ============================================================

class OfflineDataset:
    """Loads NPZ dataset into GPU tensors for fast sampling."""
    def __init__(self, path, device):
        data = np.load(path)
        self.states = torch.tensor(data['states'], dtype=torch.float32, device=device)
        self.actions = torch.tensor(data['actions'], dtype=torch.float32, device=device)
        self.rewards = torch.tensor(data['rewards'], dtype=torch.float32, device=device)
        self.next_states = torch.tensor(data['next_states'], dtype=torch.float32, device=device)
        self.dones = torch.tensor(data['dones'], dtype=torch.float32, device=device)
        self.state_mean = torch.tensor(data['state_mean'], dtype=torch.float32, device=device)
        self.state_std = torch.tensor(data['state_std'], dtype=torch.float32, device=device)
        self.size = self.states.shape[0]
        print(f"Loaded dataset: {self.size} transitions")

    def normalize_state(self, s):
        return (s - self.state_mean) / self.state_std

    def sample(self, batch_size):
        idx = torch.randint(0, self.size, (batch_size,), device=self.states.device)
        return (
            self.normalize_state(self.states[idx]),
            self.actions[idx],
            self.rewards[idx],
            self.normalize_state(self.next_states[idx]),
            self.dones[idx],
        )


# ============================================================
# Soft update
# ============================================================

@torch.no_grad()
def soft_update(target, source, tau):
    for tp, sp in zip(target.parameters(), source.parameters()):
        tp.data.copy_(tau * sp.data + (1.0 - tau) * tp.data)


# ============================================================
# IQL Trainer
# ============================================================

class IQLTrainer:
    def __init__(self, dataset, device, tau_expectile=0.7, beta=3.0,
                 lr=3e-4, discount=0.99, target_update_rate=0.005):
        self.dataset = dataset
        self.device = device
        self.tau_expectile = tau_expectile
        self.beta = beta
        self.discount = discount
        self.target_update_rate = target_update_rate

        # Networks
        self.q1 = QNetwork().to(device)
        self.q2 = QNetwork().to(device)
        self.q1_target = QNetwork().to(device)
        self.q2_target = QNetwork().to(device)
        self.q1_target.load_state_dict(self.q1.state_dict())
        self.q2_target.load_state_dict(self.q2.state_dict())

        self.v = VNetwork().to(device)
        self.policy = GaussianPolicy().to(device)

        # Move joint limits to device
        global JOINT_LOWER, JOINT_UPPER
        JOINT_LOWER = JOINT_LOWER.to(device)
        JOINT_UPPER = JOINT_UPPER.to(device)

        # Optimizers
        self.opt_q = torch.optim.Adam(
            list(self.q1.parameters()) + list(self.q2.parameters()), lr=lr)
        self.opt_v = torch.optim.Adam(self.v.parameters(), lr=lr)
        self.opt_pi = torch.optim.Adam(self.policy.parameters(), lr=lr)

    def compute_value_loss(self, norm_s, actions):
        with torch.no_grad():
            q1_t = self.q1_target(norm_s, actions)
            q2_t = self.q2_target(norm_s, actions)
            q_target = torch.min(q1_t, q2_t)

        v = self.v(norm_s)
        diff = q_target - v
        weight = torch.where(diff > 0, self.tau_expectile, 1.0 - self.tau_expectile)
        loss_v = (weight * diff ** 2).mean()
        return loss_v, v.mean().item()

    def compute_q_loss(self, norm_s, actions, rewards, norm_s_next, dones):
        with torch.no_grad():
            v_next = self.v(norm_s_next)
            target_q = rewards + self.discount * (1.0 - dones) * v_next

        q1 = self.q1(norm_s, actions)
        q2 = self.q2(norm_s, actions)
        loss_q = ((q1 - target_q) ** 2).mean() + ((q2 - target_q) ** 2).mean()
        return loss_q, q1.mean().item()

    def compute_policy_loss(self, norm_s, actions):
        with torch.no_grad():
            q1_t = self.q1_target(norm_s, actions)
            q2_t = self.q2_target(norm_s, actions)
            q_target = torch.min(q1_t, q2_t)
            v = self.v(norm_s)
            advantage = q_target - v
            exp_adv = torch.exp(self.beta * advantage).clamp(max=100.0)

        log_prob = self.policy.log_prob(norm_s, actions)
        loss_pi = -(exp_adv * log_prob).mean()
        return loss_pi, advantage.mean().item(), exp_adv.mean().item()

    def train_step(self, batch_size=256):
        norm_s, actions, rewards, norm_s_next, dones = self.dataset.sample(batch_size)

        # Update V
        self.opt_v.zero_grad()
        loss_v, v_mean = self.compute_value_loss(norm_s, actions)
        loss_v.backward()
        self.opt_v.step()

        # Update Q
        self.opt_q.zero_grad()
        loss_q, q_mean = self.compute_q_loss(norm_s, actions, rewards, norm_s_next, dones)
        loss_q.backward()
        self.opt_q.step()

        # Update Policy
        self.opt_pi.zero_grad()
        loss_pi, adv_mean, exp_adv_mean = self.compute_policy_loss(norm_s, actions)
        loss_pi.backward()
        self.opt_pi.step()

        # Soft update target Q
        soft_update(self.q1_target, self.q1, self.target_update_rate)
        soft_update(self.q2_target, self.q2, self.target_update_rate)

        return {
            'v_loss': loss_v.item(),
            'q_loss': loss_q.item(),
            'policy_loss': loss_pi.item(),
            'v_mean': v_mean,
            'q_mean': q_mean,
            'advantage_mean': adv_mean,
            'exp_advantage_mean': exp_adv_mean,
        }

    def save(self, path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        torch.save({
            'q1': self.q1.state_dict(),
            'q2': self.q2.state_dict(),
            'q1_target': self.q1_target.state_dict(),
            'q2_target': self.q2_target.state_dict(),
            'v': self.v.state_dict(),
            'policy': self.policy.state_dict(),
        }, path)

    def load(self, path):
        ckpt = torch.load(path, map_location=self.device, weights_only=True)
        self.q1.load_state_dict(ckpt['q1'])
        self.q2.load_state_dict(ckpt['q2'])
        self.q1_target.load_state_dict(ckpt['q1_target'])
        self.q2_target.load_state_dict(ckpt['q2_target'])
        self.v.load_state_dict(ckpt['v'])
        self.policy.load_state_dict(ckpt['policy'])


# ============================================================
# Main training loop
# ============================================================

def train(tau_expectile=0.7, beta=3.0, lr=3e-4, batch_size=256,
          discount=0.99, target_update_rate=0.005, num_iterations=100000,
          seed=42, log_freq=1000, save_freq=10000):

    # Seed
    torch.manual_seed(seed)
    np.random.seed(seed)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")

    # Dataset
    data_path = '/code/zxx240000/training/offline_rl/data/offline_dataset.npz'
    dataset = OfflineDataset(data_path, device)

    # Output dirs
    tau_str = str(tau_expectile).replace('.', '')
    result_dir = f'/code/zxx240000/training/offline_rl/results/iql_tau{tau_str}'
    ckpt_dir = os.path.join(result_dir, 'checkpoints')
    os.makedirs(ckpt_dir, exist_ok=True)

    # CSV log
    csv_path = os.path.join(result_dir, 'training_log.csv')
    csv_fields = ['step', 'v_loss', 'q_loss', 'policy_loss',
                  'v_mean', 'q_mean', 'advantage_mean', 'exp_advantage_mean']
    csv_file = open(csv_path, 'w', newline='')
    csv_writer = csv.DictWriter(csv_file, fieldnames=csv_fields)
    csv_writer.writeheader()

    # Trainer
    trainer = IQLTrainer(dataset, device, tau_expectile=tau_expectile,
                         beta=beta, lr=lr, discount=discount,
                         target_update_rate=target_update_rate)

    print(f"\nStarting IQL training: tau={tau_expectile}, beta={beta}, "
          f"lr={lr}, batch={batch_size}, iters={num_iterations}")
    print(f"Results dir: {result_dir}\n")

    for step in range(1, num_iterations + 1):
        metrics = trainer.train_step(batch_size)

        if step % log_freq == 0:
            row = {'step': step, **metrics}
            csv_writer.writerow(row)
            csv_file.flush()
            print(f"[{step:>6d}] v_loss={metrics['v_loss']:.4f}  "
                  f"q_loss={metrics['q_loss']:.4f}  "
                  f"pi_loss={metrics['policy_loss']:.4f}  "
                  f"V={metrics['v_mean']:.4f}  Q={metrics['q_mean']:.4f}  "
                  f"adv={metrics['advantage_mean']:.4f}  "
                  f"exp_adv={metrics['exp_advantage_mean']:.2f}")

        if step % save_freq == 0:
            ckpt_path = os.path.join(ckpt_dir, f'iql_step{step}.pt')
            trainer.save(ckpt_path)
            print(f"  -> Saved checkpoint: {ckpt_path}")

    csv_file.close()

    # Final save
    final_path = os.path.join(ckpt_dir, 'iql_final.pt')
    trainer.save(final_path)
    print(f"\nTraining complete. Final checkpoint: {final_path}")

    # Validation: sample actions and check limits
    print("\n--- Validation ---")
    norm_s_sample = dataset.normalize_state(dataset.states[:100])
    with torch.no_grad():
        sampled_actions, _ = trainer.policy.sample(norm_s_sample)
        q1_vals = trainer.q1(norm_s_sample, dataset.actions[:100])
        v_vals = trainer.v(norm_s_sample)

    print(f"Sampled action min: {sampled_actions.min(dim=0).values.cpu().numpy()}")
    print(f"Sampled action max: {sampled_actions.max(dim=0).values.cpu().numpy()}")
    print(f"Joint lower limits: {JOINT_LOWER.cpu().numpy()}")
    print(f"Joint upper limits: {JOINT_UPPER.cpu().numpy()}")
    within = (sampled_actions >= JOINT_LOWER) & (sampled_actions <= JOINT_UPPER)
    print(f"Actions within limits: {within.all().item()}")
    print(f"Q-values range: [{q1_vals.min().item():.4f}, {q1_vals.max().item():.4f}]")
    print(f"V-values range: [{v_vals.min().item():.4f}, {v_vals.max().item():.4f}]")
    print(f"V < Q (expected): V_mean={v_vals.mean().item():.4f}, Q_mean={q1_vals.mean().item():.4f}")

    # Verify checkpoint loads
    test_trainer = IQLTrainer(dataset, device)
    test_trainer.load(final_path)
    with torch.no_grad():
        test_v = test_trainer.v(norm_s_sample[:1])
    print(f"Checkpoint reload test: V(s0) = {test_v.item():.4f} (should match)")

    return trainer


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--tau', type=float, default=0.7)
    parser.add_argument('--beta', type=float, default=3.0)
    parser.add_argument('--lr', type=float, default=3e-4)
    parser.add_argument('--batch_size', type=int, default=256)
    parser.add_argument('--discount', type=float, default=0.99)
    parser.add_argument('--target_update_rate', type=float, default=0.005)
    parser.add_argument('--num_iterations', type=int, default=100000)
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--log_freq', type=int, default=1000)
    parser.add_argument('--save_freq', type=int, default=10000)
    args = parser.parse_args()

    train(
        tau_expectile=args.tau,
        beta=args.beta,
        lr=args.lr,
        batch_size=args.batch_size,
        discount=args.discount,
        target_update_rate=args.target_update_rate,
        num_iterations=args.num_iterations,
        seed=args.seed,
        log_freq=args.log_freq,
        save_freq=args.save_freq,
    )