from RLTest import Env
from includes import *
from common import *

SCORE_FIELD = "__score"

"""
VECTOR SPACE LAYOUT:
====================

The test data creates a 2D vector space with 4 documents positioned as follows:

    doc:3 ●────────────● doc:4
          │            │
          │            │
          │            │  ● Query
          │            │    Vector
    doc:1 ●────────────● doc:2
        Query
        Vector

    Coordinates:
    - doc:1: (0.0, 0.0) - "red shoes"
    - doc:2: (1.0, 0.0) - "red running shoes"
    - doc:3: (0.0, 1.0) - "running gear"
    - doc:4: (1.0, 1.0) - "blue shoes"
    - Query Vector: (0.0, 0.0)

"""

# Test data with deterministic vectors
test_data = {
    'doc:1': {
        'description': "red shoes",
        'embedding': np.array([0.0, 0.0]).astype(np.float32).tobytes()
    },
    'doc:2': {
        'description': "red running shoes",
        'embedding': np.array([1.0, 0.0]).astype(np.float32).tobytes()
    },
    'doc:3': {
        'description': "running gear",
        'embedding': np.array([0.0, 1.0]).astype(np.float32).tobytes()
    },
    'doc:4': {
        'description': "blue shoes",
        'embedding': np.array([1.0, 1.0]).astype(np.float32).tobytes()
    }
}

def setup_basic_index(env):
    """Setup basic index with test data"""
    conn = env.getClusterConnectionIfNeeded()
    env.expect('FT.CREATE idx SCHEMA description TEXT embedding VECTOR FLAT 6 TYPE FLOAT32 DIM 2 DISTANCE_METRIC L2').ok

    # Load test data
    for doc_id, doc_data in test_data.items():
        conn.execute_command('HSET', doc_id, 'description', doc_data['description'], 'embedding', doc_data['embedding'])

def calculate_l2_distance_normalized(vec1_bytes, vec2_bytes):
    """Calculate L2 distance between two vector byte arrays and normalize"""
    def VectorNorm_L2(distance):
        return 1.0 / (1.0 + distance)

    vec1 = np.frombuffer(vec1_bytes, dtype=np.float32)
    vec2 = np.frombuffer(vec2_bytes, dtype=np.float32)
    return VectorNorm_L2(np.linalg.norm(vec1 - vec2)**2)

def calculate_l2_distance_raw(vec1_bytes, vec2_bytes):
    """Calculate raw L2 distance between two vector byte arrays"""
    vec1 = np.frombuffer(vec1_bytes, dtype=np.float32)
    vec2 = np.frombuffer(vec2_bytes, dtype=np.float32)
    return np.linalg.norm(vec1 - vec2)**2


def test_hybrid_vsim_knn_yield_score_as():
    """Test VSIM KNN with YIELD_SCORE_AS parameter"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()

    response = env.cmd(
        'FT.HYBRID', 'idx',
        'SEARCH', 'shoes',
        'VSIM', '@embedding', '$BLOB',
            'KNN', '2', 'K', '10',
            'YIELD_SCORE_AS', 'vector_score',
        'PARAMS', '2', 'BLOB', query_vector)
    results, _ = get_results_from_hybrid_response(response)

    # Validate the score field for all returned results
    env.assertGreater(len(results.keys()), 0)  # Should return docs with "shoes" in description

    for doc_key in results:
        doc_result = results[doc_key]
        env.assertTrue('vector_score' in doc_result)
        returned_distance = float(doc_result['vector_score'])
        expected_distance = calculate_l2_distance_normalized(query_vector, test_data[doc_key]['embedding'])
        env.assertAlmostEqual(returned_distance, expected_distance, delta=1e-6)


def test_hybrid_vsim_range_yield_score_as():
    """Test VSIM RANGE with YIELD_SCORE_AS parameter"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()
    radius = 2

    response = env.cmd(
        'FT.HYBRID', 'idx',
        'SEARCH', 'shoes',
        'VSIM', '@embedding', '$BLOB',
            'RANGE', '2', 'RADIUS', str(radius),
            'YIELD_SCORE_AS', 'vector_score',
        'PARAMS', '2', 'BLOB', query_vector)
    results, _ = get_results_from_hybrid_response(response)

    # Validate the vector_score field for all returned results
    env.assertGreater(len(results.keys()), 0)

    for doc_key in results:
        doc_result = results[doc_key]
        env.assertTrue('vector_score' in doc_result)
        returned_distance = float(doc_result['vector_score'])
        expected_distance = calculate_l2_distance_normalized(query_vector, test_data[doc_key]['embedding'])
        env.assertAlmostEqual(returned_distance, expected_distance, delta=1e-6)


