#!/usr/bin/env python

"""Tests for `more_math` package."""

import os
import sys
import json

# Ensure test runner (Visual Studio) can import the package regardless of working dir.
# If repository uses `src/` layout, add that to sys.path; otherwise add project root.
_here = os.path.abspath(os.path.dirname(__file__))
_project_root = os.path.abspath(os.path.join(_here, os.pardir))
_src_path = os.path.join(_project_root, "src")

if os.path.isdir(_src_path) and _src_path not in sys.path:
    sys.path.insert(0, _src_path)
elif _project_root not in sys.path:
    sys.path.insert(0, _project_root)

# Add ComfyUI root to path to find 'comfy' and 'comfy_api' packages
_comfy_root = os.path.abspath(os.path.join(_here, "../../.."))
if _comfy_root not in sys.path:
    sys.path.insert(0, _comfy_root)

import torch
from more_math.ConditioningMathNode import ConditioningMathNode
from more_math.LatentMathNode import LatentMathNode
from more_math.ImageMathNode import ImageMathNode
from more_math.FloatMathNode import FloatMathNode
from more_math.AudioMathNode import AudioMathNode


# ==========================================
# Node Initialization and Metadata Tests
# ==========================================


def test_conditioning_math_node_initialization():
    node = ConditioningMathNode()
    assert isinstance(node, ConditioningMathNode)


def test_conditioning_math_node_metadata():
    assert ConditioningMathNode.RETURN_TYPES == ["CONDITIONING", "STACK"]
    assert ConditioningMathNode.FUNCTION == "EXECUTE_NORMALIZED"
    assert ConditioningMathNode.CATEGORY == "More math/math"


def test_latent_math_node_initialization():
    node = LatentMathNode()
    assert isinstance(node, LatentMathNode)


def test_latent_math_node_metadata():
    assert LatentMathNode.RETURN_TYPES == ["LATENT", "STACK"]
    assert LatentMathNode.FUNCTION == "EXECUTE_NORMALIZED"
    assert LatentMathNode.CATEGORY == "More math/math"


def test_image_math_node_initialization():
    node = ImageMathNode()
    assert isinstance(node, ImageMathNode)


def test_image_math_node_metadata():
    assert ImageMathNode.RETURN_TYPES == ["IMAGE", "STACK"]
    assert ImageMathNode.FUNCTION == "EXECUTE_NORMALIZED"
    assert ImageMathNode.CATEGORY == "More math/math"


# ==========================================
# FFT Tests
# ==========================================


def test_fft_invertibility():
    # Create random input latent (Batch, Channel, Height, Width)
    input_tensor = torch.randn(1, 4, 32, 32, dtype=torch.float32)
    input_dict = {"samples": input_tensor}
    # Execute ifft(fft(a))
    input_V = {"V0": input_dict}
    result_list, stack = LatentMathNode.execute(Expression="ifft(fft(a))", V=input_V, F={}, batching=0)
    output_tensor = result_list[0]["samples"]
    assert torch.allclose(input_tensor, output_tensor, atol=1e-5), f"Max difference: {(input_tensor - output_tensor).abs().max()}"


def test_image_fft_dims():
    # Image input is (Batch, Height, Width, Channel)
    input_tensor = torch.randn(1, 32, 32, 3, dtype=torch.float32)
    result_list, stack = ImageMathNode.execute(Expression="ifft(fft(a))", V={"V0": input_tensor}, F={}, batching=0)
    output_tensor = result_list[0]
    assert input_tensor.shape == output_tensor.shape
    assert torch.allclose(input_tensor, output_tensor, atol=1e-5), (
        f"Image FFT round trip failed. Max diff: {(input_tensor - output_tensor).abs().max()}"
    )


# ==========================================
# Latent Math Basic Functions (Evaluated on Tensors)
# ==========================================


