# Test of memory profiling


from io import StringIO

import numpy as np

import pytensor.tensor as pt
from pytensor.compile import ProfileStats
from pytensor.compile.function import function
from pytensor.configdefaults import config
from pytensor.ifelse import ifelse
from pytensor.tensor.type import fvector, scalars


class TestProfiling:
    # Test of PyTensor profiling with min_peak_memory=True

    def test_profiling(self):
        config1 = config.profile
        config2 = config.profile_memory
        config3 = config.profiling__min_peak_memory
        try:
            config.profile = True
            config.profile_memory = True
            config.profiling__min_peak_memory = True

            x = [fvector(f"val{i}") for i in range(3)]

            z = []
            z += [pt.outer(x[i], x[i + 1]).sum(axis=1) for i in range(len(x) - 1)]
            z += [x[i] + x[i + 1] for i in range(len(x) - 1)]

            p = ProfileStats(False, gpu_checks=False)

            f = function(x, z, profile=p, name="test_profiling", mode="CVM")

            inp = [np.arange(1024, dtype="float32") + 1 for i in range(len(x))]
            f(*inp)

            buf = StringIO()
            f.profile.summary(buf)

            # regression testing for future algo speed up
            the_string = buf.getvalue()
            lines1 = [l for l in the_string.split("\n") if "Max if linker" in l]
            lines2 = [l for l in the_string.split("\n") if "Minimum peak" in l]
            # NODE: The specific numbers can change for distinct (but correct) toposort orderings
            #  Update the test values if a different algorithm is used
            assert "CPU: 4112KB (4112KB)" in the_string, (lines1, lines2)
            assert "CPU: 8204KB (8204KB)" in the_string, (lines1, lines2)
            assert "CPU: 8208KB" in the_string, (lines1, lines2)
            assert (
                "Minimum peak from all valid apply node order is 4104KB" in the_string
            ), (lines1, lines2)

        finally:
            config.profile = config1
            config.profile_memory = config2
            config.profiling__min_peak_memory = config3

    def test_ifelse(self):
        config1 = config.profile
        config2 = config.profile_memory

        try:
            config.profile = True
            config.profile_memory = True

            a, b = scalars("a", "b")
            x, y = scalars("x", "y")

            z = ifelse(pt.lt(a, b), x * 2, y * 2)

            p = ProfileStats(False, gpu_checks=False)

            if config.mode in ("DebugMode", "DEBUG_MODE", "FAST_COMPILE"):
                m = "FAST_RUN"
            else:
                m = None

            f_ifelse = function([a, b, x, y], z, profile=p, name="test_ifelse", mode=m)

            val1 = 0.0
            val2 = 1.0
            big_mat1 = 10
            big_mat2 = 11

            f_ifelse(val1, val2, big_mat1, big_mat2)

        finally:
            config.profile = config1
            config.profile_memory = config2