def test_hybrid_search_yield_score_as():
    """Test SEARCH with YIELD_SCORE_AS parameter"""
    env = Env()
    setup_basic_index(env)

    response = env.cmd('FT.HYBRID', 'idx', 'SEARCH', '*', 'YIELD_SCORE_AS', 'search_score',
                        'VSIM', '@embedding', '$BLOB',
                        'PARAMS', '2', 'BLOB', np.array([0.0, 0.0]).astype(np.float32).tobytes())
    results, _ = get_results_from_hybrid_response(response)

    # Validate the search_score field for all returned results
    env.assertGreater(len(results.keys()), 0)

    for doc_key in results:
        doc_result = results[doc_key]
        env.assertTrue('search_score' in doc_result)
        # Search score should be a valid float
        search_score = float(doc_result['search_score'])
        env.assertGreater(search_score, 0)


def test_hybrid_search_and_vsim_yield_parameters():
    """Test using SEARCH YIELD_SCORE_AS with VSIM YIELD_SCORE_AS together"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()

    response = env.cmd(
        'FT.HYBRID', 'idx',
        'SEARCH', '*',
            'YIELD_SCORE_AS', 'search_score',
        'VSIM', '@embedding', '$BLOB',
            'KNN', '2', 'K', '10',
        'YIELD_SCORE_AS', 'vector_distance',
        'PARAMS', '2', 'BLOB', query_vector)
    results, _ = get_results_from_hybrid_response(response)

    # Validate both search_score and vector_distance fields
    env.assertGreater(len(results.keys()), 0)

    for doc_key in results:
        doc_result = results[doc_key]
        # Should have either search_score or vector_distance (or both)
        has_search_score = 'search_score' in doc_result
        has_vector_distance = 'vector_distance' in doc_result
        env.assertTrue(has_search_score or has_vector_distance)

def test_hybrid_vsim_knn_both_yield_distance_and_score():
    """Test VSIM KNN with both YIELD_DISTANCE_AS and YIELD_SCORE_AS together -
    should fail because YIELD_DISTANCE_AS is not supported in VSIM"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()

    # YIELD_DISTANCE_AS is not supported in VSIM clauses and should return an error
    env.expect(
        'FT.HYBRID', 'idx', 'SEARCH', 'shoes', 'VSIM', '@embedding', '$BLOB',
        'KNN', '6', 'K', '10', 'YIELD_DISTANCE_AS', 'vector_distance',
        'YIELD_SCORE_AS', 'vector_score',
        'PARAMS', '2', 'BLOB', query_vector)\
            .error().contains('Unknown argument `YIELD_DISTANCE_AS` in KNN')