def test_latent_lerp():
    node = LatentMathNode()
    l_a = {"samples": torch.zeros(1, 4, 32, 32)}
    l_b = {"samples": torch.full((1, 4, 32, 32), 10.0)}
    result_list, stack = node.execute(Expression="lerp(a, b, 0.5)", V={"V0": l_a, "V1": l_b}, F={}, batching=0)
    res_lerp = result_list[0]["samples"]
    assert torch.allclose(res_lerp, torch.full_like(res_lerp, 5.0))


def test_latent_step_true():
    node = LatentMathNode()
    # step(x, edge) where x=0.8, edge=0.5 -> 1
    result_list, stack = node.execute(Expression="step(a, 0.5)", V={"V0": {"samples": torch.full((1, 1, 1, 1), 0.8)}}, F={}, batching=0)
    res_step = result_list[0]["samples"]
    assert torch.allclose(res_step, torch.ones_like(res_step))


def test_latent_step_false():
    node = LatentMathNode()
    # step(x, edge) where x=0.2, edge=0.5 -> 0
    result_list, stack = node.execute(Expression="step(a, 0.5)", V={"V0": {"samples": torch.full((1, 1, 1, 1), 0.2)}}, F={}, batching=0)
    res_step2 = result_list[0]["samples"]
    assert torch.allclose(res_step2, torch.zeros_like(res_step2))


def test_latent_swap():
    node = LatentMathNode()
    t_lat = torch.tensor([0.0, 10.0, 20.0, 30.0]).view(1, 4, 1, 1)
    # Swap channels 0 and 3 -> 30, 10, 20, 0
    l_swap = {"samples": t_lat}
    result_list, stack = node.execute(Expression="swap(a, 1, 0, 3)", V={"V0": l_swap}, F={}, batching=0)
    res_swap = result_list[0]["samples"]
    expected = torch.tensor([30.0, 10.0, 20.0, 0.0]).view(1, 4, 1, 1)
    assert torch.allclose(res_swap, expected)


def test_latent_relu():
    node = LatentMathNode()
    l_a = {"samples": torch.zeros(1, 4, 32, 32)}
    result_list, stack = node.execute(Expression="relu(-5.0)", V={"V0": l_a}, F={}, batching=0)
    res_relu = result_list[0]["samples"]
    assert torch.allclose(res_relu, torch.zeros_like(res_relu))


def test_latent_sign():
    node = LatentMathNode()
    l_a = {"samples": torch.zeros(1, 4, 32, 32)}
    result_list, stack = node.execute(Expression="sign(-5.0)", V={"V0": l_a}, F={}, batching=0)
    res_sign = result_list[0]["samples"]
    assert torch.allclose(res_sign, torch.full_like(res_sign, -1.0))


def test_latent_fract():
    node = LatentMathNode()
    l_a = {"samples": torch.zeros(1, 4, 32, 32)}
    result_list, stack = node.execute(Expression="fract(1.5)", V={"V0": l_a}, F={}, batching=0)
    res_fract = result_list[0]["samples"]
    assert torch.allclose(res_fract, torch.full_like(res_fract, 0.5))


# ==========================================
# Float Math Basic Functions (Evaluated on Scalars)
# ==========================================


def test_float_lerp():
    node = FloatMathNode()
    res = node.execute(FloatFunc="lerp(a, b, 0.5)", V={"V0": 0.0, "V1": 10.0})[0]
    assert abs(res - 5.0) < 1e-5


def test_float_step():
    node = FloatMathNode()
    res = node.execute(FloatFunc="step(a, 0.5)", V={"V0": 0.8})[0]
    assert abs(res - 1.0) < 1e-5


def test_float_relu():
    node = FloatMathNode()
    res = node.execute(FloatFunc="relu(a)", V={"V0": -5.0})[0]
    assert abs(res - 0.0) < 1e-5


def test_float_smoothstep():
    node = FloatMathNode()
    res = node.execute(FloatFunc="smoothstep(a, 0, 1)", V={"V0": 0.5})[0]
    assert abs(res - 0.5) < 1e-5


# ==========================================
# Float Math Extended Functions
# ==========================================


