"""Tests for gherbal.training module."""

import os
import pytest
import pandas as pd
from gherbal.training import (
    train_fasttext,
    quantize_model,
    DEFAULT_DICT_PROPS,
    DEFAULT_TRAIN_OPTS,
)


class TestDefaults:
    """Tests for default training configuration."""

    def test_default_dict_props_keys(self):
        expected = {"minCount", "minCountLabel", "wordNgrams", "bucket",
                    "minn", "maxn", "t", "label"}
        assert set(DEFAULT_DICT_PROPS.keys()) == expected

    def test_default_train_opts_keys(self):
        expected = {"lr", "lrUpdateRate", "dim", "ws", "epoch", "neg", "loss", "thread"}
        assert set(DEFAULT_TRAIN_OPTS.keys()) == expected

    def test_default_label_prefix(self):
        assert DEFAULT_DICT_PROPS["label"] == "__label__"

    def test_defaults_are_reasonable(self):
        assert DEFAULT_TRAIN_OPTS["dim"] > 0
        assert DEFAULT_TRAIN_OPTS["epoch"] > 0
        assert DEFAULT_TRAIN_OPTS["lr"] > 0
        assert DEFAULT_DICT_PROPS["wordNgrams"] >= 1


class TestTrainFasttext:
    """Tests for train_fasttext function."""

    def test_basic_train(self, train_df, tmp_dir):
        model, labels, model_path, labels_path = train_fasttext(
            train_df, tmp_dir,
            train_opts={"epoch": 2, "dim": 20},
        )
        assert model is not None
        assert len(labels) == 3
        assert os.path.exists(model_path)
        assert os.path.exists(labels_path)

    def test_train_creates_bin_by_default(self, train_df, tmp_dir):
        _, _, model_path, _ = train_fasttext(
            train_df, tmp_dir,
            train_opts={"epoch": 2, "dim": 20},
        )
        assert model_path.endswith(".bin")

    def test_train_quantized(self, train_df, tmp_dir):
        _, _, model_path, _ = train_fasttext(
            train_df, tmp_dir,
            train_opts={"epoch": 2, "dim": 20},
            quantize=True,
        )
        assert model_path.endswith(".ftz")

    def test_labels_file_contents(self, train_df, tmp_dir):
        _, labels, _, labels_path = train_fasttext(
            train_df, tmp_dir,
            train_opts={"epoch": 2, "dim": 20},
        )
        with open(labels_path) as f:
            file_labels = [line.strip() for line in f if line.strip()]
        assert file_labels == labels

    def test_labels_sorted(self, train_df, tmp_dir):
        _, labels, _, _ = train_fasttext(
            train_df, tmp_dir,
            train_opts={"epoch": 2, "dim": 20},
        )
        assert labels == sorted(labels)

    def test_temp_file_cleaned(self, train_df, tmp_dir):
        train_fasttext(
            train_df, tmp_dir,
            train_opts={"epoch": 2, "dim": 20},
        )
        temp_files = [f for f in os.listdir(tmp_dir) if f.startswith("_train_")]
        assert len(temp_files) == 0

    def test_custom_dict_props(self, train_df, tmp_dir):
        model, _, _, _ = train_fasttext(
            train_df, tmp_dir,
            dict_props={"wordNgrams": 2, "minn": 2, "maxn": 4},
            train_opts={"epoch": 2, "dim": 20},
        )
        assert model is not None

    def test_custom_train_opts(self, train_df, tmp_dir):
        model, _, _, _ = train_fasttext(
            train_df, tmp_dir,
            train_opts={"epoch": 1, "dim": 10, "lr": 0.5},
        )
        assert model is not None

    def test_model_can_predict(self, train_df, tmp_dir):
        model, labels, _, _ = train_fasttext(
            train_df, tmp_dir,
            train_opts={"epoch": 2, "dim": 20},
        )
        preds, probs = model.predict(["hello world"], k=1)
        assert len(preds) == 1
        assert len(probs) == 1

    def test_does_not_mutate_input_df(self, train_df, tmp_dir):
        original_cols = set(train_df.columns)
        train_fasttext(train_df, tmp_dir, train_opts={"epoch": 2, "dim": 20})
        # train_fasttext works on a copy so original should be unchanged
        assert set(train_df.columns) == original_cols


class TestQuantizeModel:
    """Tests for quantize_model function."""

    def test_quantize(self, train_df, tmp_dir):
        import tempfile
        train_dir = tempfile.mkdtemp()
        model, labels, _, _ = train_fasttext(
            train_df, train_dir,
            train_opts={"epoch": 2, "dim": 20},
        )

        q_model, q_path, q_labels_path = quantize_model(model, labels, tmp_dir)
        assert os.path.exists(q_path)
        assert q_path.endswith(".ftz")
        assert os.path.exists(q_labels_path)

    def test_quantized_model_predicts(self, train_df, tmp_dir):
        import tempfile
        train_dir = tempfile.mkdtemp()
        model, labels, _, _ = train_fasttext(
            train_df, train_dir,
            train_opts={"epoch": 2, "dim": 20},
        )

        q_model, _, _ = quantize_model(model, labels, tmp_dir)
        preds, probs = q_model.predict(["hello world"], k=1)
        assert len(preds) == 1

    def test_quantized_labels_match(self, train_df, tmp_dir):
        import tempfile
        train_dir = tempfile.mkdtemp()
        model, labels, _, _ = train_fasttext(
            train_df, train_dir,
            train_opts={"epoch": 2, "dim": 20},
        )

        _, _, q_labels_path = quantize_model(model, labels, tmp_dir)
        with open(q_labels_path) as f:
            q_labels = [line.strip() for line in f if line.strip()]
        assert q_labels == labels
