import random

from tqdm import tqdm

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from vp_suite.base.base_model import VideoPredictionModel
from vp_suite.model_blocks import Autoencoder
from vp_suite.model_blocks.predrnn import SpatioTemporalLSTMCell, ActionConditionalSpatioTemporalLSTMCell
from vp_suite.model_blocks.phydnet import PhyCell_Cell, K2M


class STPhy(VideoPredictionModel):
    r"""

    """

    # model-specific constants
    NAME = "ST-Phy"
    CAN_HANDLE_ACTIONS = True

    # model hyperparameters
    num_layers = 3  #: TODO
    phycell_kernel_size = (7, 7)  #: TODO
    phycell_channels = 49  #: TODO
    st_cell_channels = 64  #: TODO
    inflated_action_dim = 3  #: TODO

    # model training hyperparams
    reconstruction_loss_scale = 0.1  #: TODO
    decoupling_loss_scale = 100.0  #: TODO
    moment_loss_scale = 1.0  #: TODO
    teacher_forcing_decay = 0.003

    def __init__(self, device, **model_kwargs):
        r"""

        Args:
            device ():
            **model_kwargs ():
        """
        super(STPhy, self).__init__(device, **model_kwargs)

        self.dim_st_hidden = [self.st_cell_channels] * self.num_layers
        self.dim_phy_hidden = [self.phycell_channels] * self.num_layers

        self.autoencoder = Autoencoder(self.img_shape, self.st_cell_channels, self.device)
        _, _, self.enc_h, self.enc_w = self.autoencoder.encoded_shape
        self.recurrent_cell = SpatioTemporalLSTMCell

        if self.action_conditional:
            self.recurrent_cell = ActionConditionalSpatioTemporalLSTMCell
            self.action_inflate = nn.Linear(in_features=self.action_size,
                                            out_features=self.inflated_action_dim * self.enc_h * self.enc_w,
                                            bias=False)
            self.action_conv_h = nn.Conv2d(in_channels=self.inflated_action_dim, out_channels=self.st_cell_channels,
                                           kernel_size=(5, 1), padding=(2, 0), bias=False)
            self.action_conv_w = nn.Conv2d(in_channels=self.inflated_action_dim, out_channels=self.st_cell_channels,
                                           kernel_size=(1, 5), padding=(0, 2), bias=False)

        st_cells, phycells, hidden_convs = [], [], []
        for i in range(self.num_layers):
            cell_in_channel = self.dim_st_hidden[0] if i == 0 else self.dim_st_hidden[i - 1]
            st_cells.append(self.recurrent_cell(cell_in_channel, self.dim_st_hidden[i], self.enc_h, self.enc_w,
                                                filter_size=5, stride=1, layer_norm=True))

            phycells.append(PhyCell_Cell(input_dim=cell_in_channel, action_conditional=self.action_conditional,
                                         action_size=self.action_size, hidden_dim=self.dim_phy_hidden[i],
                                         kernel_size=self.phycell_kernel_size).to(self.device))

            hc_bias = i < self.num_layers - 1
            hidden_convs.append(nn.Conv2d(in_channels=self.st_cell_channels + self.dim_st_hidden[i],
                                          out_channels=self.st_cell_channels, kernel_size=(1, 1), bias=hc_bias))

        self.st_cell_list = nn.ModuleList(st_cells)
        self.phycell_list = nn.ModuleList(phycells)
        self.hidden_conv_list = nn.ModuleList(hidden_convs)

        # shared adapter
        adapter_dim_hidden = self.dim_st_hidden[0]
        self.adapter = nn.Conv2d(adapter_dim_hidden, adapter_dim_hidden, 1, stride=1, padding=0, bias=False)

        self.constraints = torch.zeros((self.phycell_channels, *self.phycell_kernel_size), device=self.device)
        ind = 0
        for i in range(0, self.phycell_kernel_size[0]):
            for j in range(0, self.phycell_kernel_size[1]):
                self.constraints[ind, i, j] = 1
                ind += 1

    def pred_1(self, x, **kwargs):
        r"""

        Args:
            x ():
            **kwargs ():

        Returns:

        """
        return self(x, pred_frames=1, **kwargs)[0].squeeze(dim=1)

    def forward(self, x, pred_frames=1, **kwargs):
        r"""

        Args:
            x ():
            pred_frames ():
            **kwargs ():

        Returns:

        """
        # in training mode (default: False), returned sequence starts with 2nd context frame,
        # and the moment regularization loss is calculated.
        train = kwargs.get("train", False)
        teacher_forcing = kwargs.get("teacher_forcing", False) and train
        batch_size, context_frames = x.shape[:2]
        if train:  # during training, the whole sequence is passed
            context_frames -= pred_frames
        empty_actions = torch.zeros(x.shape[0], context_frames+pred_frames-1, device=self.device)
        actions = kwargs.get("actions", empty_actions)
        if self.action_conditional:
            if actions.equal(empty_actions) or actions.shape[-1] != self.action_size:
                raise ValueError("Given actions are None or of the wrong size!")

        # init ST and Phy
        phy_h_t = []
        st_h_t = []
        st_c_t = []
        delta_c_list = []
        delta_m_list = []

        decouple_loss = []

        for i in range(self.num_layers):
            zeros = torch.zeros([batch_size, self.dim_st_hidden[i], self.enc_h, self.enc_w]).to(self.device)
            st_h_t.append(zeros)
            st_c_t.append(zeros)
            delta_c_list.append(zeros)
            delta_m_list.append(zeros)
            phy_h_t.append(torch.zeros(batch_size, self.st_cell_channels, self.enc_h, self.enc_w).to(self.device))

        st_memory = torch.zeros([batch_size, self.dim_st_hidden[0], self.enc_h, self.enc_w]).to(self.device)
        out_frames = []
        x_gen = None

        # fwd prop through time
        for t in range(context_frames + pred_frames - 1):

            # get input, depending on training/inference stage
            if t < context_frames or teacher_forcing:
                next_input = self.autoencoder.encode(x[:, t])
            else:
                next_input = x_gen

            for i, (st_cell, phycell) in enumerate(zip(self.st_cell_list, self.phycell_list)):

                # phy
                phy_h_t[i] = phycell(next_input, actions[:, t], phy_h_t[i])  # returns (hidden, hidden)

                # st
                if self.action_conditional:
                    ac = self.action_inflate(actions[:, t]).view(-1, self.inflated_action_dim, self.enc_h, self.enc_w)
                    inflated_action = self.action_conv_h(ac) + self.action_conv_w(ac)
                    st_h_t[i], st_c_t[i], st_memory, delta_c, delta_m \
                        = st_cell(next_input, st_h_t[i], st_c_t[i], st_memory, inflated_action)
                else:
                    st_h_t[i], st_c_t[i], st_memory, delta_c, delta_m \
                        = st_cell(next_input, st_h_t[i], st_c_t[i], st_memory)
                delta_c_list[i] = F.normalize(self.adapter(delta_c).view(delta_c.shape[0], delta_c.shape[1], -1), dim=2)
                delta_m_list[i] = F.normalize(self.adapter(delta_m).view(delta_m.shape[0], delta_m.shape[1], -1), dim=2)

                # merge
                x_gen = self.hidden_conv_list[i](torch.cat([st_h_t[i], phy_h_t[i]], dim=-3))

            if train or t >= (context_frames - 1):
                out_frame = self.autoencoder.decode(x_gen)
                out_frames.append(out_frame)

            # Spatio-temporal decoupling loss in STCell during training
            if train:
                for i in range(0, self.num_layers):
                    decouple_loss.append(
                        torch.mean(torch.abs(torch.cosine_similarity(delta_c_list[i], delta_m_list[i], dim=2))))

        out_frames = torch.stack(out_frames, dim=1)

        if train:
            # Moment regularization loss during training
            k2m = K2M(self.phycell_kernel_size).to(self.device)
            moment_loss = 0
            for b in range(0, self.phycell_list[0].input_dim):
                filters = self.phycell_list[0].F.conv1.weight[:, b]
                moment = k2m(filters.double()).float()
                moment_loss += torch.mean(self.moment_loss_scale * (moment - self.constraints) ** 2)
            model_losses = {
                "moment regularization loss": moment_loss,
                "memory decoupling loss": torch.mean(torch.stack(decouple_loss, dim=0)),
            }
        else:
            model_losses = None
        return out_frames, model_losses

    def train_iter(self, config, data_loader, optimizer, loss_provider, epoch):
        r"""

        Args:
            config ():
            data_loader ():
            optimizer ():
            loss_provider ():
            epoch ():

        Returns:

        """
        teacher_forcing_ratio = np.maximum(0, 1 - epoch * self.teacher_forcing_decay)
        loop = tqdm(data_loader)
        for batch_idx, data in enumerate(loop):

            # fwd
            input_frames, _, actions = self.unpack_data(data, config, complete=True)
            teacher_forcing = True if random.random() < teacher_forcing_ratio else False
            predictions, model_losses = self(input_frames, pred_frames=config["pred_frames"],
                                             actions=actions, train=True, teacher_forcing=teacher_forcing)

            # loss
            targets = input_frames[:, 1:]
            _, total_loss = loss_provider.get_losses(predictions, targets)
            if model_losses is not None:
                for value in model_losses.values():
                    total_loss += value

            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()

            loop.set_postfix(loss=total_loss.item())