def test_float_fract():
    node = FloatMathNode()
    res = node.execute(FloatFunc="fract(a)", V={"V0": 1.5})[0]
    assert abs(res - 0.5) < 1e-5


def test_float_softplus():
    node = FloatMathNode()
    res = node.execute(FloatFunc="softplus(a)", V={"V0": 0.0})[0]
    assert abs(res - 0.69314718) < 1e-5


def test_float_sign_negative():
    node = FloatMathNode()
    assert node.execute(FloatFunc="sign(a)", V={"V0": -10.0})[0] == -1.0


def test_float_sign_positive():
    node = FloatMathNode()
    assert node.execute(FloatFunc="sign(a)", V={"V0": 10.0})[0] == 1.0


def test_float_sign_zero():
    node = FloatMathNode()
    assert node.execute(FloatFunc="sign(a)", V={"V0": 0.0})[0] == 0.0


def test_float_gelu():
    node = FloatMathNode()
    assert node.execute(FloatFunc="gelu(a)", V={"V0": 0.0})[0] == 0.0


# ==========================================
# Latent Math Extended Functions
# ==========================================


def test_latent_smoothstep():
    node = LatentMathNode()
    l_a = {"samples": torch.zeros(1, 4, 32, 32)}
    result_list, stack = node.execute(Expression="smoothstep(0.5, 0, 1)", V={"V0": l_a}, F={}, batching=0)
    res = result_list[0]["samples"]
    assert torch.allclose(res, torch.full_like(res, 0.5))


def test_latent_softplus():
    node = LatentMathNode()
    l_a = {"samples": torch.zeros(1, 4, 32, 32)}
    result_list, stack = node.execute(Expression="softplus(0.0)", V={"V0": l_a}, F={}, batching=0)
    res = result_list[0]["samples"]
    assert torch.allclose(res, torch.full_like(res, 0.69314718))


def test_latent_gelu():
    node = LatentMathNode()
    l_a = {"samples": torch.zeros(1, 4, 32, 32)}
    result_list, stack = node.execute(Expression="gelu(0.0)", V={"V0": l_a}, F={}, batching=0)
    res = result_list[0]["samples"]
    assert torch.allclose(res, torch.zeros_like(res))


# ==========================================
# Image Math Operations
# ==========================================


def test_image_lerp():
    node = ImageMathNode()
    img_red = torch.tensor([1.0, 0.0, 0.0]).view(1, 1, 1, 3)
    img_blue = torch.tensor([0.0, 0.0, 1.0]).view(1, 1, 1, 3)
    result_list, stack = node.execute(Expression="lerp(a, b, 0.5)", V={"V0": img_red, "V1": img_blue}, F={})
    res_blend = result_list[0]
    expected = torch.tensor([0.5, 0.0, 0.5]).view(1, 1, 1, 3)
    assert torch.allclose(res_blend, expected)


def test_image_swap():
    node = ImageMathNode()
    img_red = torch.tensor([1.0, 0.0, 0.0]).view(1, 1, 1, 3)
    img_blue = torch.tensor([0.0, 0.0, 1.0]).view(1, 1, 1, 3)
    result_list, stack = node.execute(Expression="swap(a, 3, 0, 2)", V={"V0": img_red}, F={})
    res_swap = result_list[0]
    assert torch.allclose(res_swap, img_blue)


# ==========================================
# Audio Math Operations
# ==========================================


def test_audio_math_basic():
    node = AudioMathNode()
    waveform = torch.randn(1, 1, 1024)
    audio = {"waveform": waveform, "sample_rate": 44100}
    # result = a * 2.0
    result_list, stack = node.execute(Expression="a * 2.0", V={"V0": audio}, F={})
    res = result_list[0]
    assert isinstance(res, dict)
    assert "waveform" in res
    assert res["sample_rate"] == 44100
    assert torch.allclose(res["waveform"], waveform * 2.0)


# ==========================================
# Nested Expressions
# ==========================================


