import itertools
from torch import nn
from torch.nn import functional as F
from torch import optim

import numpy as np
import torch
from torch import distributions

import src.pytorch_util as ptu


class ValueCritic(nn.Module):
    """Value network, which takes an observation and outputs a value for that observation."""

    def __init__(
        self,
        ob_dim: int,
        n_layers: int,
        layer_size: int,
        learning_rate: float,
    ):
        super().__init__()

        self.network = ptu.build_mlp(
            input_size=ob_dim,
            output_size=1,
            n_layers=n_layers,
            size=layer_size,
        ).to(ptu.device)

        self.optimizer = optim.Adam(
            self.network.parameters(),
            learning_rate,
        )

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        # implement the forward pass of the critic network

        values = None

        ############################
        # YOUR IMPLEMENTATION HERE #
        assert isinstance(obs, torch.Tensor), "obs must be a torch tensor"
        # squeeze the last dimension to get the values as 1D tensor
        values=self.network.forward(obs).squeeze(dim=-1)
        ############################

        return values
        

    def update(self, obs: np.ndarray, q_values: np.ndarray) -> dict:
        obs = ptu.from_numpy(obs)
        q_values = ptu.from_numpy(q_values)

        # compute loss, update the critic using the observations and q_values
        loss = None
        ############################
        # YOUR IMPLEMENTATION HERE #
        values = self.forward(obs)
        # use mean squared error loss
        loss = torch.mean(torch.square(q_values-values))

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

        return {
            "Baseline Loss": ptu.to_numpy(loss),
        }