import copy

import pytorch_lightning as pl
import torch
from torch import nn, optim
from torch.nn import functional as F

from GAT.model import GAT
from libs import utils
from TCN.model import TemporalConvNet
from libs.utils import StandardScaler


def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class Encoder(nn.Module):
    """Some Information about Encoder"""

    def __init__(self, in_features, thiddin_features, tnhid, tkernel_size, shid_features, alpha, adjs, dropout, seq_len,
                 spatial=True, temporal=True, N=6,):
        super(Encoder, self).__init__()
        self.net = clones(EncoderLayer(in_features, thiddin_features, tnhid,
                                       tkernel_size, shid_features, alpha, adjs=adjs, dropout=dropout, seq_len=seq_len,
                                       spatial=spatial, temporal=temporal), N,)

    def forward(self, x):
        for layer in self.net:
            x = layer(x)
        return x


class EncoderLayer(nn.Module):
    """Some Information about Encoder"""

    def __init__(self, in_features, thiddin_features, tnhid, tkernel_size, shid_features, alpha, adjs, dropout, seq_len, spatial, temporal):
        super(EncoderLayer, self).__init__()
        self.spatial = spatial
        self.temporal = temporal
        if temporal:
            self.temp_corr = TemporalConvNet(in_features, [
                *([thiddin_features] * tnhid), in_features],
                kernel_size=tkernel_size, dropout=dropout,)

        if spatial:
            if len(adjs.size()) == 2:
                adjs = adjs.unsqueeze(0)  # add dimension if adjs is a matrix

            self.register_buffer("adjs", adjs)
            nheads = adjs.size(0)
            self.spatial_corr = clones(GAT(in_features, shid_features, in_features,
                                           dropout, alpha, nheads), 2)
        # self.linear1 = nn.Linear(n_nodes, n_nodes)
        # self.linear2 = nn.Linear(seq_len, seq_len)

    def forward(self, x):
        # spatial correlation
        if self.spatial:
            x = x.permute(0, 3, 1, 2)
            res = x
            for layer in self.spatial_corr:
                x = layer(x, self.adjs)
            x = res + x
            x = x.permute(0, 2, 3, 1)
        # # x = self.linear1(x + out)

        # temporal correlation
        if self.temporal:
            res = x
            x = self.temp_corr(x)
            x = x + res
        # x = self.linear2(x)+res
        return x


class Decoder(nn.Module):
    """Some Information about Decoder"""

    def __init__(self, in_features, out_features):
        super(Decoder, self).__init__()
        self.mlp = nn.Sequential(
            nn.Linear(in_features, out_features * in_features),
            nn.ReLU(),
            nn.Linear(out_features * in_features, out_features),
        )

    def forward(self, x):
        x = self.mlp(x)
        return x


