# Copyright The PyTorch Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re
from unittest import mock

import pandas as pd
import pytest
import torch
from pytorch_lightning import Trainer

from flash.__main__ import main
from flash.core.data.io.input import DataKeys
from flash.core.utilities.imports import _TABULAR_AVAILABLE
from flash.tabular.classification.data import TabularClassificationData
from flash.tabular.classification.model import TabularClassifier
from tests.helpers.utils import _SERVE_TESTING, _TABULAR_TESTING

# ======== Mock functions ========


class DummyDataset(torch.utils.data.Dataset):
    def __init__(self, num_num=16, num_cat=16):
        super().__init__()
        self.num_num = num_num
        self.num_cat = num_cat

    def __getitem__(self, index):
        target = torch.randint(0, 10, size=(1,)).item()
        cat_vars = torch.randint(0, 10, size=(self.num_cat,))
        num_vars = torch.rand(self.num_num)
        return {DataKeys.INPUT: (cat_vars, num_vars), DataKeys.TARGET: target}

    def __len__(self) -> int:
        return 100


# ==============================


@pytest.mark.skipif(not _TABULAR_TESTING, reason="tabular libraries aren't installed.")
def test_init_train(tmpdir):
    train_dl = torch.utils.data.DataLoader(DummyDataset(), batch_size=16)
    model = TabularClassifier(num_classes=10, num_features=16 + 16, embedding_sizes=16 * [(10, 32)])
    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
    trainer.fit(model, train_dl)


@pytest.mark.skipif(not _TABULAR_TESTING, reason="tabular libraries aren't installed.")
def test_init_train_no_num(tmpdir):
    train_dl = torch.utils.data.DataLoader(DummyDataset(num_num=0), batch_size=16)
    model = TabularClassifier(num_classes=10, num_features=16, embedding_sizes=16 * [(10, 32)])
    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
    trainer.fit(model, train_dl)


@pytest.mark.skipif(not _TABULAR_TESTING, reason="tabular libraries aren't installed.")
def test_init_train_no_cat(tmpdir):
    train_dl = torch.utils.data.DataLoader(DummyDataset(num_cat=0), batch_size=16)
    model = TabularClassifier(num_classes=10, num_features=16, embedding_sizes=[])
    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
    trainer.fit(model, train_dl)


@pytest.mark.skipif(_TABULAR_AVAILABLE, reason="tabular libraries are installed.")
def test_module_import_error(tmpdir):
    with pytest.raises(ModuleNotFoundError, match="[tabular]"):
        TabularClassifier(num_classes=10, num_features=16, embedding_sizes=[])


@pytest.mark.skipif(not _TABULAR_TESTING, reason="tabular libraries aren't installed.")
def test_jit(tmpdir):
    model = TabularClassifier(num_classes=10, num_features=8, embedding_sizes=4 * [(10, 32)])
    model.eval()

    # torch.jit.script doesn't work with tabnet
    model = torch.jit.trace(model, ((torch.randint(0, 10, size=(1, 4)), torch.rand(1, 4)),))

    # TODO: torch.jit.save doesn't work with tabnet
    # path = os.path.join(tmpdir, "test.pt")
    # torch.jit.save(model, path)
    # model = torch.jit.load(path)

    out = model((torch.randint(0, 10, size=(1, 4)), torch.rand(1, 4)))
    assert isinstance(out, torch.Tensor)
    assert out.shape == torch.Size([1, 10])


@pytest.mark.skipif(not _SERVE_TESTING, reason="serve libraries aren't installed.")
@mock.patch("flash._IS_TESTING", True)
def test_serve():
    train_data = {"num_col": [1.4, 2.5], "cat_col": ["positive", "negative"], "target": [1, 2]}
    datamodule = TabularClassificationData.from_data_frame(
        "cat_col",
        "num_col",
        "target",
        train_data_frame=pd.DataFrame.from_dict(train_data),
    )
    model = TabularClassifier.from_data(datamodule)
    # TODO: Currently only servable once a input_transform has been attached
    model._input_transform = datamodule.input_transform
    model._input_transform._state = datamodule.train_dataset._state
    model.eval()
    model.serve()


@pytest.mark.skipif(_TABULAR_AVAILABLE, reason="tabular libraries are installed.")
def test_load_from_checkpoint_dependency_error():
    with pytest.raises(ModuleNotFoundError, match=re.escape("'lightning-flash[tabular]'")):
        TabularClassifier.load_from_checkpoint("not_a_real_checkpoint.pt")


@pytest.mark.skipif(not _TABULAR_TESTING, reason="tabular libraries aren't installed.")
def test_cli():
    cli_args = ["flash", "tabular_classification", "--trainer.fast_dev_run", "True"]
    with mock.patch("sys.argv", cli_args):
        try:
            main()
        except SystemExit:
            pass
