# NLinear 정의
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class NLinear_Model(nn.Module):
    """
    2023 SOTA 단순 Linear layer 기반 모델. 시퀀스 마지막 값을 추후 더하는 형식
    """
    def __init__(self, configs):
        super(NLinear_Model, self).__init__()
        self.seq_len = configs.seq_len
        self.pred_len = configs.pred_len
        self.channels = configs.enc_in
        self.individual = configs.individual
        if self.individual:
            self.Linear = nn.ModuleList()
            for i in range(self.channels):
                self.Linear.append(nn.Linear(self.seq_len,self.pred_len))
        else:
            self.Linear = nn.Linear(self.seq_len, self.pred_len)

    def forward(self, x):
        # 입력 sequence 중 마지막 시퀀스 입력
        seq_last = x[:,-1:,:].detach()
        x = x - seq_last
        #individual은 Bool type. 단일 값 예측 or not
        if self.individual:
            output = torch.zeros([x.size(0),self.pred_len,x.size(2)],dtype=x.dtype).to(x.device)
            for i in range(self.channels):
                output[:,:,i] = self.Linear[i](x[:,:,i])
            x = output
        else:
            x = self.Linear(x.permute(0,2,1)).permute(0,2,1)
        x = x + seq_last
        return x # [시퀀스 총 개수 , 개별 시퀀스 길이 , 특성 개수]
