"""Implementation of the sequence-based models using PyTorch."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/models/torch/seq.ipynb.

# %% auto 0
__all__ = ['S2P', 'S2PDataset', 'S2PDatasetMains']

# %% ../../../nbs/models/torch/seq.ipynb 3
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

import enilm.models.torch.utils as utils

# %% ../../../nbs/models/torch/seq.ipynb 9
class S2P(nn.Module):
    def __init__(self, seq_len):
        super(S2P, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=1,  out_channels=30, kernel_size=10, stride=1)
        self.conv2 = nn.Conv1d(in_channels=30, out_channels=30, kernel_size=8,  stride=1)
        self.conv3 = nn.Conv1d(in_channels=30, out_channels=40, kernel_size=6,  stride=1)
        self.conv4 = nn.Conv1d(in_channels=40, out_channels=50, kernel_size=5,  stride=1)
        self.drop1 = nn.Dropout(0.2)
        self.conv5 = nn.Conv1d(in_channels=50, out_channels=50, kernel_size=5,  stride=1)
        self.drop2 = nn.Dropout(0.2)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(
            in_features=((seq_len - self.conv1.out_channels) + 1) * self.conv5.out_channels, 
            out_features=1024,
        )
        # self.fc1 = nn.Linear(50 * ((seq_len - 4*10 - 3*8 - 2*6 - 5*1 + 12) // 1), 1024)
        self.drop3 = nn.Dropout(0.2)
        self.fc2 = nn.Linear(in_features=self.fc1.out_features, out_features=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = self.drop1(x)
        x = F.relu(self.conv5(x))
        x = self.drop2(x)
        x = self.flatten(x)
        x = F.relu(self.fc1(x))
        x = self.drop3(x)
        x = self.fc2(x)
        return x

# %% ../../../nbs/models/torch/seq.ipynb 15
class S2PDataset(Dataset):
    def __init__(
        self,
        mains: np.ndarray,
        appliance: np.ndarray,
        sequence_length: int,
        pad=True,  # whether to pad the sequence or not (e.g. 1, 2, 3, 4, 5, ... -> 0, 0, 1, 2, 3 if pad is True for sequence_length=5 and 1, 2, 3, 4, 5 if pad is False)
        reshape=True,  # whether to reshape the inputs and targets to the expected shape of the model
    ):
        assert sequence_length % 2 == 1
        assert isinstance(mains, np.ndarray)
        assert isinstance(appliance, np.ndarray)
        assert mains.shape == appliance.shape
        self.sequence_length = sequence_length
        self.mains = torch.tensor(mains)
        self.appliance = torch.tensor(appliance)
        self.pad = pad
        self.reshape = reshape

    def __len__(self):
        if self.pad:
            return len(self.mains)
        return len(self.mains) - self.sequence_length + 1

    def _reshape(self, mains, ground_truth):
        return (
            # reshape inputs: add a channel dimension
            mains.reshape(-1, self.sequence_length),
            # reshape targets: since the model outputs a single value per sequence, add the extra singularity dimension
            ground_truth.reshape(
                -1,
            ),
        )

    def __getitem__(self, idx):
        # stop iteration if idx is out of bounds
        if idx == len(self.mains):
            raise StopIteration

        # padding or truncating
        if self.pad:
            mains, gt = (
                utils.get_padded_sequence(self.mains, self.sequence_length, idx),
                self.appliance[idx],
            )
        else:
            if idx > len(self.mains) - self.sequence_length:
                raise StopIteration
            half_seq_len = self.sequence_length // 2
            idx += half_seq_len
            mains, gt = (
                self.mains[idx - half_seq_len : idx + half_seq_len + 1],
                self.appliance[idx],
            )

        # reshape?
        if self.reshape:
            mains, gt = self._reshape(mains, gt)

        return mains, gt

# %% ../../../nbs/models/torch/seq.ipynb 82
class S2PDatasetMains(Dataset):
    def __init__(
        self,
        mains: np.ndarray,
        sequence_length: int,
        pad=True,  # whether to pad the sequence or not (e.g. 1, 2, 3, 4, 5, ... -> 0, 0, 1, 2, 3 if pad is True for sequence_length=5 and 1, 2, 3, 4, 5 if pad is False)
        reshape=True,  # whether to reshape the inputs and targets to the expected shape of the model
    ):
        assert sequence_length % 2 == 1
        assert isinstance(mains, np.ndarray)
        self.sequence_length = sequence_length
        self.mains = torch.tensor(mains)
        self.pad = pad
        self.reshape = reshape

    def __len__(self):
        if self.pad:
            return len(self.mains)
        return len(self.mains) - self.sequence_length + 1

    def _reshape(self, mains):
        # reshape inputs: add a channel dimension
        return mains.reshape(-1, self.sequence_length)

    def __getitem__(self, idx):
        # stop iteration if idx is out of bounds
        if idx == len(self.mains):
            raise StopIteration

        # padding or truncating
        if self.pad:
            mains = utils.get_padded_sequence(
                self.mains, self.sequence_length, idx
            )
        else:
            if idx > len(self.mains) - self.sequence_length:
                raise StopIteration
            half_seq_len = self.sequence_length // 2
            idx += half_seq_len
            mains = self.mains[idx - half_seq_len : idx + half_seq_len + 1]

        # reshape?
        if self.reshape:
            mains = self._reshape(mains)

        return mains
