from pathlib import Path
from unittest import TestCase

import accelforge as af

INPUT_FILES_DIR = Path(__file__).parent / "input_files" / "networked"


class TestParsing(TestCase):
    def test_hierarchical(self):
        spec = af.Spec.from_yaml(
            # af.examples.arches.networked.hierarchical,
            INPUT_FILES_DIR
            / "hierarchical.yaml",
        )
        self.assertIn("PeArray", spec.arch.nodes)
        self.assertEqual(spec.arch.nodes["PeArray"].get_fanout(), 1)
        self.assertIn("Scratchpad", spec.arch.nodes)
        self.assertEqual(spec.arch.nodes["Scratchpad"].get_fanout(), 4)
        self.assertIn("MacArray", spec.arch.nodes)
        self.assertEqual(spec.arch.nodes["MacArray"].get_fanout(), 1)

        try:
            spec = spec.calculate_component_area_energy_latency_leak()
        except af.EvaluationError as e:
            self.fail(e.message)

    def test_flat(self):
        spec = af.Spec.from_yaml(
            # af.examples.arches.networked.flat,
            INPUT_FILES_DIR
            / "flat.yaml",
        )
        self.assertIn("NoC", spec.arch.nodes)
        self.assertEqual(spec.arch.nodes["NoC"].get_fanout(), 1)
        self.assertEqual(
            {n.name for n in spec.arch.get_nodes_of_type(af.spec.Leaf)},
            {
                "MainMemory",
                "GlobalBuffer",
                "NoC",
                "RowBuffer",
                "ColumnBuffer",
                "DistributedBuffer",
                "Scratchpad",
                "MAC",
            },
        )

        try:
            spec = spec.calculate_component_area_energy_latency_leak()
        except af.EvaluationError as e:
            self.fail(e.message)


