from time import sleep

import numpy as np
import pytest

import torch
from coreax.kernels import LinearKernel
from pixeltable import Error
from sklearn.decomposition import PCA

import pixeltable as pxt
from torch.utils.data import Dataset

from goldener.pxt_utils import GoldPxtTorchDataset
from goldener.reduce import GoldReducer
from goldener.select import GoldSelector
from goldener.select import (
    GoldGreedyClosestPointSelection,
    GoldGreedyKernelPoints,
)


class DummyDataset(Dataset):
    def __init__(self, samples):
        self._samples = samples

    def __len__(self):
        return len(self._samples)

    def __getitem__(self, idx):
        return (
            self._samples[idx].copy()
            if isinstance(self._samples[idx], dict)
            else self._samples[idx]
        )


class TestGoldSelector:
    def test_selection_table_creation_from_table(self):
        pxt.drop_dir("unit_test", force=True)

        src_path = "unit_test.src_table_input"
        desc_path = "unit_test.test_select_from_table"

        source_rows = [
            {
                "idx": 0,
                "vectorized": torch.zeros(1, 5).numpy(),
                "label": "dummy",
                "idx_vector": 0,
            },
            {
                "idx_vector": 1,
                "vectorized": torch.zeros(1, 5).numpy(),
                "label": "dummy",
                "idx": 0,
            },
        ]

        pxt.create_dir("unit_test", if_exists="ignore")
        src_table = pxt.create_table(
            src_path, source=source_rows, if_exists="replace_force"
        )

        selector = GoldSelector(table_path=desc_path, allow_existing=False)

        pxt_table = selector._selection_table_from_table(
            src_table, old_selection_table=None
        )

        assert set(pxt_table.columns()) == {
            selector.selection_key,
            "idx",
            "idx_vector",
            "chunked",
        }
        row_indices = [
            row["idx_vector"]
            for row in pxt_table.select(pxt_table.idx_vector).collect()
        ]
        assert set(row_indices) == {0, 1}

        pxt.drop_dir("unit_test", force=True)

    def test_selection_table_from_table_when_missing_vectorized(self):
        pxt.drop_dir("unit_test", force=True)

        src_path = "unit_test.src_table_invalid"
        pxt.create_dir("unit_test", if_exists="ignore")

        src_table = pxt.create_table(
            src_path,
            source=[{"idx": 0, "notvec": [1, 2, 3]}],
            if_exists="replace_force",
        )

        selector = GoldSelector(table_path="unit_test.test_select", allow_existing=True)

        with pytest.raises(ValueError, match="does not contain the required column"):
            selector._selection_table_from_table(
                select_from=src_table, old_selection_table=None
            )

        pxt.drop_dir("unit_test", force=True)

    def test_selection_table_from_table_when_invalid_old(self):
        pxt.drop_dir("unit_test", force=True)

        src_path = "unit_test.src_table_invalid"
        pxt.create_dir("unit_test", if_exists="ignore")

        src_table = pxt.create_table(
            src_path,
            source=[
                {
                    "idx": 0,
                }
            ],
            if_exists="replace_force",
        )

        selector = GoldSelector(table_path=src_path, allow_existing=True)

        with pytest.raises(ValueError, match="The table is missing required"):
            selector._selection_table_from_table(
                select_from=src_table, old_selection_table=src_table
            )

        pxt.drop_dir("unit_test", force=True)

    def test_selection_table_from_dataset(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_initialize"

        sample = {
            "vectorized": torch.rand(1, 5),
            "idx": 0,
        }
        dataset = DummyDataset([sample, sample])

        selector = GoldSelector(table_path=table_path, allow_existing=False)

        pxt_table = selector._selection_table_from_dataset(
            dataset, old_selection_table=None
        )

        assert set(pxt_table.columns()) == {
            selector.selection_key,
            selector.vectorized_key,
            "idx",
            "idx_vector",
            "chunked",
        }
        row_indices = [
            row["idx_vector"]
            for row in pxt_table.select(pxt_table.idx_vector).collect()
        ]
        assert set(row_indices) == {0, 1}

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_from_dataset(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_from_dataset"

        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path, allow_existing=True, batch_size=10, max_batches=None
        )

        selection_table = selector.select_in_table(
            dataset, select_size=10, value="train"
        )

        assert selection_table.count() == 100
        assert (
            selection_table.where(
                selection_table[selector.selection_key] == "train"
            ).count()
            == 10
        )

        # validate running with already filled work
        selector.select_in_table(dataset, select_size=10, value="train")

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_with_wrong_size(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_from_dataset"

        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path, allow_existing=True, batch_size=10, max_batches=None
        )

        with pytest.raises(
            ValueError, match="When select_size is a float, it must be in the range"
        ):
            selector.select_in_table(dataset, select_size=1.0, value="train")

        with pytest.raises(ValueError, match="select_size must be a positive integer"):
            selector.select_in_table(dataset, select_size=0, value="train")

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_from_dataset_with_ratio(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_from_dataset"

        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path, allow_existing=True, batch_size=10, max_batches=None
        )

        selection_table = selector.select_in_table(
            dataset, select_size=0.1, value="train"
        )

        assert selection_table.count() == 100
        assert (
            selection_table.where(
                selection_table[selector.selection_key] == "train"
            ).count()
            == 10
        )

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_from_dataset_with_small_ratio(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_from_dataset"

        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx_sample": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path, allow_existing=True, batch_size=10, max_batches=None
        )

        selection_table = selector.select_in_table(
            dataset, select_size=0.0001, value="train"
        )

        assert selection_table.count() == 100
        assert (
            selection_table.where(
                selection_table[selector.selection_key] == "train"
            ).count()
            == 1
        )

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_from_dataset_with_class(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_from_dataset"

        dataset = DummyDataset(
            [
                {"vectorized": torch.rand(5), "idx": idx, "label": str(idx % 2)}
                for idx in range(100)
            ]
        )

        selector = GoldSelector(
            table_path=table_path,
            allow_existing=True,
            class_key="label",
            batch_size=10,
            max_batches=None,
        )

        selection_table = selector.select_in_table(
            dataset,
            select_size=10,
            value="train",
        )

        assert selection_table.count() == 100
        assert (
            len(
                selector.get_selected_sample_indices(
                    selection_table,
                    "train",
                    selector.selection_key,
                    class_key=selector.class_key,
                    class_value="0",
                )
            )
            == 5
        )
        assert (
            len(
                selector.get_selected_sample_indices(
                    selection_table,
                    "train",
                    selector.selection_key,
                    class_key=selector.class_key,
                    class_value="1",
                )
            )
            == 5
        )

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_with_chunk(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_chunk"

        dataset = DummyDataset(
            [
                {
                    "vectorized": torch.rand(
                        5,
                    ),
                    "idx": idx,
                }
                for idx in range(100)
            ]
        )

        selector = GoldSelector(
            table_path=table_path, allow_existing=True, chunk=26, batch_size=10
        )

        selection_table = selector.select_in_table(
            dataset, select_size=10, value="train"
        )

        assert selection_table.count() == 100
        assert (
            selection_table.where(
                selection_table[selector.selection_key] == "train"
            ).count()
            == 10
        )

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_with_reducer(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_reducer"

        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path,
            allow_existing=True,
            reducer=GoldReducer(PCA(n_components=3)),
            batch_size=10,
        )

        selection_table = selector.select_in_table(
            dataset, select_size=10, value="train"
        )

        assert selection_table.count() == 100
        assert (
            selection_table.where(
                selection_table[selector.selection_key] == "train"
            ).count()
            == 10
        )

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_with_max_batches(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_max_batches"

        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path, allow_existing=True, batch_size=10, max_batches=2
        )

        selection_table = selector.select_in_table(
            dataset, select_size=10, value="train"
        )

        assert selection_table.count() == 20
        assert (
            selection_table.where(
                selection_table[selector.selection_key] == "train"
            ).count()
            == 10
        )

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_with_restart(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_max_batches"

        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path, allow_existing=True, batch_size=10, max_batches=2
        )

        selector.select_in_table(dataset, select_size=10, value="train")

        selector.max_batches = None
        selection_table = selector.select_in_table(
            dataset, select_size=20, value="train"
        )

        assert selection_table.count() == 100
        assert (
            selection_table.where(
                selection_table[selector.selection_key] == "train"
            ).count()
            == 20
        )

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_with_restart_disallowed(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_max_batches"

        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path, allow_existing=False, batch_size=10, max_batches=2
        )

        selector.select_in_table(dataset, select_size=10, value="train")

        selector.max_batches = None

        # calling select_in_table when a table exists and allow_existing is False should raise
        with pytest.raises(
            ValueError, match="already exists and allow_existing is set to"
        ):
            selector.select_in_table(dataset, select_size=20, value="train")

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_with_not_enough_sample(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_not_enough"
        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path, allow_existing=False, batch_size=10, max_batches=2
        )

        # calling select_in_table when a table exists and allow_existing is False should raise
        with pytest.raises(
            ValueError, match="Cannot select more unique data points than available"
        ):
            selector.select_in_table(dataset, select_size=21, value="train")

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_table_from_table(self):
        pxt.drop_dir("unit_test", force=True)

        src_path = "unit_test.test_select_in_table"

        pxt.create_dir("unit_test", if_exists="ignore")
        src_table = pxt.create_table(
            src_path,
            source=[
                {
                    "vectorized": torch.rand(5).numpy().astype(np.float32),
                    "idx": idx % 10,
                    "idx_vector": idx,
                }
                for idx in range(100)
            ],
            if_exists="replace_force",
            primary_key="idx_vector",
        )

        selector = GoldSelector(
            table_path="unit_test.test_select",
            allow_existing=True,
            batch_size=10,
            max_batches=2,
        )

        selector.select_in_table(src_table, select_size=3, value="train")

        selector.max_batches = None
        selection_table = selector.select_in_table(
            src_table, select_size=6, value="train"
        )

        assert selection_table.count() == 100
        assert (
            len(
                selector.get_selected_sample_indices(
                    selection_table, "train", selector.selection_key
                )
            )
            == 6
        )

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_dataset_from_dataset(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_from_dataset"

        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path, allow_existing=False, batch_size=10, max_batches=2
        )

        dataset = selector.select_in_dataset(dataset, select_count=3, value="train")

        assert isinstance(dataset, GoldPxtTorchDataset)

        selected_count = 0
        total_count = 0
        already_selected = set()
        for item in dataset:
            total_count += 1
            if item["selected"] == "train" and item["idx"] not in already_selected:
                selected_count += 1
                already_selected.add(item["idx"])

        assert total_count == 20
        assert selected_count == 3

        dataset.keep_cache = False

        pxt.drop_dir("unit_test", force=True)

    def test_select_in_dataset_with_drop_table(self):
        pxt.drop_dir("unit_test", force=True)

        table_path = "unit_test.test_select_from_dataset"

        dataset = DummyDataset(
            [{"vectorized": torch.rand(5), "idx": idx} for idx in range(100)]
        )

        selector = GoldSelector(
            table_path=table_path,
            allow_existing=False,
            batch_size=10,
            max_batches=2,
            drop_table=True,
        )

        dataset = selector.select_in_dataset(dataset, select_count=3, value="train")

        sleep(1)
        with pytest.raises(
            Error, match="Path 'unit_test.test_select_from_dataset' does not exist"
        ):
            pxt.get_table(table_path, if_not_exists="error")

        dataset.keep_cache = False

        pxt.drop_dir("unit_test", force=True)

    def test_get_selected_sample_indices(self):
        pxt.drop_dir("unit_test", force=True)

        src_path = "unit_test.test_select"

        pxt.create_dir("unit_test", if_exists="ignore")
        src_table = pxt.create_table(
            src_path,
            source=[
                {
                    "vectorized": torch.rand(5).numpy().astype(np.float32),
                    "idx": idx,
                    "selected": "train" if idx < 50 else None,
                    "label": "value" if idx < 25 else "other",
                }
                for idx in range(100)
            ],
            if_exists="replace_force",
        )

        sample_indices = GoldSelector.get_selected_sample_indices(
            table=src_table,
            value="train",
            selection_key="selected",
        )

        assert sample_indices == set(range(50))

        sample_indices = GoldSelector.get_selected_sample_indices(
            table=src_table,
            value=None,
            selection_key="selected",
        )
        assert sample_indices == set(range(50, 100, 1))

        sample_indices = GoldSelector.get_selected_sample_indices(
            table=src_table,
            value="train",
            selection_key="selected",
            class_key="label",
            class_value="value",
        )
        assert sample_indices == set(range(25))

        sample_indices = GoldSelector.get_selected_sample_indices(
            table=src_table,
            value="train",
            selection_key="selected",
            class_key="label",
            class_value="other",
        )
        assert sample_indices == set(range(25, 50, 1))

        with pytest.raises(
            ValueError, match="class_key and class_value must be set together"
        ):
            GoldSelector.get_selected_sample_indices(
                table=src_table,
                value="train",
                selection_key="selected",
                class_key="label",
                class_value=None,
            )

        with pytest.raises(
            ValueError, match="class_key and class_value must be set together"
        ):
            GoldSelector.get_selected_sample_indices(
                table=src_table,
                value="train",
                selection_key="selected",
                class_key=None,
                class_value="other",
            )

        pxt.drop_dir("unit_test", force=True)


class TestGoldGreedyClosestPointSelection:
    def test_basic_selection_properties(self) -> None:
        x = torch.tensor([[0.0], [1.0], [10.0], [10.1], [1.1]], dtype=torch.float32)

        tool = GoldGreedyClosestPointSelection(device="cpu")
        indices = tool.select(x, k=2)

        assert indices == [1, 2]

    def test_select_all_points(self) -> None:
        x = torch.tensor(
            [[0.0], [1.0], [2.0], [3.0]],
            dtype=torch.float32,
        )

        tool = GoldGreedyClosestPointSelection(device="cpu")
        indices = tool.select(x, k=x.size(0))

        assert indices == [0, 1, 2, 3]

    def test_selection_with_multidimensional_features(self) -> None:
        x = torch.tensor(
            [[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]],
            dtype=torch.float32,
        )

        tool = GoldGreedyClosestPointSelection(device="cpu")
        indices = tool.select(x, k=2)

        assert indices == [0, 1]

    def test_with_k_greater_than_size(self) -> None:
        x = torch.tensor([[0.0], [1.0]], dtype=torch.float32)

        tool = GoldGreedyClosestPointSelection(device="cpu")
        with pytest.raises(
            ValueError, match="k cannot be greater than the number of data points in x"
        ):
            tool.select(x, k=5)


class TestGoldGreedyKernelPoints:
    def test_basic_selection_with_linear_kernel(self) -> None:
        tool = GoldGreedyKernelPoints(
            feature_kernel=LinearKernel(output_scale=1, constant=0)
        )

        x = torch.arange(12, dtype=torch.float32).view(4, 3)
        k = 2

        indices = tool.select(x, k)

        assert len(indices) == k
        assert len(set(indices)) == k

    def test_select_all_points_with_linear_kernel(self) -> None:
        tool = GoldGreedyKernelPoints(
            feature_kernel=LinearKernel(output_scale=1, constant=0)
        )

        x = torch.arange(12, dtype=torch.float32).view(4, 3)
        k = 4

        indices = tool.select(x, k)

        assert len(indices) == k
        assert len(set(indices)) == k

    def test_with_k_greater_than_size(self) -> None:
        x = torch.tensor([[0.0], [1.0]], dtype=torch.float32)

        tool = GoldGreedyKernelPoints(
            feature_kernel=LinearKernel(output_scale=1, constant=0)
        )
        with pytest.raises(ValueError, match="must be less than"):
            tool.select(x, k=5)
