import string

import numpy as np
import pytest
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader

from pytorch_widedeep.models import Wide, TabMlp
from pytorch_widedeep.metrics import Accuracy, MultipleMetrics
from pytorch_widedeep.training._finetune import FineTune
from pytorch_widedeep.utils.general_utils import setup_device, to_device_model
from pytorch_widedeep.models.image._layers import conv_layer
from pytorch_widedeep.training._wd_dataset import WideDeepDataset

device = setup_device()


# Define a series of simple models to quickly test the FineTune class
class TextModeTestClass(nn.Module):
    def __init__(self):
        super(TextModeTestClass, self).__init__()
        self.word_embed = nn.Embedding(5, 16, padding_idx=0)
        self.rnn = nn.LSTM(16, 8, batch_first=True)
        self.linear = nn.Linear(8, 1)

    def forward(self, X):
        embed = self.word_embed(X.long())
        o, (h, c) = self.rnn(embed)
        return self.linear(h).view(-1, 1)


class ImageModeTestClass(nn.Module):
    def __init__(self):
        super(ImageModeTestClass, self).__init__()

        self.conv_block = nn.Sequential(
            conv_layer(3, 64, 3),
            conv_layer(64, 128, 1, maxpool=False, adaptiveavgpool=True),
        )
        self.linear = nn.Linear(128, 1)

    def forward(self, X):
        x = self.conv_block(X)
        x = x.view(x.size(0), -1)
        return self.linear(x)


# Remember that the FineTune class will be instantiated inside WideDeep and
# will take, among others, the activation_fn and the loss_fn of that class as
# parameters. Therefore, we define equivalent classes to replicate the
# scenario
# def activ_fn(inp):
#     return torch.sigmoid(inp)


def loss_fn(y_pred, y_true):
    return F.binary_cross_entropy_with_logits(y_pred, y_true.view(-1, 1))


# Define the data components:

# target
target = np.random.randint(0, 2, (100, 1)).astype(float)

# wide
X_wide = np.random.randint(1, 20, (100, 4)).astype(float)

# deep
colnames = list(string.ascii_lowercase)[:10]
embed_cols = [np.random.choice(np.arange(5), 100) for _ in range(5)]
cont_cols = [np.random.rand(100) for _ in range(5)]
embed_input = [(u, i, j) for u, i, j in zip(colnames[:5], [5] * 5, [16] * 5)]
column_idx = {k: v for v, k in enumerate(colnames[:10])}
continuous_cols = colnames[-5:]
X_tab = np.vstack(embed_cols + cont_cols).transpose()

# text
X_text = np.concatenate(
    (np.zeros([100, 1]), np.random.randint(1, 4, (100, 4)).astype(float)), axis=1
)

# image
X_image = np.random.rand(100, 28, 28, 3)

# Define the model components

# wide/linear
wide = Wide(np.unique(X_wide).size, 1)
wide = to_device_model(wide, device)

# tabular
tab_mlp = TabMlp(
    column_idx=column_idx,
    cat_embed_input=embed_input,
    continuous_cols=continuous_cols,
    mlp_hidden_dims=[32, 16, 8],
    mlp_dropout=0.2,
)
tab_mlp = nn.Sequential(tab_mlp, nn.Linear(8, 1))  # type: ignore[assignment]
tab_mlp = to_device_model(tab_mlp, device)

# text
deeptext = TextModeTestClass()
deeptext = to_device_model(deeptext, device)

# image
deepimage = ImageModeTestClass()
ddeimage = to_device_model(deepimage, device)

# Define the loader
mmset = WideDeepDataset(X_wide, X_tab, X_text, X_image, target)
mmloader = DataLoader(mmset, batch_size=10, shuffle=True)

# Instantiate the FineTune class
finetuner = FineTune(loss_fn, MultipleMetrics([Accuracy()]), "binary", False)

# List the layers for the finetune_gradual method
# tab_mlp childrens -> TabMmlp and the final Linear layer
# TabMlp children -> Embeddings and MLP
# MLP children -> dense layers
# so here we go...
last_linear = list(tab_mlp.children())[1]
inverted_mlp_layers = [
    list(list(list(tab_mlp.named_modules())[12][1].children())[2].children())[0]
]
tab_layers = [last_linear] + inverted_mlp_layers
text_layers = [c for c in list(deeptext.children())[1:]][::-1]
image_layers = [c for c in list(deepimage.children())][::-1]


###############################################################################
# Simply test that finetune_all runs
###############################################################################
@pytest.mark.parametrize(
    "model, modelname, loader, n_epochs, max_lr",
    [
        (wide, "wide", mmloader, 1, 0.01),
        (tab_mlp, "deeptabular", mmloader, 1, 0.01),
        (deeptext, "deeptext", mmloader, 1, 0.01),
        (deepimage, "deepimage", mmloader, 1, 0.01),
    ],
)
def test_finetune_all(model, modelname, loader, n_epochs, max_lr):
    has_run = True
    try:
        finetuner.finetune_all(model, modelname, loader, n_epochs, max_lr)
    except Exception:
        has_run = False
    assert has_run


###############################################################################
# Simply test that finetune_gradual runs
###############################################################################
@pytest.mark.parametrize(
    "model, modelname, loader, max_lr, layers, routine",
    [
        (tab_mlp, "deeptabular", mmloader, 0.01, tab_layers, "felbo"),
        (tab_mlp, "deeptabular", mmloader, 0.01, tab_layers, "howard"),
        (deeptext, "deeptext", mmloader, 0.01, text_layers, "felbo"),
        (deeptext, "deeptext", mmloader, 0.01, text_layers, "howard"),
        (deepimage, "deepimage", mmloader, 0.01, image_layers, "felbo"),
        (deepimage, "deepimage", mmloader, 0.01, image_layers, "howard"),
    ],
)
def test_finetune_gradual(model, modelname, loader, max_lr, layers, routine):
    has_run = True
    try:
        finetuner.finetune_gradual(model, modelname, loader, max_lr, layers, routine)
    except Exception as e:
        has_run = False
        print(e)
    assert has_run