class TCNPredictor(pl.LightningModule):
    def __init__(self, in_features=1, thid_features=1, nhid=2, kernel_size=3, seq_len=12, horizon=12, shid_features=16,
                 alpha=0.2, dropout=0.3, lr=1e-3, weight_decay=1e-2, scaler=StandardScaler(), adjs=None, N=6, spatial=False,
                 temporal=False, **kwargs,):
        super(TCNPredictor, self).__init__()
        self.encode = Encoder(in_features, thid_features, nhid, kernel_size,
                              shid_features, alpha, adjs=adjs, dropout=dropout, N=N, seq_len=seq_len,
                              spatial=spatial, temporal=temporal)
        self.decode = Decoder(seq_len, horizon)
        self.lr = lr
        self.weight_decay = weight_decay
        self.scaler = scaler

    def forward(self, x):
        # encoder
        x = self.encode(x)

        # decoder
        pred = self.decode(x)
        return pred.squeeze(1)

    def training_step(self, batch, idx):
        x, y = batch
        x, y = x.float(), y.float()
        pred = self.forward(x)
        loss = F.l1_loss(self.scaler.inverse_transform(pred), y)
        self.log("Training/mae", loss, prog_bar=True)
        return loss

    def validation_step(self, batch, idx):
        x, y = batch
        x, y = x.float(), y.float()
        pred = self.forward(x)
        loss = F.l1_loss(self.scaler.inverse_transform(pred), y)
        return loss

    def validation_epoch_end(self, outputs):
        loss = torch.stack(outputs).mean()
        self.log('Validation/mae', loss, prog_bar=True, on_epoch=True)
        return loss

    def test_step(self, batch, idx):
        x, y = batch
        x, y = x.float(), y.float()
        pred = self.scaler.inverse_transform(self.forward(x))
        loss = {
            "mae": utils.masked_MAE(pred, y, dim=-2),
            "rmse": utils.masked_RMSE(pred, y, dim=-2),
        }
        return loss

    def test_epoch_end(self, outputs):
        loss = {"mae": 0, "rmse": 0}
        for m in loss:
            loss[m] = torch.cat([output[m]
                                 for output in outputs], dim=0).nanmean(dim=0)
        for h in range(len(loss["mae"])):
            print(f"Horizon {h+1} ({5*(h+1)} min) - ", end="")
            print(f"MAE: {loss['mae'][h]:.2f}", end=", ")
            print(f"RMSE: {loss['rmse'][h]:.2f}")
            self.logger.experiment.add_scalar(f"Test/mae", loss['mae'][h], h)
            self.logger.experiment.add_scalar(f"Test/rmse", loss['rmse'][h], h)
        print("Aggregation - ", end="")
        print(f"MAE: {loss['mae'].mean():.2f}", end=", ")
        print(f"RMSE: {loss['rmse'].mean():.2f}")

    def configure_optimizers(self):
        return optim.Adam(self.parameters(),
                          lr=self.lr,
                          weight_decay=self.weight_decay)

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = parent_parser.add_argument_group("TCNPredictor")
        parser.add_argument("--batch_size", type=int, default=64, help="")
        parser.add_argument("--horizon", type=int, default=12, help="")
        parser.add_argument("--lookback_window", type=int, default=2, help="")
        parser.add_argument("--dropout", type=int, default=0.3, help="")
        parser.add_argument("--lr", type=float, default=1e-3, help="")
        parser.add_argument("--weight_decay",
                            type=float,
                            default=1e-2,
                            help="")
        parser.add_argument("--in_features", type=int, default=1)
        parser.add_argument("--thid_features", type=int, default=5)
        parser.add_argument("--shid_features", type=int, default=5)
        parser.add_argument('--spatial', dest='spatial',
                            action='store_true')
        parser.add_argument('--temporal', dest='temporal',
                            action='store_true')
        parser.add_argument("--nhid", type=int, default=3)
        parser.add_argument("--kernel_size", type=int, default=3)
        parser.add_argument("--N", type=int, default=3)
        return parent_parser

import copy

import pytorch_lightning as pl
import torch
from torch import nn, optim
from torch.nn import functional as F

from GAT.model import GAT
from libs import utils
from TCN.model import TemporalConvNet
from libs.utils import StandardScaler


def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class Encoder(nn.Module):
    """Some Information about Encoder"""

    def __init__(self, in_features, thiddin_features, tnhid, tkernel_size, shid_features, alpha, adjs, dropout, seq_len,
                 spatial=True, temporal=True, N=6,):
        super(Encoder, self).__init__()
        self.net = clones(EncoderLayer(in_features, thiddin_features, tnhid,
                                       tkernel_size, shid_features, alpha, adjs=adjs, dropout=dropout, seq_len=seq_len,
                                       spatial=spatial, temporal=temporal), N,)

    def forward(self, x):
        for layer in self.net:
            x = layer(x)
        return x


class EncoderLayer(nn.Module):
    """Some Information about Encoder"""

    def __init__(self, in_features, thiddin_features, tnhid, tkernel_size, shid_features, alpha, adjs, dropout, seq_len, spatial, temporal):
        super(EncoderLayer, self).__init__()
        self.spatial = spatial
        self.temporal = temporal
        if temporal:
            self.temp_corr = TemporalConvNet(in_features, [
                *([thiddin_features] * tnhid), in_features],
                kernel_size=tkernel_size, dropout=dropout,)

        if spatial:
            if len(adjs.size()) == 2:
                adjs = adjs.unsqueeze(0)  # add dimension if adjs is a matrix

            self.register_buffer("adjs", adjs)
            nheads = adjs.size(0)
            self.spatial_corr = clones(GAT(in_features, shid_features, in_features,
                                           dropout, alpha, nheads), 2)
        # self.linear1 = nn.Linear(n_nodes, n_nodes)
        # self.linear2 = nn.Linear(seq_len, seq_len)

    def forward(self, x):
        # spatial correlation
        if self.spatial:
            x = x.permute(0, 3, 1, 2)
            res = x
            for layer in self.spatial_corr:
                x = layer(x, self.adjs)
            x = res + x
            x = x.permute(0, 2, 3, 1)
        # # x = self.linear1(x + out)

        # temporal correlation
        if self.temporal:
            res = x
            x = self.temp_corr(x)
            x = x + res
        # x = self.linear2(x)+res
        return x


