import os
from pathlib import Path

import numpy as np
import pytest
from jax import numpy as jnp

from slzero import dataset, hmm
from slzero.hmm.functions import Params, predict


@pytest.fixture
def dataset_path() -> Path:
    path = Path(os.path.dirname(__file__)) / "data" / "test.conllu"
    return path


@pytest.fixture
def dataset_and_vocab(
    dataset_path: Path,
) -> tuple[list, list, dataset.Vocab, dataset.Vocab]:
    X_train, y_train = dataset.read_file(str(dataset_path))
    vocab_token = dataset.build_vocab(X_train, unk_token="unk", max_size=None)
    vocab_label = dataset.build_vocab(y_train, unk_token=None)
    return (X_train, y_train, vocab_token, vocab_label)


class TestHiddenMarkovModel:
    def test_smooth(self, dataset_and_vocab: tuple) -> None:
        smooth_factor = 0.2
        (
            X_train,
            y_train,
            vocab_token,
            vocab_label,
        ) = dataset_and_vocab

        num_tokens = len(vocab_token)
        num_labels = len(vocab_label)

        pi, a, b = hmm.train(
            hmm.init_params(num_labels, num_tokens),
            dataset.get_batch_stream(vocab_token, vocab_label, X_train, y_train),
            smooth_factor=smooth_factor,
        )

        label_to_index = vocab_label._token_to_index
        token_to_index = vocab_token._token_to_index

        # Definition
        # https://en.wikipedia.org/wiki/Additive_smoothing
        def smoothed_div(x: int, N: int, d: int) -> float:
            return (x + smooth_factor) / (N + (smooth_factor * d))

        assert np.allclose(
            a[label_to_index["AUX"], label_to_index["PUNCT"]],
            smoothed_div(2, 8, num_labels),
        )
        assert np.allclose(
            b[label_to_index["NOUN"], token_to_index["ソング"]],
            smoothed_div(1, 9, num_tokens),
        )
        assert np.allclose(pi[label_to_index["CCONJ"]], smoothed_div(1, 2, num_labels))

    def test_train(self, dataset_and_vocab: tuple) -> None:
        (
            X_train,
            y_train,
            vocab_token,
            vocab_label,
        ) = dataset_and_vocab

        num_tokens = len(vocab_token)
        num_labels = len(vocab_label)

        pi, a, b = hmm.train(
            hmm.init_params(num_labels, num_tokens),
            dataset.get_batch_stream(vocab_token, vocab_label, X_train, y_train),
            smooth_factor=0.0,
        )

        label_to_index = vocab_label._token_to_index
        token_to_index = vocab_token._token_to_index

        assert num_labels == 11
        assert num_tokens == (33 + 1)  # +1 is unk
        assert a.shape == (num_labels, num_labels)
        for transition in a:
            assert np.allclose(transition.sum(), 1.0)
        assert np.allclose(a[label_to_index["VERB"], label_to_index["AUX"]], 2 / 3)
        assert a[label_to_index["AUX"], label_to_index["PUNCT"]], 2 / 8
        assert a[label_to_index["AUX"], label_to_index["ADJ"]] == 0.0
        assert b.shape == (num_labels, num_tokens)
        assert b[:, 0].sum() == 0.0  # unk token
        for emission in b:
            assert np.allclose(emission.sum(), 1.0)
        assert b[label_to_index["AUX"], token_to_index["た"]] == 4 / 8
        assert np.allclose(b[label_to_index["SCONJ"], token_to_index["けど"]], 1 / 3)
        assert b[label_to_index["PUNCT"], token_to_index["ソング"]] == 0.0
        assert len(pi) == num_labels
        assert pi.sum() == 1.0
        assert pi[label_to_index["CCONJ"]] == 1 / 2
        assert pi[label_to_index["ADP"]] == 0.0

    def test_predict(self) -> None:
        pi = [0.3, 0.3, 0.3]
        a = [[0.1, 0.8, 0.1], [0.1, 0.1, 0.8], [0.8, 0.1, 0.1]]
        b = [[0.1, 0.7, 0.1, 0.1], [0.1, 0.1, 0.7, 0.1], [0.1, 0.1, 0.1, 0.7]]

        params = Params(*map(jnp.array, (pi, a, b)))

        X_batch = jnp.array([[1, 2, 3]])
        masks = jnp.full(X_batch.shape, True)
        last_probabilities, predictions = predict(params, X_batch, masks)
        assert jnp.allclose(predictions[0], jnp.array([0, 1, 2]))
        assert jnp.allclose(
            last_probabilities[0], jnp.array([0.001176, 0.001176, 0.065856])
        )

        X_batch = jnp.array([[1, 2, 3, 0]])
        masks = X_batch > 0
        last_probabilities, predictions = predict(params, X_batch, masks)
        assert jnp.allclose(predictions[0], jnp.array([0, 1, 2, -1]))
        assert jnp.allclose(
            last_probabilities[0], jnp.array([0.001176, 0.001176, 0.065856])
        )

        X_batch = jnp.array([[1, 2, 0]])
        masks = jnp.full(X_batch.shape, True)
        last_probabilities, predictions = predict(params, X_batch, masks)
        assert jnp.allclose(predictions[0], jnp.array([0, 1, 2]))
        assert jnp.allclose(
            last_probabilities[0], jnp.array([0.001176, 0.001176, 0.009408])
        )

import os
from pathlib import Path

