# -*- coding: utf-8 -*-
"""
Created on Wed Apr 21 21:36:00 2021

@author: arash
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
#from torchvision import transforms




class Net(nn.Module):
    def __init__(self,input_dim=17481, window=100):
        super().__init__()
        self.window=window
        self.input = nn.Sequential(nn.Linear(input_dim, 1000), nn.ReLU(), nn.Linear(1000, 500), nn.ReLU())
        self.num_layers=2
        self.rnn = nn.LSTM(500, 500, num_layers=self.num_layers, batch_first=True, dropout=0.2, bidirectional=True)
        self.batch_size=256
        self.output = nn.Sequential(nn.Linear(1000,500), nn.ReLU(), nn.Linear(500, 1))
        self.tar=nn.Identity()
        
 
    def forward(self, x):
         
        self.window=x.size()[0]
        H=self.input(x)
        y,(h, c)=self.rnn(H, (self.initHidden(),self.initHidden()))
        
        return self.output(torch.cat((y[:,0,500:],y[:,-1,:500]),1))
    
    def initHidden(self):
        return torch.zeros((2*self.num_layers, self.window,500)).cuda()
    