from typing import Tuple
import numpy as np
import torch
from torch.utils.data import TensorDataset, DataLoader


class TorchModel(torch.nn.Module):
    def predict(self, x: np.ndarray, batch_size: int = 1) -> Tuple[np.ndarray, np.ndarray]:
        x = torch.from_numpy(x).cuda().float()
        dataset = TensorDataset(x)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
        self.eval()
        decodeds = []
        risks = []
        for x in dataloader:
            decoded, risk_pred = self(x[0])
            decodeds.append(decoded)
            risks.append(risk_pred)
        decodeds = torch.cat(decodeds, dim=0)
        decodeds = decodeds.cpu().detach().numpy()
        risks = torch.cat(risks, dim=0)
        risks = risks.cpu().detach().numpy()
        return decodeds, risks


class BaselineAutoencoder(TorchModel):
    def __init__(self, input_shape: int, dropout: float, num_ae_units1: int, num_ae_units2: int):
        super().__init__()
        num_ae_units1 = round(num_ae_units1)
        num_ae_units2 = round(num_ae_units2)
        self.encoder = torch.nn.Sequential(
            torch.nn.Dropout(dropout),
            torch.nn.Linear(input_shape, num_ae_units1),
            torch.nn.ReLU(),
            torch.nn.Linear(num_ae_units1, num_ae_units2),
            torch.nn.ReLU(),
        )

        self.decoder = torch.nn.Sequential(
            torch.nn.Linear(num_ae_units2, num_ae_units1),
            torch.nn.ReLU(),
            torch.nn.Linear(num_ae_units1, input_shape)
        )
        self.risk_regressor = torch.nn.Linear(num_ae_units2, 1)

    def forward(self, x):
        encoded = self.encoder(x)
        risk_pred = self.risk_regressor(encoded)
        decoded = self.decoder(encoded)
        return decoded, risk_pred


class BaselineBNAutoencoder(TorchModel):
    def __init__(self, input_shape: int, dropout: float, num_ae_units1: int, num_ae_units2: int):
        super().__init__()
        num_ae_units1 = round(num_ae_units1)
        num_ae_units2 = round(num_ae_units2)
        self.encoder = torch.nn.Sequential(
            torch.nn.Dropout(dropout),
            torch.nn.Linear(input_shape, num_ae_units1),
            torch.nn.BatchNorm1d(num_ae_units1),
            torch.nn.ReLU(),
            torch.nn.Linear(num_ae_units1, num_ae_units2),
            torch.nn.BatchNorm1d(num_ae_units2),
            torch.nn.ReLU(),
        )

        self.decoder = torch.nn.Sequential(
            torch.nn.Linear(num_ae_units2, num_ae_units1),
            torch.nn.BatchNorm1d(num_ae_units1),
            torch.nn.ReLU(),
            torch.nn.Linear(num_ae_units1, input_shape)
        )
        self.risk_regressor = torch.nn.Sequential(
            torch.nn.Linear(num_ae_units2, num_ae_units2//2),
            torch.nn.BatchNorm1d(num_ae_units2//2),
            torch.nn.ReLU(),
            torch.nn.Linear(num_ae_units2//2, 1)
        )

    def forward(self, x):
        encoded = self.encoder(x)
        risk_pred = self.risk_regressor(encoded)
        decoded = self.decoder(encoded)
        return decoded, risk_pred


def model_factory(model_name: str, **kwargs) -> TorchModel:
    # Before defining network architecture, clear current computation graph (if one exists)
    torch.cuda.empty_cache()
    if model_name == "baseline_autoencoder":
        model = BaselineAutoencoder(**kwargs)
    elif model_name == "baseline_bn_autoencoder":
        model = BaselineBNAutoencoder(**kwargs)
    else:
        raise ValueError("Model name {} has not been implemented.".format(model_name))
    return model