"""
Implementation of buffered end-to-end filter.
"""
from typing import Tuple

import torch

from motrack_motion.datasets import transforms
from motrack_motion.filter.base import StateModelFilter, State
from motrack_motion.filter.buffer import ObservationBuffer
from motrack_motion.models.architectures.end_to_end import FilterModule


class BufferedE2EFilter(StateModelFilter):
    """
    Implementation of buffered end-to-end
    """
    def __init__(
        self,
        model: FilterModule,
        transform: transforms.InvertibleTransformWithVariance,
        accelerator: str,

        buffer_size: int,
        buffer_min_size: int = 1,
        buffer_min_history: int = 5,
        dtype: torch.dtype = torch.float32,

        recursive_inverse: bool = False
    ):
        """
        Args:
            model: Model
            transform: Transform function
            accelerator: cpu/gpu accelerator
            buffer_size: Buffer size
            buffer_min_size: Minimum size to run model (default 1)
            buffer_min_history: Minimum history to save
            dtype: Torch datatype
            recursive_inverse: Apply recursive inverse transform
        """
        self._model = model
        self._transform = transform
        self._accelerator = accelerator
        self._dtype = dtype

        self._buffer_size = buffer_size
        self._buffer_min_size = buffer_min_size
        self._buffer_min_history = buffer_min_history

        # Model
        self._model.to(self._accelerator)
        self._model.eval()

        self._recursive_inverse = recursive_inverse

    def _get_ts(self) -> Tuple[torch.Tensor, torch.Tensor]:
        t_obs = torch.tensor([0], dtype=self._dtype).view(1, 1)
        t_unobs = torch.tensor([1], dtype=self._dtype).view(1, 1)
        return t_obs, t_unobs

    def initiate(self, measurement: torch.Tensor) -> State:
        buffer = ObservationBuffer(
            size=self._buffer_size,
            min_size=self._buffer_min_size,
            min_history=self._buffer_min_history,
            dtype=self._dtype
        )
        buffer.push(measurement)
        return buffer, None, None, None

    @torch.no_grad()
    def multistep_predict(self, state: State, n_steps: int) -> State:
        assert n_steps == 1, 'Multistep prediction is not supported!'

        buffer, _, _, _ = state
        if not buffer.has_input:
            raise BufferError('Buffer does not have an input!')

        x_obs, ts_obs, ts_unobs = buffer.get_input(n_steps)
        if ts_obs.shape[0] == 1:
            # Only one bbox in history - using baseline (propagate last bbox) instead of NN model
            x_unobs_mean_hat = torch.stack([x_obs[-1].clone() for _ in range(ts_unobs.shape[0])]).to(ts_obs)
            x_unobs_std_hat = 10 * torch.ones_like(x_unobs_mean_hat).to(ts_obs)

            z0, _ = self._model.initialize(batch_size=1, device=self._accelerator)
            t_obs, t_unobs = self._get_ts()
            t_obs, t_unobs = t_obs.to(self._accelerator), t_unobs.to(self._accelerator)
            _, _, z1_prior = self._model.estimate_prior(z0, t_obs, t_unobs)
            z1_prior = z1_prior.detach().cpu()

            return buffer, x_unobs_mean_hat[:, 0, :], x_unobs_std_hat[:, 0, :], z1_prior

        ts_first = ts_obs[0, 0, 0]
        ts_obs = ts_obs - ts_first + 1
        ts_unobs = ts_unobs - ts_first + 1
        t_obs_last, t_unobs_last = ts_obs[-1, 0, 0], ts_unobs[-1, 0, 0]
        if self._recursive_inverse:
            ts_unobs = torch.tensor(list(range(int(t_obs_last) + 1, int(t_unobs_last) + 1)), dtype=torch.float32).view(-1, 1, 1)

        t_x_obs, _, t_ts_obs, t_ts_unobs, *_ = self._transform.apply(data=[x_obs, None, ts_obs, ts_unobs, None], shallow=False)
        t_x_obs, t_ts_obs, t_ts_unobs = t_x_obs.to(self._accelerator), t_ts_obs.to(self._accelerator), t_ts_unobs.to(
            self._accelerator)
        x_unobs_dummy = torch.zeros(1, 1, 4).to(self._accelerator)
        t_x_prior_mean, t_x_prior_var, _, _, z_prior = \
            self._model.forward(t_x_obs, t_ts_obs, x_unobs_dummy, t_ts_unobs, mask=False)
        x_unobs_dummy.detach().cpu()
        t_x_prior_mean = t_x_prior_mean.detach().cpu()
        t_x_prior_mean = t_x_prior_mean.detach().cpu()
        z_prior = z_prior.detach().cpu()

        _, prior_mean, *_ = self._transform.inverse(data=[x_obs, t_x_prior_mean, None], shallow=False)

        prior_mean = prior_mean[-1]  # Removing temporal dimension after inverse transform
        prior_var = self._transform.inverse_var(t_x_prior_var, additional_data=[x_obs, None], shallow=False)
        prior_std = torch.sqrt(prior_var)[0]

        return buffer, prior_mean, prior_std, z_prior

    def predict(self, state: State) -> State:
        buffer, prior_mean, prior_std, z1_prior = self.multistep_predict(state, n_steps=1)
        return buffer, prior_mean[0], prior_std[0], z1_prior

    def singlestep_to_multistep_state(self, state: State) -> State:
        buffer, prior_mean, prior_std, z1_prior = state
        return buffer, prior_mean.unsqueeze(0), prior_std.unsqueeze(0), z1_prior

    @torch.no_grad()
    def update(self, state: State, measurement: torch.Tensor) -> State:
        buffer, _, _, z_prior = state
        x_obs, ts_obs, ts_unobs = buffer.get_input(1)

        if ts_obs.shape[0] == 1:
            # Only one bbox in history - using baseline (propagate last bbox) instead of NN model
            x_unobs_mean_hat = torch.stack([x_obs[-1].clone() for _ in range(ts_unobs.shape[0])]).to(ts_obs)
            x_unobs_std_hat = 10 * torch.ones_like(x_unobs_mean_hat).to(ts_obs)

            z0, _ = self._model.initialize(batch_size=1, device=self._accelerator)
            t_obs, t_unobs = self._get_ts()
            t_obs, t_unobs = t_obs.to(self._accelerator), t_unobs.to(self._accelerator)
            _, _, z1_prior = self._model.estimate_prior(z0, t_obs, t_unobs)
            z1_prior = z1_prior.detach().cpu()

            buffer.push(measurement)
            return buffer, x_unobs_mean_hat[0, 0, :], x_unobs_std_hat[0, 0, :], z1_prior

        _, t_measurement, *_ = self._transform.apply(data=[x_obs, measurement.view(1, 1, -1), ts_obs, ts_unobs, None], shallow=False)
        t_measurement, z_prior = t_measurement[0].to(self._accelerator), z_prior.to(self._accelerator)
        z_evidence = self._model.encode_unobs(t_measurement)
        t_posterior_mean, t_posterior_log_var, z_posterior = self._model.estimate_posterior(z_prior, z_evidence)
        t_posterior_mean, t_posterior_log_var, z_posterior = \
            t_posterior_mean.detach().cpu(), t_posterior_log_var.detach().cpu(), z_posterior.detach().cpu()
        t_posterior_var = self._model.postprocess_log_var(t_posterior_log_var)
        _, posterior_mean, *_ = self._transform.inverse(data=[x_obs, t_posterior_mean.view(1, 1, -1), None], shallow=False)
        posterior_var = self._transform.inverse_var(t_posterior_var, additional_data=[posterior_mean, None], shallow=False)
        posterior_std = torch.sqrt(posterior_var)

        posterior_mean = posterior_mean[0, 0, :]
        posterior_std = posterior_std[0, :]

        buffer.push(measurement)

        return buffer, posterior_mean, posterior_std, z_posterior

    def missing(self, state: State) -> State:
        buffer, mean, std, z_state = state
        buffer.increment()
        return buffer, mean, std, z_state

    def project(self, state: State) -> Tuple[torch.Tensor, torch.Tensor]:
        _, mean, std, _ = state
        var = torch.square(std)
        return mean, var


def run_test() -> None:
    # pylint: disable=import-outside-toplevel
    from motrack_motion.datasets.transforms.base import IdentityTransform
    from motrack_motion.models.architectures.rnn_filter import RNNFilterModel

    smf = BufferedE2EFilter(
        model=RNNFilterModel(
            observable_dim=4,
            latent_dim=4,
            output_dim=4
        ),
        buffer_size=5,
        accelerator='cpu',
        transform=IdentityTransform()
    )

    measurements = torch.randn(10, 4, dtype=torch.float32)

    # Initiate test
    state = smf.initiate(measurements[0])

    # Predict test
    state = smf.predict(state)
    mean_hat, cov_hat = smf.project(state)
    assert mean_hat.shape == (4,) and cov_hat.shape == (4,)

    # Update test
    state = smf.update(state, measurements[1])
    mean_updated, cov_updated = smf.project(state)
    assert mean_updated.shape == (4,) and cov_updated.shape == (4,)


if __name__ == '__main__':
    run_test()