def test_hybrid_vsim_range_both_yield_distance_and_score():
    """Test VSIM RANGE with both YIELD_DISTANCE_AS and YIELD_SCORE_AS together -
    should fail because YIELD_DISTANCE_AS is not supported in VSIM"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()
    radius = 2

    # YIELD_DISTANCE_AS is not supported in VSIM clauses and should return an error
    env.expect(
        'FT.HYBRID', 'idx', 'SEARCH', 'shoes', 'VSIM', '@embedding', '$BLOB',
        'RANGE', '6', 'RADIUS', str(radius),
        'YIELD_DISTANCE_AS', 'vector_distance',
        'YIELD_SCORE_AS', 'vector_score',
        'PARAMS', '2', 'BLOB', query_vector)\
            .error().contains('Unknown argument `YIELD_DISTANCE_AS` in RANGE')


def test_hybrid_yield_score_as_after_combine_error():
    """Test that YIELD_SCORE_AS after COMBINE keyword fails"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()

    # This should fail because YIELD_SCORE_AS appears after COMBINE
    env.expect(
        'FT.HYBRID', 'idx', 'SEARCH', 'shoes', 'VSIM', '@embedding', '$BLOB',
        'KNN', '4', 'K', '10', 'COMBINE', 'RRF', '2', 'CONSTANT', '60',
        'YIELD_SCORE_AS', 'vector_distance',
        'PARAMS', '2', 'BLOB', query_vector)\
            .error().contains('Unknown argument `COMBINE` in KNN')

def test_hybrid_search_yield_score_as_after_combine():
    """Test that SEARCH YIELD_SCORE_AS after COMBINE keyword works"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()

    # YIELD_SCORE_AS after COMBINE should work
    response = env.cmd(
        'FT.HYBRID', 'idx', 'SEARCH', 'shoes', 'VSIM', '@embedding', '$BLOB',
        'COMBINE', 'RRF', '2', 'CONSTANT', '60',
            'YIELD_SCORE_AS', 'search_score',
        'PARAMS', '2', 'BLOB', query_vector)
    results, _ = get_results_from_hybrid_response(response)

    # Validate the search_score field
    env.assertGreater(len(results.keys()), 0)

    for doc_key in results:
        doc_result = results[doc_key]
        env.assertTrue('search_score' in doc_result)
        search_score = float(doc_result['search_score'])
        env.assertGreater(search_score, 0)

def test_hybrid_combine_yield_score_as_both_forms():
    """YIELD_SCORE_AS after COMBINE is accepted in two equivalent forms for
    backward compatibility: counted inside the method argument count (legacy),
    and positional after the method block (current). Both must succeed and
    produce the same results."""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()

    # Counted form: count 4 covers CONSTANT 60 YIELD_SCORE_AS search_score
    counted, _ = get_results_from_hybrid_response(env.cmd(
        'FT.HYBRID', 'idx', 'SEARCH', 'shoes', 'VSIM', '@embedding', '$BLOB',
        'COMBINE', 'RRF', '4', 'CONSTANT', '60', 'YIELD_SCORE_AS', 'search_score',
        'PARAMS', '2', 'BLOB', query_vector))

    # Positional form: count 2 covers CONSTANT 60, YIELD_SCORE_AS follows
    positional, _ = get_results_from_hybrid_response(env.cmd(
        'FT.HYBRID', 'idx', 'SEARCH', 'shoes', 'VSIM', '@embedding', '$BLOB',
        'COMBINE', 'RRF', '2', 'CONSTANT', '60', 'YIELD_SCORE_AS', 'search_score',
        'PARAMS', '2', 'BLOB', query_vector))

    env.assertGreater(len(counted.keys()), 0)
    for results in (counted, positional):
        for doc_key in results:
            env.assertTrue('search_score' in results[doc_key])
            env.assertGreater(float(results[doc_key]['search_score']), 0)
    env.assertEqual(counted, positional,
                    message="Counted and positional YIELD_SCORE_AS must be equivalent")

def test_hybrid_combine_duplicate_yield_score_as_error():
    """A duplicate YIELD_SCORE_AS after the COMBINE clause is rejected."""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()

    env.expect(
        'FT.HYBRID', 'idx',
        'SEARCH', 'shoes',
        'VSIM', '@embedding', '$BLOB',
        'COMBINE', 'RRF', '2', 'CONSTANT', '60',
            'YIELD_SCORE_AS', 'score1',
            'YIELD_SCORE_AS', 'score2',
        'PARAMS', '2', 'BLOB', query_vector)\
            .error().contains('YIELD_SCORE_AS: Unknown argument')

def test_hybrid_combine_missing_argument():
    """Test that missing argument value for YIELD_SCORE_AS after COMBINE clause
    results in an error"""
    env = Env()
    setup_basic_index(env)

    env.expect(
        'FT.HYBRID', 'idx',
        'SEARCH', 'shoes',
        'VSIM', '@embedding', '$BLOB',
        'COMBINE', 'RRF', '2', 'CONSTANT', '60',
            'YIELD_SCORE_AS')\
            .error().contains('Missing argument value for YIELD_SCORE_AS')

def test_hybrid_combine_without_fusion():
    """Test that COMBINE without a fusion method fails"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()

    env.expect(
        'FT.HYBRID', 'idx',
        'SEARCH', 'shoes',
        'VSIM', '@embedding', '$BLOB',
        'COMBINE',
            'YIELD_SCORE_AS', 'score1',
        'PARAMS', '2', 'BLOB', query_vector)\
            .error().contains('COMBINE: Invalid value for argument')