class TestModel(TestCase):
    def test_hierarchical_1d(self):
        M = 8
        KN = 8
        MAC_TILE = 2
        M_TILE = 4
        BITS_PER_VALUE = 8

        spec = af.Spec.from_yaml(
            af.examples.workloads.matmuls,
            # af.examples.arches.networked.hierarchical,
            INPUT_FILES_DIR / "hierarchical_1d.yaml",
            # af.examples.mappings.one_matmul_to_networked_hierarchical,
            INPUT_FILES_DIR / "one_matmul_to_networked_hierarchical_1d.yaml",
            jinja_parse_data={
                "N_EINSUMS": 1,
                "M": 8,
                "KN": 8,
                "MAC_TILE": MAC_TILE,
                "M_TILE": M_TILE,
            },
        )
        result = spec.evaluate_mapping()
        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>MacArray<SEP>T0<SEP>hops"].iloc[0],
            (M / M_TILE)
            * (KN / MAC_TILE)  # number of used Scratchpad
            * M_TILE
            * KN  # temporal for n1 in mapping
            * sum(i + 1 for i in range(MAC_TILE))  # unicast along X-axis of MacArray
            * BITS_PER_VALUE,
        )
        # NOTE: assuming XY routing (as defined in mapping)
        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>MacArray<SEP>T1<SEP>hops"].iloc[0],
            (M / M_TILE)
            * (KN / MAC_TILE)
            * M_TILE
            * KN  # temporal for n1 in mapping
            * MAC_TILE  # multicast along X-axis of MacArray
            * BITS_PER_VALUE,
        )
        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>MacArray<SEP>W0<SEP>hops"].iloc[0],
            (M / M_TILE)
            * (KN / MAC_TILE)
            * M_TILE
            * KN
            * sum(i + 1 for i in range(MAC_TILE))
            * BITS_PER_VALUE,
        )

        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>PeArray<SEP>T0<SEP>hops"].iloc[0],
            (M / M_TILE)
            * sum(
                i + 1 for i in range(KN // MAC_TILE)
            )  # unicast along X-axis of PeArray
            * M_TILE
            * MAC_TILE
            * BITS_PER_VALUE,
        )
        # NOTE: assuming XY routing (as defined in mapping)
        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>PeArray<SEP>T1<SEP>hops"].iloc[0],
            (M / M_TILE)
            * KN
            // MAC_TILE  # multicast along X-axis of PeArray
            * M_TILE
            * KN
            * BITS_PER_VALUE,
        )
        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>PeArray<SEP>W0<SEP>hops"].iloc[0],
            (M / M_TILE)
            * sum(i + 1 for i in range(KN // MAC_TILE))  # unicast along PeArray
            * MAC_TILE
            * KN
            * BITS_PER_VALUE,
        )

    def test_hierarchical(self):
        M = 8
        KN = 8
        MAC_TILE = 2
        PE_TILE = KN // MAC_TILE
        M_TILE = 4
        BITS_PER_VALUE = 8

        spec = af.Spec.from_yaml(
            af.examples.workloads.matmuls,
            # af.examples.arches.networked.hierarchical,
            INPUT_FILES_DIR / "hierarchical.yaml",
            # af.examples.mappings.one_matmul_to_networked_hierarchical,
            INPUT_FILES_DIR / "one_matmul_to_networked_hierarchical.yaml",
            jinja_parse_data={
                "N_EINSUMS": 1,
                "M": 8,
                "KN": 8,
                "MAC_TILE": MAC_TILE,
                "M_TILE": M_TILE,
            },
        )
        result = spec.evaluate_mapping()
        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>MacArray<SEP>T0<SEP>hops"].iloc[0],
            (M / M_TILE)
            * (KN / MAC_TILE) ** 2
            * M_TILE
            * (
                sum(i + 1 for i in range(MAC_TILE))  # unicasting along X
                + MAC_TILE * MAC_TILE  # multicast along Y for each column
            )
            * BITS_PER_VALUE,
        )
        # NOTE: assuming XY routing (as defined in mapping)
        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>MacArray<SEP>T1<SEP>hops"].iloc[0],
            (M / M_TILE)
            * (KN / MAC_TILE) ** 2
            * M_TILE
            * (
                MAC_TILE
                * MAC_TILE  # multicast along X (the tile is shape N1, which is MAC_TILE here)
                + MAC_TILE
                * sum(i + 1 for i in range(MAC_TILE))  # unicasting along Y for each row
            )
            * BITS_PER_VALUE,
        )
        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>MacArray<SEP>W0<SEP>hops"].iloc[0],
            (M / M_TILE)
            * (KN / MAC_TILE) ** 2
            * M_TILE
            * (
                MAC_TILE
                * sum(
                    i + 1 for i in range(MAC_TILE)
                )  # unicast along X (the tile is shape N1, which is MAC_TILE here)
                + MAC_TILE
                * sum(i + 1 for i in range(MAC_TILE))  # unicasting along Y for each row
            )
            * BITS_PER_VALUE,
        )

        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>PeArray<SEP>T0<SEP>hops"].iloc[0],
            (M / M_TILE) * (sum(i + 1 for i in range(PE_TILE)) + PE_TILE * PE_TILE)
            # tile shape
            * M_TILE * MAC_TILE * BITS_PER_VALUE,
        )
        # NOTE: assuming XY routing (as defined in mapping)
        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>PeArray<SEP>T1<SEP>hops"].iloc[0],
            (M / M_TILE)
            * (PE_TILE * PE_TILE + PE_TILE * sum(i + 1 for i in range(PE_TILE)))
            * M_TILE
            * MAC_TILE
            * BITS_PER_VALUE,
        )
        self.assertEqual(
            result.data["Matmul0<SEP>action<SEP>PeArray<SEP>W0<SEP>hops"].iloc[0],
            (M / M_TILE)
            * (
                PE_TILE * sum(i + 1 for i in range(PE_TILE))
                + PE_TILE * sum(i + 1 for i in range(PE_TILE))
            )
            * MAC_TILE**2
            * BITS_PER_VALUE,
        )


class TestMapper(TestCase):
    def test_hierarchical(self):
        M = 8
        KN = 8
        MAC_TILE = 2
        PE_TILE = KN // MAC_TILE
        M_TILE = 4
        BITS_PER_VALUE = 8

        spec = af.Spec.from_yaml(
            af.examples.workloads.matmuls,
            # af.examples.arches.networked.hierarchical,
            INPUT_FILES_DIR / "hierarchical.yaml",
            jinja_parse_data={
                "N_EINSUMS": 1,
                "M": 8,
                "KN": 8,
                "MAC_TILE": MAC_TILE,
                "M_TILE": M_TILE,
            },
        )
        result = spec.map_workload_to_arch()
