import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl


class MLP(pl.LightningModule):
    def __init__(self, input_size, xcol='emb', ycol='avg_rating', batch_norm=True):
        super().__init__()
        self.input_size = input_size
        self.xcol = xcol
        self.ycol = ycol
        # self.layers = nn.Sequential(
        #     nn.Linear(self.input_size, 2048),
        #     nn.ReLU(),
        #     nn.BatchNorm1d(2048),
        #     nn.Dropout(0.4),

        #     nn.Linear(2048, 512),
        #     nn.ReLU(),
        #     nn.BatchNorm1d(512),
        #     nn.Dropout(0.3),

        #     nn.Linear(512, 256),
        #     nn.ReLU(),
        #     nn.BatchNorm1d(256),
        #     nn.Dropout(0.2),

        #     nn.Linear(256, 128),
        #     nn.ReLU(),
        #     nn.BatchNorm1d(128),
        #     nn.Dropout(0.1),

        #     nn.Linear(128, 32),
        #     nn.ReLU(),
        #     nn.Linear(32, 1)
        # )
        self.layers = nn.Sequential(
            nn.Linear(self.input_size, 1024),
            # nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(1024, 128),
            # nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            # nn.ReLU(),
            nn.Dropout(0.1),

            nn.Linear(64, 16),
            # nn.ReLU(),

            nn.Linear(16, 1)
        )

    def forward(self, x):
        return self.layers(x)

    def training_step(self, batch, batch_idx):
        x = batch[self.xcol]
        y = batch[self.ycol].reshape(-1, 1)
        x_hat = self.layers(x)
        loss = F.mse_loss(x_hat, y)
        return loss

    def validation_step(self, batch, batch_idx):
        x = batch[self.xcol]
        y = batch[self.ycol].reshape(-1, 1)
        x_hat = self.layers(x)
        loss = F.mse_loss(x_hat, y)
        return loss

    # def configure_optimizers(self):
    #     optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
    #     return optimizer


class ResidualBlock(nn.Module):
    def __init__(self, input_size, output_size, batch_norm=True, dropout_rate=0.0):
        super(ResidualBlock, self).__init__()
        self.linear = nn.Linear(input_size, output_size)
        self.relu = nn.ReLU()
        self.batch_norm = nn.BatchNorm1d(output_size) if batch_norm else nn.Identity()
        self.dropout = nn.Dropout(dropout_rate)
        self.adjust_dims = nn.Linear(input_size, output_size) if input_size != output_size else nn.Identity()

    def forward(self, x):
        identity = self.adjust_dims(x)
        out = self.linear(x)
        out = self.relu(out)
        out = self.batch_norm(out)
        out = self.dropout(out)
        out += identity
        out = self.relu(out)
        return out


class ResMLP(pl.LightningModule):
    def __init__(self, input_size, xcol='emb', ycol='avg_rating', batch_norm=True):
        super().__init__()
        self.input_size = input_size
        self.xcol = xcol
        self.ycol = ycol
        self.layers = nn.Sequential(
            ResidualBlock(input_size, 2048, batch_norm, dropout_rate=0.3),
            ResidualBlock(2048, 512, batch_norm, dropout_rate=0.3),
            ResidualBlock(512, 256, batch_norm, dropout_rate=0.2),
            ResidualBlock(256, 128, batch_norm, dropout_rate=0.1),
            nn.Linear(128, 32),
            nn.ReLU(),
            nn.Linear(32, 1)
        )

    def forward(self, x):
        return self.layers(x)

    def training_step(self, batch, batch_idx):
        x = batch[self.xcol]
        y = batch[self.ycol].reshape(-1, 1)
        x_hat = self.layers(x)
        loss = F.mse_loss(x_hat, y)
        return loss

    def validation_step(self, batch, batch_idx):
        x = batch[self.xcol]
        y = batch[self.ycol].reshape(-1, 1)
        x_hat = self.layers(x)
        loss = F.mse_loss(x_hat, y)
        return loss