import unittest
from pathlib import Path

from accelforge.frontend.spec import Spec
from accelforge.model.main import evaluate_mapping
from accelforge.util.parallel import set_n_parallel_jobs

set_n_parallel_jobs(1)


EXAMPLES_DIR = Path(__file__).parent.parent / "examples"


class TestModel(unittest.TestCase):
    def test_one_matmul(self):
        spec = Spec.from_yaml(
            EXAMPLES_DIR / "arches" / "simple.yaml",
            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
            EXAMPLES_DIR / "mappings" / "unfused_matmuls_to_simple.yaml",
            jinja_parse_data={"N_EINSUMS": 1, "M": 64, "KN": 64},
        )

        result = evaluate_mapping(spec)

    def test_two_matmuls(self):
        spec = Spec.from_yaml(
            EXAMPLES_DIR / "arches" / "simple.yaml",
            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
            EXAMPLES_DIR / "mappings" / "unfused_matmuls_to_simple.yaml",
            jinja_parse_data={"N_EINSUMS": 2, "M": 64, "KN": 64},
        )

        result = evaluate_mapping(spec)


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