import numpy as np
import pytest
from jax import numpy as jnp

from slzero import dataset, hmm
from slzero.hmm.functions import Params, predict


@pytest.fixture
def dataset_path() -> Path:
    path = Path(os.path.dirname(__file__)) / "data" / "test.conllu"
    return path


@pytest.fixture
def dataset_and_vocab(
    dataset_path: Path,
) -> tuple[list, list, dataset.Vocab, dataset.Vocab]:
    X_train, y_train = dataset.read_file(str(dataset_path))
    vocab_token = dataset.build_vocab(X_train, unk_token="unk", max_size=None)
    vocab_label = dataset.build_vocab(y_train, unk_token=None)
    return (X_train, y_train, vocab_token, vocab_label)


class TestHiddenMarkovModel:
    def test_smooth(self, dataset_and_vocab: tuple) -> None:
        smooth_factor = 0.2
        (
            X_train,
            y_train,
            vocab_token,
            vocab_label,
        ) = dataset_and_vocab

        num_tokens = len(vocab_token)
        num_labels = len(vocab_label)

        pi, a, b = hmm.train(
            hmm.init_params(num_labels, num_tokens),
            dataset.get_batch_stream(vocab_token, vocab_label, X_train, y_train),
            smooth_factor=smooth_factor,
        )

        label_to_index = vocab_label._token_to_index
        token_to_index = vocab_token._token_to_index

        # Definition
        # https://en.wikipedia.org/wiki/Additive_smoothing
        def smoothed_div(x: int, N: int, d: int) -> float:
            return (x + smooth_factor) / (N + (smooth_factor * d))

        assert np.allclose(
            a[label_to_index["AUX"], label_to_index["PUNCT"]],
            smoothed_div(2, 8, num_labels),
        )
        assert np.allclose(
            b[label_to_index["NOUN"], token_to_index["ソング"]],
            smoothed_div(1, 9, num_tokens),
        )
        assert np.allclose(pi[label_to_index["CCONJ"]], smoothed_div(1, 2, num_labels))

    def test_train(self, dataset_and_vocab: tuple) -> None:
        (
            X_train,
            y_train,
            vocab_token,
            vocab_label,
        ) = dataset_and_vocab

        num_tokens = len(vocab_token)
        num_labels = len(vocab_label)

        pi, a, b = hmm.train(
            hmm.init_params(num_labels, num_tokens),
            dataset.get_batch_stream(vocab_token, vocab_label, X_train, y_train),
            smooth_factor=0.0,
        )

        label_to_index = vocab_label._token_to_index
        token_to_index = vocab_token._token_to_index

        assert num_labels == 11
        assert num_tokens == (33 + 1)  # +1 is unk
        assert a.shape == (num_labels, num_labels)
        for transition in a:
            assert np.allclose(transition.sum(), 1.0)
        assert np.allclose(a[label_to_index["VERB"], label_to_index["AUX"]], 2 / 3)
        assert a[label_to_index["AUX"], label_to_index["PUNCT"]], 2 / 8
        assert a[label_to_index["AUX"], label_to_index["ADJ"]] == 0.0
        assert b.shape == (num_labels, num_tokens)
        assert b[:, 0].sum() == 0.0  # unk token
        for emission in b:
            assert np.allclose(emission.sum(), 1.0)
        assert b[label_to_index["AUX"], token_to_index["た"]] == 4 / 8
        assert np.allclose(b[label_to_index["SCONJ"], token_to_index["けど"]], 1 / 3)
        assert b[label_to_index["PUNCT"], token_to_index["ソング"]] == 0.0
        assert len(pi) == num_labels
        assert pi.sum() == 1.0
        assert pi[label_to_index["CCONJ"]] == 1 / 2
        assert pi[label_to_index["ADP"]] == 0.0

    def test_predict(self) -> None:
        pi = [0.3, 0.3, 0.3]
        a = [[0.1, 0.8, 0.1], [0.1, 0.1, 0.8], [0.8, 0.1, 0.1]]
        b = [[0.1, 0.7, 0.1, 0.1], [0.1, 0.1, 0.7, 0.1], [0.1, 0.1, 0.1, 0.7]]

        params = Params(*map(jnp.array, (pi, a, b)))

        X_batch = jnp.array([[1, 2, 3]])
        masks = jnp.full(X_batch.shape, True)
        last_probabilities, predictions = predict(params, X_batch, masks)
        assert jnp.allclose(predictions[0], jnp.array([0, 1, 2]))
        assert jnp.allclose(
            last_probabilities[0], jnp.array([0.001176, 0.001176, 0.065856])
        )

        X_batch = jnp.array([[1, 2, 3, 0]])
        masks = X_batch > 0
        last_probabilities, predictions = predict(params, X_batch, masks)
        assert jnp.allclose(predictions[0], jnp.array([0, 1, 2, -1]))
        assert jnp.allclose(
            last_probabilities[0], jnp.array([0.001176, 0.001176, 0.065856])
        )

        X_batch = jnp.array([[1, 2, 0]])
        masks = jnp.full(X_batch.shape, True)
        last_probabilities, predictions = predict(params, X_batch, masks)
        assert jnp.allclose(predictions[0], jnp.array([0, 1, 2]))
        assert jnp.allclose(
            last_probabilities[0], jnp.array([0.001176, 0.001176, 0.009408])
        )

# Definition # https://en.wikipedia.org/wiki/Additive_smoothing # +1 is unk # unk token