"""
Action-conditioned world model with reconstruction.

The model learns:
    observation_t -> latent_t
    latent_t + action_t -> latent_t+1
    latent_t+1 -> predicted observation_t+1
"""

import torch
from torch import nn


class ActionConditionedWorldModel(nn.Module):
    def __init__(
        self,
        obs_dim: int,
        action_dim: int,
        latent_dim: int = 32,
        hidden_dim: int = 128,
    ):
        super().__init__()

        self.encoder = nn.Sequential(
            nn.Linear(obs_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim),
        )

        self.action_encoder = nn.Sequential(
            nn.Linear(action_dim, 32),
            nn.ReLU(),
        )

        self.dynamics = nn.Sequential(
            nn.Linear(latent_dim + 32, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim),
        )

        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, obs_dim),
        )

    def step(
        self,
        latent: torch.Tensor,
        action: torch.Tensor,
    ) -> torch.Tensor:
        action_features = self.action_encoder(action)
        x = torch.cat([latent, action_features], dim=-1)
        delta = self.dynamics(x)
        return latent + delta

    def forward(
        self,
        obs: torch.Tensor,
        action: torch.Tensor,
    ):
        latent = self.encoder(obs)
        next_latent = self.step(latent, action)
        next_obs_pred = self.decoder(next_latent)
        return next_obs_pred, next_latent


if __name__ == "__main__":
    torch.manual_seed(0)

    model = ActionConditionedWorldModel(
        obs_dim=12,
        action_dim=3,
        latent_dim=24,
    )

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    obs_t = torch.randn(32, 12)
    action_t = torch.randn(32, 3)
    obs_t1 = torch.randn(32, 12)

    pred_obs_t1, _ = model(obs_t, action_t)
    loss = nn.functional.mse_loss(pred_obs_t1, obs_t1)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    print("training loss:", float(loss))