import torch
import torch.nn as nn
import torch.nn.functional as F
from base import BaseModel

class GRUModel(BaseModel):
    def __init__(self, input_channels, output_channels, input_len, pred_len, hidden_channels , layers = 1):
        super(GRUModel, self).__init__()
        self.input_channels = input_channels
        self.output_channels = output_channels
        self.input_len = input_len
        self.pred_len = pred_len
        self.hidden_channels = hidden_channels
        self.layers = layers

        self.GRU = nn.GRU(input_size = self.input_channels,
                            hidden_size = self.hidden_channels,
                            num_layers = self.layers,
                            batch_first = True)

        self.out_linear = nn.Linear(in_features = self.hidden_channels,out_features = self.output_channels)

    def forward(self, x):
        x , hidden = self.GRU(x)
        x = self.out_linear(x)
        B , L , C = x.shape
        x = x.reshape(B,-1)

        return x

    def set_device(self,device):
        self.device = device

# input = torch.randn((32,96,5))
# model = GRUModel(input_channels = 5, output_channels = 1, input_len = 96, pred_len = 96, hidden_channels = 64 , layers = 1)
# out = model(input)
# print(out.shape)