import torch
import torch.nn as nn
from torch.nn.utils import weight_norm
from torch.nn import functional as F

class Chomp1d(nn.Module):
    def __init__(self, chomp_size):
        super(Chomp1d, self).__init__()
        self.chomp_size = chomp_size

    def forward(self, x):
        return x[:, :, : -self.chomp_size].contiguous()


class Conv(nn.Module):
    """Causal convolution layer."""

    def __init__(
        self,
        n_inputs,
        n_outputs,
        kernel_size,
        stride,
        dilation,
        padding,
        causal=True,
    ):
        super().__init__()
        self.conv = nn.Conv1d(
            n_inputs,
            n_outputs,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
        )
        self.chomp = Chomp1d(padding)
        self.causal = causal

    def forward(self, x):
        out = self.conv(x)
        if self.causal:
            out = self.chomp(out)
        return out


class TemporalBlock(nn.Module):
    def __init__(
        self,
        n_inputs,
        n_outputs,
        kernel_size,
        stride,
        dilation,
        padding,
        dropout=0.2,
    ):
        super(TemporalBlock, self).__init__()
        self.conv1 = Conv(
            n_inputs,
            n_outputs,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
        )
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(dropout)

        self.conv2 = Conv(
            n_outputs,
            n_outputs,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
        )
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(dropout)

        self.net = nn.Sequential(
            self.conv1, self.relu1, self.dropout1, self.conv2, self.relu2, self.dropout2
        )
        self.downsample = (
            Conv(
                n_inputs,
                n_outputs,
                1,
                stride=1,
                padding=0,
                dilation=1,
                causal=False,
            )
            if n_inputs != n_outputs
            else None
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        out = x
        for i, l in enumerate(self.net):
            out = l(out)
        res = x if self.downsample is None else self.downsample(x)
        return self.relu(out + res)


class TemporalConvNet(nn.Module):
    """
    See:
    - https://arxiv.org/pdf/1803.01271.pdf
    - https://github.com/locuslab/TCN
    """
    def __init__(
        self,
        num_inputs,
        num_channels,
        kernel_size=2,
        dropout=0.2,
        embedding_dim=2,
    ):
        super(TemporalConvNet, self).__init__()
        layers = []
        num_levels = len(num_channels)
        for i in range(num_levels):
            dilation_size = 2 ** i
            in_channels = num_inputs if i == 0 else num_channels[i - 1]
            out_channels = num_channels[i]
            layers += [
                TemporalBlock(
                    in_channels,
                    out_channels,
                    kernel_size,
                    stride=1,
                    dilation=dilation_size,
                    padding=(kernel_size - 1) * dilation_size,
                    dropout=dropout,
                )
            ]
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        out = x
        for l in self.network:
            out = l(out)
        return out


class TCNSeq(nn.Module):
    """
    See:
    - https://arxiv.org/pdf/1803.01271.pdf
    - https://github.com/locuslab/TCN
    """
    def __init__(
        self,
        x_dim,
        y_dim,
        hidden_size=32,
        nlayers=6,
        kernel_size=2,
        dropout=0.2,
    ):
        super().__init__()
        self.tcn = TemporalConvNet(
            num_inputs=x_dim + y_dim,
            kernel_size=kernel_size,
            num_channels=[hidden_size] * nlayers,
            dropout=dropout)
        self._min_std = 0.01
        self.mean = nn.Linear(hidden_size, y_dim)
        self.std = nn.Linear(hidden_size, y_dim)

    def forward(self, past_x, past_y, future_x, future_y=None):
        device = next(self.parameters()).device
        B, S, _ = future_x.shape
        future_y_fake = past_y[:, -1:, :].repeat(1, S, 1).to(device)
        context = torch.cat([past_x, past_y], -1)
        target = torch.cat([future_x, future_y_fake], -1)
        x = torch.cat([context, target * 1], 1).detach()

        out = self.tcn(x.permute(0, 2, 1)).permute(0, 2, 1)
        
        # Seems to help a little, especially with extrapolating out of bounds
        steps = past_y.shape[1]
        mean = self.mean(out)[:, steps:, :]
        log_sigma = self.std(out)[:, steps:, :]
        
        sigma = self._min_std + (1 - self._min_std) * F.softplus(log_sigma)
        return torch.distributions.Normal(mean, sigma), {}