class Decoder(nn.Module):
    """Some Information about Decoder"""

    def __init__(self, in_features, out_features):
        super(Decoder, self).__init__()
        self.mlp = nn.Sequential(
            nn.Linear(in_features, out_features * in_features),
            nn.ReLU(),
            nn.Linear(out_features * in_features, out_features),
        )

    def forward(self, x):
        x = self.mlp(x)
        return x


class TCNPredictor(pl.LightningModule):
    def __init__(self, in_features=1, thid_features=1, nhid=2, kernel_size=3, seq_len=12, horizon=12, shid_features=16,
                 alpha=0.2, dropout=0.3, lr=1e-3, weight_decay=1e-2, scaler=StandardScaler(), adjs=None, N=6, spatial=False,
                 temporal=False, **kwargs,):
        super(TCNPredictor, self).__init__()
        self.encode = Encoder(in_features, thid_features, nhid, kernel_size,
                              shid_features, alpha, adjs=adjs, dropout=dropout, N=N, seq_len=seq_len,
                              spatial=spatial, temporal=temporal)
        self.decode = Decoder(seq_len, horizon)
        self.lr = lr
        self.weight_decay = weight_decay
        self.scaler = scaler

    def forward(self, x):
        # encoder
        x = self.encode(x)

        # decoder
        pred = self.decode(x)
        return pred.squeeze(1)

    def training_step(self, batch, idx):
        x, y = batch
        x, y = x.float(), y.float()
        pred = self.forward(x)
        loss = F.l1_loss(self.scaler.inverse_transform(pred), y)
        self.log("Training/mae", loss, prog_bar=True)
        return loss

    def validation_step(self, batch, idx):
        x, y = batch
        x, y = x.float(), y.float()
        pred = self.forward(x)
        loss = F.l1_loss(self.scaler.inverse_transform(pred), y)
        return loss

    def validation_epoch_end(self, outputs):
        loss = torch.stack(outputs).mean()
        self.log('Validation/mae', loss, prog_bar=True, on_epoch=True)
        return loss

    def test_step(self, batch, idx):
        x, y = batch
        x, y = x.float(), y.float()
        pred = self.scaler.inverse_transform(self.forward(x))
        loss = {
            "mae": utils.masked_MAE(pred, y, dim=-2),
            "rmse": utils.masked_RMSE(pred, y, dim=-2),
        }
        return loss

    def test_epoch_end(self, outputs):
        loss = {"mae": 0, "rmse": 0}
        for m in loss:
            loss[m] = torch.cat([output[m]
                                 for output in outputs], dim=0).nanmean(dim=0)
        for h in range(len(loss["mae"])):
            print(f"Horizon {h+1} ({5*(h+1)} min) - ", end="")
            print(f"MAE: {loss['mae'][h]:.2f}", end=", ")
            print(f"RMSE: {loss['rmse'][h]:.2f}")
            self.logger.experiment.add_scalar(f"Test/mae", loss['mae'][h], h)
            self.logger.experiment.add_scalar(f"Test/rmse", loss['rmse'][h], h)
        print("Aggregation - ", end="")
        print(f"MAE: {loss['mae'].mean():.2f}", end=", ")
        print(f"RMSE: {loss['rmse'].mean():.2f}")

    def configure_optimizers(self):
        return optim.Adam(self.parameters(),
                          lr=self.lr,
                          weight_decay=self.weight_decay)

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = parent_parser.add_argument_group("TCNPredictor")
        parser.add_argument("--batch_size", type=int, default=64, help="")
        parser.add_argument("--horizon", type=int, default=12, help="")
        parser.add_argument("--lookback_window", type=int, default=2, help="")
        parser.add_argument("--dropout", type=int, default=0.3, help="")
        parser.add_argument("--lr", type=float, default=1e-3, help="")
        parser.add_argument("--weight_decay",
                            type=float,
                            default=1e-2,
                            help="")
        parser.add_argument("--in_features", type=int, default=1)
        parser.add_argument("--thid_features", type=int, default=5)
        parser.add_argument("--shid_features", type=int, default=5)
        parser.add_argument('--spatial', dest='spatial',
                            action='store_true')
        parser.add_argument('--temporal', dest='temporal',
                            action='store_true')
        parser.add_argument("--nhid", type=int, default=3)
        parser.add_argument("--kernel_size", type=int, default=3)
        parser.add_argument("--N", type=int, default=3)
        return parent_parser

Some Information about Encoder Some Information about Encoder # add dimension if adjs is a matrix # self.linear1 = nn.Linear(n_nodes, n_nodes) # self.linear2 = nn.Linear(seq_len, seq_len) # spatial correlation # # x = self.linear1(x + out) # temporal correlation # x = self.linear2(x)+res Some Information about Decoder # encoder # decoder