def test_hybrid_linear_combine_and_fused_score():
    """Test that COMBINE with LINEAR method and count of 0 is supported"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([1.2, 0.3]).astype(np.float32).tobytes()

    res = env.cmd(
        'FT.HYBRID', 'idx',
            'SEARCH', 'blue',
            'VSIM', '@embedding', '$BLOB',
            'COMBINE', 'LINEAR', '0',
            'YIELD_SCORE_AS', 'fuse_score',
        'PARAMS', '2', 'BLOB', query_vector)
    results, _ = get_results_from_hybrid_response(res)

    # Execute the same command with COMBINE with LINEAR method to verify that
    # the results are the same
    res_with_linear = env.cmd(
        'FT.HYBRID', 'idx',
            'SEARCH', 'blue',
            'VSIM', '@embedding', '$BLOB',
            'COMBINE', 'LINEAR', '4', 'ALPHA', '0.3', 'BETA', '0.7',
            'YIELD_SCORE_AS', 'fuse_score',
        'PARAMS', '2', 'BLOB', query_vector)
    results_with_linear, _ = get_results_from_hybrid_response(res_with_linear)

    env.assertEqual(results, results_with_linear,
                    message="Results with COMBINE LINEAR count 0 should match results with COMBINE LINEAR count 4")

def test_hybrid_rrf_combine_and_fused_score():
    """Test that COMBINE with RRF method and count of 0 is supported"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([1.2, 0.3]).astype(np.float32).tobytes()

    res = env.cmd(
        'FT.HYBRID', 'idx',
        'SEARCH', 'blue',
        'VSIM', '@embedding', '$BLOB',
        'COMBINE', 'RRF', '0',
        'YIELD_SCORE_AS', 'fused_score',
        'PARAMS', '2', 'BLOB', query_vector)
    results, _ = get_results_from_hybrid_response(res)

    # Execute the same command with COMBINE with RRF method to verify that
    # the results are the same
    res_with_rrf = env.cmd(
        'FT.HYBRID', 'idx',
        'SEARCH', 'blue',
        'VSIM', '@embedding', '$BLOB',
        'COMBINE', 'RRF', '2', 'CONSTANT', '60',
        'YIELD_SCORE_AS', 'fused_score',
        'PARAMS', '2', 'BLOB', query_vector)
    results_with_rrf, _ = get_results_from_hybrid_response(res_with_rrf)

    env.assertEqual(results, results_with_rrf,
                    message="Results with COMBINE RRF count 0 should match results with COMBINE RRF constant 60")

def test_hybrid_multiple_yield_after_combine_error():
    """Test that multiple YIELD parameters after COMBINE keyword fail"""
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()

    # This should fail because both YIELD parameters appear after COMBINE
    env.expect(
        'FT.HYBRID', 'idx', 'SEARCH', 'shoes', 'VSIM', '@embedding', '$BLOB',
        'KNN', '4', 'K', '10', 'COMBINE', 'LINEAR', '8', 'ALPHA', '0.5', 'BETA', '0.5',
        'YIELD_SCORE_AS', 'vector_distance', 'YIELD_SCORE_AS', 'vector_score',
        'PARAMS', '2', 'BLOB', query_vector)\
            .error().contains('Unknown argument `COMBINE` in KNN')