def test_float_nested_expressions_true():
    node = FloatMathNode()
    # lerp(0, 10, step(0.8, 0.5)) -> lerp(0, 10, 1) -> 10
    res = node.execute(FloatFunc="lerp(0, 10, step(0.8, 0.5))", V={"V0": 0.0})[0]
    assert res == 10.0


def test_float_nested_expressions_false():
    node = FloatMathNode()
    # lerp(0, 10, step(0.2, 0.5)) -> lerp(0, 10, 0) -> 0
    res2 = node.execute(FloatFunc="lerp(0, 10, step(0.2, 0.5))", V={"V0": 0.0})[0]
    assert res2 == 0.0


# ==========================================
# 5D Tensor Support
# ==========================================


def test_5d_tensors_identity():
    node = LatentMathNode()
    samples = torch.randn(1, 5, 4, 32, 32)
    l_in = {"samples": samples}
    result_list, stack = node.execute(Expression="a * 1.0", V={"V0": l_in}, F={}, batching=0)
    res = result_list[0]["samples"]
    assert res.shape == (1, 5, 4, 32, 32)
    assert torch.allclose(res, samples)


def test_5d_tensors_variable_T():
    node = LatentMathNode()
    samples = torch.randn(1, 5, 4, 32, 32)
    l_in = {"samples": samples}
    # In 5D, T maps to dim -4 (size 5)
    result_list, stack = node.execute(Expression="a + T", V={"V0": l_in}, F={}, batching=0)
    res_t = result_list[0]["samples"]
    assert torch.allclose(res_t, samples + 5.0)


def test_5d_tensors_fft():
    node = LatentMathNode()
    samples = torch.randn(1, 5, 4, 32, 32)
    l_in = {"samples": samples}
    result_list, stack = node.execute(Expression="ifft(fft(a))", V={"V0": l_in}, F={}, batching=0)
    res_fft = result_list[0]["samples"]
    assert torch.allclose(res_fft, samples, atol=1e-5)


# ==========================================
# Noise Math Node 5D Support
# ==========================================


def test_noise_math_5d():
    from more_math.NoiseMathNode import NoiseMathNode

    class MockNoise:
        def __init__(self, tensor):
            self.tensor = tensor
            self.seed = 0

        def generate_noise(self, input_latent):
            return self.tensor

    node = NoiseMathNode()
    samples = torch.randn(1, 5, 4, 32, 32)
    noise_a = MockNoise(samples)

    result_executor = node.execute("a + T", V={"V0": noise_a}, F={})[0]
    dummy_latent = {"samples": samples}
    res = result_executor.generate_noise(dummy_latent)

    assert res.shape == (1, 5, 4, 32, 32)
    assert torch.allclose(res, samples + 5.0)


def test_noise_math_autogrow():
    from more_math.NoiseMathNode import NoiseMathNode

    class MockNoise:
        def __init__(self, tensor):
            self.tensor = tensor
            self.seed = 0

        def generate_noise(self, input_latent):
            return self.tensor

    node = NoiseMathNode()
    samples = torch.randn(1, 4, 32, 32)
    noise_a = MockNoise(samples * 1.0)
    noise_b = MockNoise(samples * 2.0)

    # Expression uses 'a', 'b' (aliases for V0, V1) and 'w' (alias for F0)
    result_executor = node.execute("a + b * w", V={"V0": noise_a, "V1": noise_b}, F={"F0": 0.5})[0]
    dummy_latent = {"samples": samples}
    res = result_executor.generate_noise(dummy_latent)

    assert res.shape == (1, 4, 32, 32)
    # 1.0 + 2.0 * 0.5 = 2.0
    assert torch.allclose(res, samples * 2.0)


# ==========================================
# NestedTensor Support
# ==========================================


