"""Arthur end to end in the container: loads every Arthur folder in MODELS_DIR and runs real decisions.

Needs torch (the container has it) and at least one Arthur folder in the models folder; otherwise skipped with the reason.
"""
import importlib.util
import os
import unittest

from app.registry import model_specs

TORCH = importlib.util.find_spec("torch") is not None
ARTHURS = [s for s in model_specs(os.environ) if s["kind"] == "arthur" and s.get("source") == "local"]   # Hub copies need the network
QUESTIONS = {
    "team": {"type": "choice", "instructions": "Which team should handle this?",
             "criteria": {"billing": "Payments and refunds", "tech": "Bugs and errors", "sales": "Pricing and plans"}},
    "angry": {"type": "noul", "instructions": "The customer sounds angry"},
    "urgency": {"type": "score", "instructions": "How urgent is this?", "criteria": ["Low", "Medium", "High", "Critical"]},
}
STATE = {"channel": "email", "message": "I was charged twice for March. Refund the duplicate today or I cancel."}


@unittest.skipUnless(TORCH and ARTHURS, "needs torch and an Arthur folder in the models folder")
class ArthurRuntimeTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        from app import arthur
        cls.arthur = arthur
        cls.loaded = [(s, *arthur.load_arthur(s["path"], "cpu")) for s in ARTHURS]

    def test_weights_load_strictly_and_answers_are_probability_distributions(self):
        for spec, net, temps, _ in self.loaded:
            out = self.arthur.decide(net, temps, STATE, QUESTIONS)
            with self.subTest(model=spec["name"]):
                self.assertEqual(set(out), set(QUESTIONS))
                for r in out.values():
                    self.assertAlmostEqual(sum(r["probs"].values()), 1.0, places=5)
                self.assertEqual(set(out["team"]["probs"]), {"billing", "tech", "sales"})
                self.assertAlmostEqual(out["angry"]["p_true"], out["angry"]["probs"]["True"], places=6)
                self.assertTrue(0.0 <= out["urgency"]["expected_level"] <= 3.0)

    def test_decide_matches_the_notebooks_single_decision_path(self):
        import torch
        from app.arthur_io import normalize_question, probabilities, state_text
        for spec, net, temps, _ in self.loaded:
            for name, q in QUESTIONS.items():
                qtype, text, keys, opts = normalize_question(q)
                d = {"question": text, "options": opts, "state": state_text(STATE), "qtype": qtype}
                with torch.no_grad():
                    z = net(*self.arthur.batch_tensors([d])).float().cpu().numpy()[0, :len(opts)]
                ref = probabilities(z, qtype, temps)
                got = self.arthur.decide(net, temps, STATE, {name: q})[name]["probs"]
                with self.subTest(model=spec["name"], question=name):
                    for k, p in zip(keys, ref):
                        self.assertAlmostEqual(got[str(k)], float(p), places=5)

    def test_batching_questions_together_does_not_change_answers(self):
        for spec, net, temps, _ in self.loaded:
            together = self.arthur.decide(net, temps, STATE, QUESTIONS)
            for name, q in QUESTIONS.items():
                alone = self.arthur.decide(net, temps, STATE, {name: q})[name]["probs"]
                with self.subTest(model=spec["name"], question=name):
                    for k, p in alone.items():
                        self.assertAlmostEqual(together[name]["probs"][k], p, places=4)


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