text: from pathlib import Path
  
  import lightgbm as lgb
  import numpy as np
  import numpy.random
  import pandas as pd
  import pytest
  from hypothesis import given, settings
  from hypothesis import strategies as st
  
  from lleaves import Model
  
  numpy.random.seed(1337)
  
  
  @pytest.fixture(scope="session", params=["pure_cat", "mixed_cat"])
  def categorical_model_txt(tmpdir_factory, request):
      n_features = 5
      n_rows = 500
      # specify a tree and generate data from it
      if request.param == "pure_cat":
          n_categorical = 5
  
          def tree(axis):
              if axis[0] < 5:
                  if axis[1] < 5:
                      return 2
                  elif axis[0] < 2:
                      return 50
                  else:
                      if axis[2] > 7:
                          return 10
                      return 2
              else:
                  if axis[1] > 8:
                      return 99
                  return 0
  
      else:
          n_categorical = 3
  
          def tree(axis):
              if axis[0] < 5:
                  if axis[1] < 5:
                      return 2
                  elif axis[1] < 2:
                      return 50
                  else:
                      if axis[2] > 7:
                          return 100
                      return -32
              else:
                  if axis[4] > 0.8:
                      return 99
                  else:
                      if axis[3] < 0.2:
                          return -20
                  return 0
  
      train_data_cat = np.random.randint(1, 10, size=(n_rows, n_categorical))
      if n_features - n_categorical > 0:
          train_data_num = np.random.rand(n_rows, n_features - n_categorical)
          train_data = np.concatenate([train_data_cat, train_data_num], axis=1)
      else:
          train_data = train_data_cat
  
      label = np.apply_along_axis(tree, axis=1, arr=train_data)
      train_data = lgb.Dataset(
          train_data, label=label, categorical_feature=list(range(n_categorical))
      )
  
      param = {}
      lightgbm_model = lgb.train(
          param, train_data, 1, categorical_feature=list(range(n_categorical))
      )
  
      tmpdir = tmpdir_factory.mktemp("model")
      model_path = tmpdir / "model.txt"
      lightgbm_model.save_model(str(model_path))
      return model_path
  
  
  def test_large_categorical(tmpdir_factory):
      # test categorical var with >32 different entries
      def f(x):
          return x % 3
  
      train_data_cat = np.repeat(
          np.expand_dims(np.arange(1, 150), axis=1), repeats=40, axis=0
      )
      label = np.apply_along_axis(f, axis=1, arr=train_data_cat).flatten()
      train_data = lgb.Dataset(train_data_cat, label=label, categorical_feature=[0])
      lightgbm_model = lgb.train({}, train_data, 1, categorical_feature=[0])
  
      tmpdir = tmpdir_factory.mktemp("model")
      model_path = str(tmpdir / "model.txt")
      lightgbm_model.save_model(model_path)
      assert "decision_type=9 9 9" in Path(model_path).read_text()
  
      llvm_model = Model(model_file=model_path)
      llvm_model.compile()
      lgbm_model = lgb.Booster(model_file=model_path)
      tests_data = np.expand_dims(np.arange(0, 210), axis=1)
      numpy.testing.assert_equal(
          llvm_model.predict(tests_data), lgbm_model.predict(tests_data)
      )
  
  
  def test_predict_pandas_categorical(tmpdir_factory):
      def f(x):
          t = ord(x["C1"]) - ord(x["C2"]) ** 2 + ord(x["C3"]) ** 3 + x["A"]
          return t / 100000
  
      train_df = pd.DataFrame(
          {
              "C1": 500 * ["a"] + 500 * ["b"] + 500 * ["c"],
              "C2": 500 * ["b"] + 500 * ["c"] + 500 * ["d"],
              "C3": 10 * (50 * ["x"] + 50 * ["y"] + 25 * ["z"] + 25 * ["w"]),
              "A": 5 * (100 * [10] + 100 * [-10] + 100 * [-30]),
          }
      )
      result = train_df.apply(f, axis=1)
      train_df["C1"] = train_df["C1"].astype("category")
      train_df["C2"] = train_df["C2"].astype("category")
      train_df["C3"] = train_df["C3"].astype("category")
      train_data = lgb.Dataset(train_df, label=result, categorical_feature="auto")
      lightgbm_model = lgb.train({}, train_data, 3, categorical_feature="auto")
      assert len(lightgbm_model.pandas_categorical) == 3
  
      tmpdir = tmpdir_factory.mktemp("model")
      model_path = str(tmpdir / "model.txt")
      lightgbm_model.save_model(model_path)
  
      llvm_model = Model(model_file=model_path)
      llvm_model.compile()
      lgbm_model = lgb.Booster(model_file=model_path)
  
      df = pd.DataFrame(
          {
              "C1": 2 * ["a", "b", "c", "c"],
              "C2": 2 * ["b", "c", "d", "b"],
              "C3": 2 * ["x", "y", "z", "w"],
              "A": [10, -10, 100, -30, -30, 10, 1, 100],
          },
      ).astype({"C1": "category", "C2": "category", "C3": "category"})
      df2 = df.copy()
      # reorder categories to ensure that same letter now maps to a different pandas code
      df2["C1"] = df2["C1"].cat.set_categories(["c", "b", "a"])
      df2["C2"] = df2["C2"].cat.set_categories(["a", "z", "w"])
      assert not (list(df["C1"].cat.codes) == list(df2["C1"].cat.codes))
      assert not (list(df["C2"].cat.codes) == list(df2["C2"].cat.codes))
  
      df3 = pd.DataFrame(
          {
              "C1": 2 * ["a", "b", "c", "s"],
              "C2": 2 * ["b", "c", "v", "b"],
              "C3": 2 * ["x", "u", "z", "w"],
              "A": [10, -10, 100, -30, -30, 10, 1, 100],
          },
      ).astype({"C1": "category", "C2": "category", "C3": "category"})
  
      # prediction from lightgbm should still be equal
      numpy.testing.assert_equal(lightgbm_model.predict(df), lightgbm_model.predict(df2))
      test_data = [
          df,
          df2,
          df3,
          train_df,
          np.array([[1, 1, 1, 10], [2, -2, 2, 10], [0, 1, 2, 10], [2, 1, 0, -10]]),
      ]
      for data in test_data:
          numpy.testing.assert_equal(llvm_model.predict(data), lgbm_model.predict(data))
  
      data = (
          pd.DataFrame(
              {
                  "C1": 2 * ["b", "a", "a", "a"],
                  "C2": 2 * ["d", "c", "a", "b"],
                  "C3": 2 * ["z", "w", "z", "y"],
                  "A": [10, -10, 100, -30, -30, 10, 1, 100],
              }
          ),
      )
      with pytest.raises(ValueError):
          # no categories
          llvm_model.predict(data)
  
  
  @settings(deadline=1000)
  @given(data=st.data())
  def test_categorical_prediction_llvm_real(data, categorical_model_txt):
      lgbm_model = lgb.Booster(model_file=str(categorical_model_txt))
      llvm_model = Model(model_file=categorical_model_txt)
      llvm_model.compile()
  
      input = data.draw(
          st.lists(
              st.integers(min_value=0, max_value=2 ** 31 - 2),
              max_size=llvm_model.num_feature(),
              min_size=llvm_model.num_feature(),
          )
      )
      assert llvm_model.predict([input]) == lgbm_model.predict([input])
  
  
  def test_pure_categorical_prediction():
      llvm_model = Model("tests/models/pure_categorical/model.txt")
      llvm_model.compile()
      lgbm_model = lgb.Booster(model_file="tests/models/pure_categorical/model.txt")
  
      results = [12.616231057968633, 10.048276920678525, 9.2489478721549396]
      for data, res_idx in zip(
          [
              [0, 9, 0],
              [0, -1, 0],
              [1, 9, 0],
              [0, 6, 5],
              [1, 5, 1],
              [2, 5, 1],
              [4, 5, 1],
              [5, 5, 9],
              [6, 5, 3],
              [9, 5, 2],
              [-1, 5, 2],
          ],
          [0, 1, 0, 0, 1, 1, 1, 2, 2, 2, 2],
      ):
          assert lgbm_model.predict([data]) == [results[res_idx]]
          assert llvm_model.predict([data]) == [results[res_idx]]
  
      na = float("NaN")
      inf = float("Inf")
      for data, res_idx in zip(
          [
              [na, 9.0, 0.0],
              [na, 0.0, 0.0],
              [inf, 0.0, 0.0],
              [0.0, na, 0.0],
              [4.0, inf, 0.0],
              [na, na, 0.0],
          ],
          [0, 2, 2, 1, 1, 2],
      ):
          assert lgbm_model.predict([data]) == [results[res_idx]]
          assert llvm_model.predict([data]) == [results[res_idx]]
  