def test_nested_tensor_support():
    try:
        from comfy.nested_tensor import NestedTensor
    except ImportError:
        assert False, "Could not import comfy.nested_tensor."

    node = LatentMathNode()
    t1 = torch.full((1, 4, 32, 32), 1.0)
    t2 = torch.full((2, 4, 32, 32), 2.0)
    nt_in = NestedTensor([t1, t2])
    l_in = {"samples": nt_in}

    # The base V0 variable returns the original NestedTensor unchanged.
    result_list, stack = node.execute(Expression="V0", V={"V0": l_in}, F={}, batching=0)
    assert len(result_list) == 1
    samples = result_list[0]["samples"]
    assert getattr(samples, "is_nested", False), "V0 should return the original NestedTensor"
    assert torch.allclose(samples.tensors[0], torch.full((1, 4, 32, 32), 1.0))
    assert torch.allclose(samples.tensors[1], torch.full((2, 4, 32, 32), 2.0))

    # Component access uses V0_0, V0_1, ...
    # Arithmetic on the whole NestedTensor returns a single NestedTensor again
    result_list, stack = node.execute(Expression="V0 * 0.5 + 0.1", V={"V0": l_in}, F={}, batching=0)
    assert len(result_list) == 1
    samples = result_list[0]["samples"]
    assert getattr(samples, "is_nested", False), "Arithmetic on a NestedTensor should return a NestedTensor"
    assert torch.allclose(samples.tensors[0], torch.full((1, 4, 32, 32), 0.6))
    assert torch.allclose(samples.tensors[1], torch.full((2, 4, 32, 32), 1.1))

    result_list, stack = node.execute(Expression="V0_0 + 1.0", V={"V0": l_in}, F={}, batching=0)
    res_0 = result_list[0]["samples"]
    assert isinstance(res_0, torch.Tensor), f"Expected torch.Tensor, got {type(res_0)}"
    assert not getattr(res_0, "is_nested", False), "Result should be a regular tensor, not NestedTensor"
    assert torch.allclose(res_0, torch.full((1, 4, 32, 32), 2.0))

    result_list, stack = node.execute(Expression="V0_1 + 1.0", V={"V0": l_in}, F={}, batching=0)
    res_1 = result_list[0]["samples"]
    assert torch.allclose(res_1, torch.full((2, 4, 32, 32), 3.0))


def test_nested_tensor_support_position_v1():
    from comfy.nested_tensor import NestedTensor

    node = LatentMathNode()
    t1 = torch.full((1, 4, 32, 32), 10.0)
    t2 = torch.full((2, 4, 32, 32), 20.0)
    nt_in = NestedTensor([t1, t2])
    l_in = {"samples": nt_in}

    base_latent = {"samples": torch.zeros((1, 4, 32, 32))}

    # V1 returns the original NestedTensor unchanged for downstream compatibility
    result_list, _ = node.execute(Expression="V1", V={"V0": base_latent, "V1": l_in}, F={}, batching=0)
    assert len(result_list) == 1
    samples = result_list[0]["samples"]
    assert getattr(samples, "is_nested", False)
    assert torch.allclose(samples.tensors[0], torch.full((1, 4, 32, 32), 10.0))
    assert torch.allclose(samples.tensors[1], torch.full((2, 4, 32, 32), 20.0))

    # Arithmetic over the whole NestedTensor returns a NestedTensor again
    result_list, _ = node.execute(Expression="V1 * 0.5 + 0.1", V={"V0": base_latent, "V1": l_in}, F={}, batching=0)
    assert len(result_list) == 1
    samples = result_list[0]["samples"]
    assert getattr(samples, "is_nested", False)
    assert torch.allclose(samples.tensors[0], torch.full((1, 4, 32, 32), 5.1))
    assert torch.allclose(samples.tensors[1], torch.full((2, 4, 32, 32), 10.1))

    # Component access for a NestedTensor at position V1
    result_list, _ = node.execute(Expression="V1_0 + V1_1", V={"V0": base_latent, "V1": l_in}, F={}, batching=0)
    assert torch.allclose(result_list[0]["samples"], torch.full((2, 4, 32, 32), 30.0))

    # Legacy alias b should map to the first component of V1
    result_list, _ = node.execute(Expression="b * 2.0", V={"V0": base_latent, "V1": l_in}, F={}, batching=0)
    assert torch.allclose(result_list[0]["samples"], torch.full((1, 4, 32, 32), 20.0))

    # V list variable contains the original V inputs; NestedTensor V1 supports indexing its components
    result_list, _ = node.execute(Expression="V[1][0] + 5.0", V={"V0": base_latent, "V1": l_in}, F={}, batching=0)
    assert torch.allclose(result_list[0]["samples"], torch.full((1, 4, 32, 32), 15.0))

    # V0 is a regular latent (not NestedTensor)
    result_list, _ = node.execute(Expression="V[0] + 5.0", V={"V0": base_latent, "V1": l_in}, F={}, batching=0)
    assert torch.allclose(result_list[0]["samples"], torch.full((1, 4, 32, 32), 5.0))

    # V1_1 is a direct component variable
    result_list, _ = node.execute(Expression="V1_1", V={"V0": base_latent, "V1": l_in}, F={}, batching=0)
    assert torch.allclose(result_list[0]["samples"], torch.full((2, 4, 32, 32), 20.0))


