"""
Temporal Convolutional Network (TCN) untuk sequence classification.

Berisi implementasi model Temporal Convolutional Network sesuai README.

"""

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


class Chomp1d(nn.Module):
    """
    Remove padding dari akhir sequence.
    """
    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 TemporalBlock(nn.Module):
    """
    Temporal block dari TCN.
    Terdiri dari dua dilated conv layers dengan residual connection.
    """
    def __init__(
        self,
        n_inputs,
        n_outputs,
        kernel_size,
        stride,
        dilation,
        padding,
        dropout=0.2
    ):
        super(TemporalBlock, self).__init__()

        # First dilated convolution
        self.conv1 = weight_norm(nn.Conv1d(
            n_inputs, n_outputs, kernel_size,
            stride=stride, padding=padding, dilation=dilation
        ))
        self.chomp1 = Chomp1d(padding)
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(dropout)

        # Second dilated convolution
        self.conv2 = weight_norm(nn.Conv1d(
            n_outputs, n_outputs, kernel_size,
            stride=stride, padding=padding, dilation=dilation
        ))
        self.chomp2 = Chomp1d(padding)
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(dropout)

        # Residual connection
        self.net = nn.Sequential(
            self.conv1, self.chomp1, self.relu1, self.dropout1,
            self.conv2, self.chomp2, self.relu2, self.dropout2
        )

        # Downsample jika perlu
        self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None
        self.relu = nn.ReLU()

        # Initialize weights
        self.conv1.weight.data.normal_(0, 0.01)
        self.conv2.weight.data.normal_(0, 0.01)

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


class TemporalConvNet(nn.Module):
    """
    Temporal Convolutional Network.
    Stack dari temporal blocks dengan increasing dilation.
    """
    def __init__(
        self,
        num_inputs,
        num_channels,
        kernel_size=3,
        dropout=0.2
    ):
        """
        Args:
            num_inputs: Jumlah input channels
            num_channels: List of hidden layer sizes
            kernel_size: Kernel size untuk convolutions
            dropout: Dropout rate
        """
        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.append(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):
        return self.network(x)


class TCNClassifier(nn.Module):
    """
    TCN-based classifier untuk gesture recognition.

    Sesuai README:
    - Input: Sequence dari landmark (30, 306)
    - Output: Probability untuk setiap kelas
    """
    def __init__(
        self,
        input_size: int,
        num_classes: int,
        num_channels: list = None,
        kernel_size: int = 3,
        dropout: float = 0.2,
        fc_hidden_size: int = 512
    ):
        """
        Args:
            input_size: Ukuran input features (306)
            num_classes: Jumlah kelas output
            num_channels: Hidden layer sizes
            kernel_size: Kernel size
            dropout: Dropout rate
            fc_hidden_size: Hidden size untuk fully connected layer
        """
        super(TCNClassifier, self).__init__()

        if num_channels is None:
            num_channels = [64, 128, 256]

        self.tcn = TemporalConvNet(
            num_inputs=input_size,
            num_channels=num_channels,
            kernel_size=kernel_size,
            dropout=dropout
        )

        tcn_output_size = num_channels[-1]

        # Fully connected layers
        self.fc = nn.Sequential(
            nn.Linear(tcn_output_size, fc_hidden_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(fc_hidden_size, num_classes)
        )

    def forward(self, x):
        # x shape: (batch, seq_len, features)
        # Transpose ke (batch, features, seq_len) untuk Conv1d
        x = x.transpose(1, 2)

        # TCN
        tcn_out = self.tcn(x)

        # Global average pooling
        pooled = tcn_out.mean(dim=2)

        # Classify
        out = self.fc(pooled)

        return out