"""Default: learner checks. --reference: verify the supplied reference engine."""
import copy
import json
import math
from pathlib import Path
import sys
import tempfile
import unittest
import stage00

if "--reference" in sys.argv:
    import stage00 as solution
    sys.argv.remove("--reference")
else:
    import starter as solution

CORPUS = Path(__file__).with_name("corpus.json")


class Components(unittest.TestCase):
    def test_hand_loss(self):
        self.assertAlmostEqual(solution.mean_nll([0.5, 0.25]), 1.0397207708399179)
        self.assertAlmostEqual(math.exp(solution.mean_nll([0.5, 0.25])), math.sqrt(8))

    def test_boundary_probabilities(self):
        self.assertEqual(solution.mean_nll([1, 1]), 0)
        self.assertTrue(math.isinf(solution.mean_nll([0, 0.5])))

    def test_invalid_loss_inputs(self):
        for values in ([], [-0.1], [1.1], [float("nan")], [float("inf")]):
            with self.subTest(values=values), self.assertRaises(ValueError):
                solution.mean_nll(values)

    def test_loss_accepts_iterator(self):
        self.assertAlmostEqual(solution.mean_nll(iter([0.1, 0.7, 0.9])),
                               -math.log(0.1 * 0.7 * 0.9) / 3)

    def test_smoothing_normalizes(self):
        for alpha in (0.1, 0.5, 2):
            row = [solution.conditional_probability(c, 6, 3, alpha) for c in (3, 2, 1)]
            self.assertAlmostEqual(sum(row), 1)
        self.assertEqual(solution.conditional_probability(0, 0, 4, 0.5), 0.25)

    def test_invalid_smoothing(self):
        for args in ((-1, 3, 4, 1), (4, 3, 4, 1), (1, 3, 0, 1),
                     (1, 3, 4, 0), (1, 3, 4, float("nan"))):
            with self.subTest(args=args), self.assertRaises(ValueError):
                solution.conditional_probability(*args)


class IntegratedBaseline(unittest.TestCase):
    def setUp(self):
        self.corpus, self.splits, _ = stage00.load_corpus(CORPUS)
        self.model = stage00.fit(self.splits["train"])

    def test_document_boundaries(self):
        self.assertEqual(self.model["totals"][stage00.BOS], len(self.splits["train"]))
        self.assertEqual(self.model["unigrams"][stage00.EOS], len(self.splits["train"]))
        self.assertEqual(self.model["bigrams"][stage00.EOS, stage00.BOS], 0)

    def test_fit_rejects_held_out(self):
        with self.assertRaises(ValueError):
            stage00.fit(self.splits["validation"])

    def test_unknown_characters_do_not_grow_vocabulary(self):
        before = copy.deepcopy(self.model)
        result = stage00.score(self.model, [{"text": "Ω"}], "bigram")
        self.assertEqual(result["unknown_targets"], 1)
        self.assertEqual(result["targets"], 2)  # unknown character + EOS
        self.assertTrue(math.isfinite(result["nll_nats"]))
        self.assertEqual(self.model, before)

    def test_uniform_loss_has_known_value(self):
        result = stage00.score(self.model, self.splits["validation"], "uniform")
        self.assertAlmostEqual(result["nll_nats"], math.log(len(self.model["vocabulary"])))

    def test_token_weighting_not_document_weighting(self):
        docs = [{"text": "a"}, {"text": "aaaa"}]
        together = stage00.score(self.model, docs, "bigram")
        each = [stage00.score(self.model, [d], "bigram") for d in docs]
        weighted = sum(x["nll_nats"] * x["targets"] for x in each) / sum(x["targets"] for x in each)
        self.assertAlmostEqual(together["nll_nats"], weighted)

    def test_report_is_repeatable_and_selects_on_validation(self):
        report = stage00.run_experiment(CORPUS, solution.mean_nll, solution.conditional_probability)
        self.assertEqual(report, stage00.run_experiment(CORPUS, solution.mean_nll, solution.conditional_probability))
        best = min(report["smoothing_candidates"], key=lambda a: report["validation_candidates"][a]["nll_nats"])
        self.assertEqual(report["selected_alpha"], best)
        self.assertFalse(report["test_used_for_selection"])
        self.assertLess(report["metrics"]["validation"]["bigram"]["nll_nats"],
                        report["metrics"]["validation"]["uniform"]["nll_nats"])

    def test_changed_test_text_does_not_change_fit_or_selection(self):
        changed = copy.deepcopy(self.corpus)
        changed["documents"][-1]["text"] = "unseen Ω symbols!"
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "changed.json"
            path.write_text(json.dumps(changed), encoding="utf-8")
            before, after = stage00.run_experiment(CORPUS), stage00.run_experiment(path)
        self.assertEqual(before["vocabulary"], after["vocabulary"])
        self.assertEqual(before["selected_alpha"], after["selected_alpha"])
        self.assertEqual(before["validation_candidates"], after["validation_candidates"])
        self.assertNotEqual(before["corpus_sha256"], after["corpus_sha256"])

    def test_duplicates_and_split_omissions_are_rejected(self):
        bad_cases = []
        duplicate = copy.deepcopy(self.corpus)
        duplicate["documents"][-1]["text"] = duplicate["documents"][0]["text"]
        bad_cases.append(duplicate)
        reused_id = copy.deepcopy(self.corpus)
        reused_id["documents"][-1]["id"] = reused_id["documents"][0]["id"]
        bad_cases.append(reused_id)
        missing = copy.deepcopy(self.corpus)
        missing["documents"] = [d for d in missing["documents"] if d["split"] != "test"]
        bad_cases.append(missing)
        for corpus in bad_cases:
            with self.subTest(), tempfile.TemporaryDirectory() as directory:
                path = Path(directory) / "invalid.json"
                path.write_text(json.dumps(corpus), encoding="utf-8")
                with self.assertRaises(ValueError):
                    stage00.load_corpus(path)


if __name__ == "__main__":
    unittest.main()