# ==========================================
# Comprehensive Math Function Tests
# ==========================================


def test_trig_functions():
    node = FloatMathNode()
    # Sin/Cos checks
    # sin(0) = 0, cos(0) = 1
    assert abs(node.execute(FloatFunc="sin(0)", V={"V0": 0.0})[0] - 0.0) < 1e-5
    assert abs(node.execute(FloatFunc="cos(0)", V={"V0": 0.0})[0] - 1.0) < 1e-5
    # tan(0) = 0
    assert abs(node.execute(FloatFunc="tan(0)", V={"V0": 0.0})[0] - 0.0) < 1e-5


def test_inverse_trig_functions():
    node = FloatMathNode()
    # asin(0) = 0, acos(1) = 0, atan(0) = 0
    assert abs(node.execute(FloatFunc="asin(0)", V={"V0": 0.0})[0] - 0.0) < 1e-5
    assert abs(node.execute(FloatFunc="acos(1)", V={"V0": 0.0})[0] - 0.0) < 1e-5
    assert abs(node.execute(FloatFunc="atan(0)", V={"V0": 0.0})[0] - 0.0) < 1e-5


def test_pow_log_functions():
    node = FloatMathNode()
    # pow(2, 3) = 8
    assert abs(node.execute(FloatFunc="pow(2, 3)", V={"V0": 0.0})[0] - 8.0) < 1e-5
    # sqrt(4) = 2
    assert abs(node.execute(FloatFunc="sqrt(4)", V={"V0": 0.0})[0] - 2.0) < 1e-5
    # exp(0) = 1
    assert abs(node.execute(FloatFunc="exp(0)", V={"V0": 0.0})[0] - 1.0) < 1e-5
    # log(100) = 2 (base 10)
    assert abs(node.execute(FloatFunc="log(100)", V={"V0": 0.0})[0] - 2.0) < 1e-5
    # ln(e) = 1. Using 'e' constant logic check or approx 2.718
    assert abs(node.execute(FloatFunc="ln(2.7182818)", V={"V0": 0.0})[0] - 1.0) < 1e-4


def test_min_max_functions():

    node = FloatMathNode()
    print("Testing tmin...", flush=True)
    # tmin(2, 5) = 2, tmax(2, 5) = 5
    assert abs(node.execute(FloatFunc="tmin(2, 5)", V={"V0": 0.0})[0] - 2.0) < 1e-5
    print("Testing tmax...", flush=True)
    assert abs(node.execute(FloatFunc="tmax(2, 5)", V={"V0": 0.0})[0] - 5.0) < 1e-5

    # smin/smax (Smooth min/max? Or just multi-arg min/max? TensorEvalVisitor uses stack.min/max)
    # smin(1, 2, 3) = 1
    print("Testing smin...", flush=True)
    res_smin = node.execute(FloatFunc="smin(1, 2, 3)", V={"V0": 0.0})[0]
    print(f"smin result: {res_smin} type: {type(res_smin)}", flush=True)
    assert abs(res_smin - 1.0) < 1e-5

    print("Testing smax...", flush=True)
    res_smax = node.execute(FloatFunc="smax(1, 2, 3)", V={"V0": 0.0})[0]
    print(f"smax result: {res_smax} type: {type(res_smax)}", flush=True)
    assert abs(res_smax - 3.0) < 1e-5