def test_hybrid_yield_score_as_all_possible_scores():
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()
    alpha = 0.3
    beta = 0.7

    response = env.cmd(
        'FT.HYBRID', 'idx',
        'SEARCH', 'shoes',
            'YIELD_SCORE_AS', 's_score',
        'VSIM', '@embedding', '$BLOB',
            'KNN', '2', 'K', '10',
            'YIELD_SCORE_AS', 'v_score',
        'COMBINE', 'LINEAR', '4', 'ALPHA', alpha, 'BETA', beta,
            'YIELD_SCORE_AS', 'fused_score',
        'APPLY', f"{alpha}*case(exists(@s_score), @s_score ,0) + {beta}*case(exists(@v_score), @v_score,0)", 'AS', 'calculated_score',
        'PARAMS', '2', 'BLOB', query_vector)
    results, _ = get_results_from_hybrid_response(response)

    # Validate the vector_distance and vector_score fields
    env.assertGreater(len(results.keys()), 0)
    for doc_key, doc_result in results.items():
        # assert at least one subquery score is present
        env.assertTrue('s_score' in doc_result or 'v_score' in doc_result)

        # assert both calculated and fused scores are present
        env.assertTrue('calculated_score' in doc_result)
        env.assertTrue('fused_score' in doc_result)

        # assert fused_score and the score calculated from the subquery scores using apply are the same
        calculated_score = float(doc_result[f'calculated_score'])
        fused_score = float(doc_result[f'fused_score'])
        env.assertGreater(fused_score, 0)
        env.assertAlmostEqual(calculated_score, fused_score, delta=1e-6, message=f"Fused score and calculated score for {doc_key} do not match")

def test_vsim_yield_score_as_with_filter():
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()
    response = env.cmd(
        'FT.HYBRID', 'idx',
        'SEARCH', 'shoes',
            'YIELD_SCORE_AS', 's_score',
        'VSIM', '@embedding', '$BLOB',
            'KNN', '2', 'K', '10',
            'FILTER', '@description:blue',
            'YIELD_SCORE_AS', 'v_score',
        'PARAMS', '2', 'BLOB', query_vector)
    results, _ = get_results_from_hybrid_response(response)
    # 3 results are returned:
    # - 3 containing "shoes" -> doc:1, doc:2, doc:4 -> s_score is present
    # - 1 containing "blue"  -> doc:4 -> v_score is present
    env.assertEqual(len(results.keys()), 3)
    for doc_key, doc_result in results.items():
        if doc_key in ["doc:1", "doc:2"]:
            env.assertTrue('s_score' in doc_result)
            env.assertFalse('v_score' in doc_result)
        if doc_key == "doc:4":
            env.assertTrue('s_score' in doc_result)
            env.assertTrue('v_score' in doc_result)

def test_vsim_yield_score_as_with_filter_and_post_filter():
    env = Env()
    setup_basic_index(env)
    query_vector = np.array([0.0, 0.0]).astype(np.float32).tobytes()
    response = env.cmd(
        'FT.HYBRID', 'idx',
        'SEARCH', 'shoes',
            'YIELD_SCORE_AS', 's_score',
        'VSIM', '@embedding', '$BLOB',
            'KNN', '2', 'K', '10',
            'FILTER', '@description:blue',
            'YIELD_SCORE_AS', 'v_score',
        'FILTER', '@__key=="doc:4"',
        'PARAMS', '2', 'BLOB', query_vector)
    results, _ = get_results_from_hybrid_response(response)
    # a single result is returned, due to post-filter:
    # - doc:4 -> v_score is present
    env.assertEqual(len(results.keys()), 1)
    for doc_key, doc_result in results.items():
        if doc_key == "doc:4":
            env.assertTrue('s_score' in doc_result)
            env.assertTrue('v_score' in doc_result)