def test_basic_utilities():
    node = FloatMathNode()
    # abs(-5) = 5
    assert abs(node.execute(FloatFunc="abs(-5)", V={"V0": 0.0})[0] - 5.0) < 1e-5
    # floor(1.9) = 1
    assert abs(node.execute(FloatFunc="floor(1.9)", V={"V0": 0.0})[0] - 1.0) < 1e-5
    # ceil(1.1) = 2
    assert abs(node.execute(FloatFunc="ceil(1.1)", V={"V0": 0.0})[0] - 2.0) < 1e-5
    # round(1.6) = 2, round(1.4) = 1
    assert abs(node.execute(FloatFunc="round(1.6)", V={"V0": 0.0})[0] - 2.0) < 1e-5
    assert abs(node.execute(FloatFunc="round(1.4)", V={"V0": 0.0})[0] - 1.0) < 1e-5
    # clamp(10, 0, 5) = 5, clamp(-5, 0, 5) = 0
    assert abs(node.execute(FloatFunc="clamp(10, 0, 5)", V={"V0": 0.0})[0] - 5.0) < 1e-5
    assert abs(node.execute(FloatFunc="clamp(-5, 0, 5)", V={"V0": 0.0})[0] - 0.0) < 1e-5


def test_advanced_activations():
    node = FloatMathNode()
    # sigm(0) = 0.5
    assert abs(node.execute(FloatFunc="sigm(0)", V={"V0": 0.0})[0] - 0.5) < 1e-5


# ==========================================
# Text Functions Tests
# ==========================================

def test_text_upper():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    expr = 'upper("hello")'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({}, (1,))
    result = visitor.visit(tree)
    assert result == "HELLO"


def test_text_lower():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    expr = 'lower("HELLO")'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({}, (1,))
    result = visitor.visit(tree)
    assert result == "hello"


def test_text_trim():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    expr = 'trim("  hello world  ")'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({}, (1,))
    result = visitor.visit(tree)
    assert result == "hello world"


def test_text_split():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    expr = 'split("a,b,c", ",")'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({}, (1,))
    result = visitor.visit(tree)
    assert result == ["a", "b", "c"]


def test_text_join():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    expr = 'join(["a", "b", "c"], "-")'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({}, (1,))
    result = visitor.visit(tree)
    assert result == "a-b-c"


def test_text_substring():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    expr = 'substring("hello world", 0, 5)'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({}, (1,))
    result = visitor.visit(tree)
    assert result == "hello"


def test_text_find():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    expr = 'find("hello world", "world")'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({}, (1,))
    result = visitor.visit(tree)
    assert result == 6.0  # Position of "world"


def test_text_replace():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    expr = 'replace("hello world", "world", "python")'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({}, (1,))
    result = visitor.visit(tree)
    assert result == "hello python"


def test_dict_literal_and_stringification():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    expr = '{a:1, b:"two", nested:{c:3}}'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({}, (1,))
    result = visitor.visit(tree)

    assert result["a"] == 1.0
    assert result["b"] == "two"
    assert result["nested"]["c"] == 3.0
    assert str(result) == '{"a": 1.0, "b": "two", "nested": {"c": 3.0}}'


def test_dict_types_stringify_as_json_with_tensor_arrays():
    from more_math.Parser.UnifiedMathVisitor import ChangeTrackingDict, MathDict

    values = {
        "items": [1.0, {"nested": True}],
        "tensor": torch.tensor([[2, 3]]),
    }
    expected = {"items": [1.0, {"nested": True}], "tensor": [[2, 3]]}

    for dict_type in (MathDict, ChangeTrackingDict):
        result = dict_type(values)
        assert repr(result) == str(result)
        assert json.loads(str(result)) == expected


def test_json_functions_decode_nested_objects_and_encode_tensors():
    from more_math.Parser.UnifiedMathVisitor import MathDict, UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    text = '{"items":[1,{"name":"two"}],"enabled":true}'
    visitor = UnifiedMathVisitor({"text": text, "value": MathDict({"tensor": torch.tensor([1, 2])})}, (1,))
    decoded = visitor.visit(parse_expr("from_json(text)"))
    encoded = visitor.visit(parse_expr("to_json(value)"))

    assert isinstance(decoded, MathDict)
    assert isinstance(decoded["items"][1], MathDict)
    assert decoded["items"][0] == 1
    assert decoded["enabled"] is True
    assert json.loads(str(decoded)) == json.loads(text)
    assert json.loads(encoded) == {"tensor": [1, 2]}
    assert json.loads(visitor.visit(parse_expr("json_encode(value)"))) == {"tensor": [1, 2]}
    assert visitor.visit(parse_expr("json_decode(text)")) == decoded


def test_dict_index_assignment_and_get():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor, MathDict
    from more_math.helper_functions import parse_expr

    expr = 'd["b"] = 2; d["b"]'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({"d": MathDict({"a": 1.0})}, (1,))
    result = visitor.visit(tree)

    assert result == 2.0
    assert visitor.variables["d"]["b"] == 2.0


def test_dict_add_remove_key_helpers():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor, MathDict
    from more_math.helper_functions import parse_expr

    expr = 'add_key(d, "b", 2); remove_key(d, "a"); d'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({"d": MathDict({"a": 1.0})}, (1,))
    result = visitor.visit(tree)

    assert result == {"b": 2.0} or str(result) == "{b:2.0}"


# ==========================================
# Crop Function Tests
# ==========================================

def test_crop_basic():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    # Create a 4x4 tensor
    input_tensor = torch.ones((4, 4))
    expr = 'crop(a, [1, 1], [2, 2])'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({"a": input_tensor}, input_tensor.shape)
    result = visitor.visit(tree)

    # Result should be 2x2
    assert result.shape == (2, 2)
    assert torch.all(result == 1.0)


def test_crop_3d():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    # Create a 4x4x4 tensor
    input_tensor = torch.ones((4, 4, 4)) * 2.0
    expr = 'crop(a, [0, 0, 0], [2, 2, 2])'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({"a": input_tensor}, input_tensor.shape)
    result = visitor.visit(tree)

    # Result should be 2x2x2
    assert result.shape == (2, 2, 2)
    assert torch.all(result == 2.0)


def test_crop_with_offset():
    from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
    from more_math.helper_functions import parse_expr

    # Create a 6x6 tensor with different values
    input_tensor = torch.arange(36).reshape((6, 6)).float()
    expr = 'crop(a, [2, 2], [2, 2])'
    tree = parse_expr(expr)
    visitor = UnifiedMathVisitor({"a": input_tensor}, input_tensor.shape)
    result = visitor.visit(tree)

    # Result should be 2x2
    assert result.shape == (2, 2)
    # Values should be from the cropped region
    expected = torch.tensor([[14.0, 15.0], [20.0, 21.0]])
    assert torch.allclose(result, expected)


if __name__ == "__main__":
    import sys

    try:
        test_conditioning_math_node_initialization()
        test_conditioning_math_node_metadata()
        test_latent_math_node_initialization()
        test_latent_math_node_metadata()
        test_image_math_node_initialization()
        test_image_math_node_metadata()
        test_trig_functions()
        test_inverse_trig_functions()
        test_pow_log_functions()
        test_min_max_functions()
        test_basic_utilities()
        test_advanced_activations()
        print("All tests passed!")
    except Exception:
        import traceback

        traceback.print_exc()
        sys.exit(1)
    print("All tests in test_more_math.py passed!")