"""Hardcoded demo responses for specific functions."""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

from aiservice.common_utils import should_hack_for_demo

from .validate import instrument_tests

if TYPE_CHECKING:
    from core.shared.testgen_models import TestGenResponseSchema, TestGenSchema


async def try_demo_testgen(data: TestGenSchema, python_version: tuple[int, int, int]) -> TestGenResponseSchema | None:
    """Return a canned demo response if source matches a demo function, else None."""
    if not should_hack_for_demo(data.source_code_being_tested):
        return None
    if "find_common_tags" in data.source_code_being_tested:
        return await hack_for_demo(data, python_version)
    if "weighted_sum" in data.source_code_being_tested:
        return await hack_for_demo_gsq(data, python_version)
    return None


async def hack_for_demo(data: TestGenSchema, python_version: tuple[int, int, int]) -> TestGenResponseSchema:
    from core.shared.testgen_models import TestGenResponseSchema

    if data.test_index == 0:
        generated_test_source = '# imports\n# function to test\nfrom __future__ import annotations\n\nimport pytest  # used for our unit tests\nfrom codeflash.result.common_tags import find_common_tags\n\n# unit tests\n\ndef test_single_article():\n    # Single article should return its tags\n    articles = [{"tags": ["python", "coding", "tutorial"]}]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_multiple_articles_with_common_tags():\n    # Multiple articles with common tags should return the common tags\n    articles = [\n        {"tags": ["python", "coding"]},\n        {"tags": ["python", "data"]},\n        {"tags": ["python", "machine learning"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_empty_list_of_articles():\n    # Empty list of articles should return an empty set\n    articles = []\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_articles_with_no_common_tags():\n    # Articles with no common tags should return an empty set\n    articles = [\n        {"tags": ["python"]},\n        {"tags": ["java"]},\n        {"tags": ["c++"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_articles_with_empty_tag_lists():\n    # Articles with some empty tag lists should return an empty set\n    articles = [\n        {"tags": []},\n        {"tags": ["python"]},\n        {"tags": ["python", "java"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_all_articles_with_empty_tag_lists():\n    # All articles with empty tag lists should return an empty set\n    articles = [\n        {"tags": []},\n        {"tags": []},\n        {"tags": []}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_tags_with_special_characters():\n    # Tags with special characters should be handled correctly\n    articles = [\n        {"tags": ["python!", "coding"]},\n        {"tags": ["python!", "data"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_case_sensitivity():\n    # Tags with different cases should not be considered the same\n    articles = [\n        {"tags": ["Python", "coding"]},\n        {"tags": ["python", "data"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_large_number_of_articles():\n    # Large number of articles with a common tag should return that tag\n    articles = [{"tags": ["common_tag", f"tag{i}"]} for i in range(1000)]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_large_number_of_tags():\n    # Large number of tags with some common tags should return the common tags\n    articles = [\n        {"tags": [f"tag{i}" for i in range(1000)]},\n        {"tags": [f"tag{i}" for i in range(500, 1500)]}\n    ]\n    expected = {f"tag{i}" for i in range(500, 1000)}\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_mixed_length_of_tag_lists():\n    # Articles with mixed length of tag lists should return the common tags\n    articles = [\n        {"tags": ["python", "coding"]},\n        {"tags": ["python"]},\n        {"tags": ["python", "coding", "tutorial"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_tags_with_different_data_types():\n    # Tags with different data types should only consider strings\n    articles = [\n        {"tags": ["python", 123]},\n        {"tags": ["python", "123"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_performance_with_large_data():\n    # Performance with large data should return the common tag\n    articles = [{"tags": ["common_tag", f"tag{i}"]} for i in range(10000)]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_scalability_with_increasing_tags():\n    # Scalability with increasing tags should return the common tag\n    articles = [{"tags": ["common_tag"] + [f"tag{i}" for i in range(j)]} for j in range(1, 1001)]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation'
    else:
        generated_test_source = '# imports\n# function to test\nfrom __future__ import annotations\n\nimport pytest  # used for our unit tests\nfrom codeflash.result.common_tags import find_common_tags\n\n# unit tests\n\ndef test_empty_input_list():\n    # Test with an empty list\n    codeflash_output = find_common_tags([])\n    # Outputs were verified to be equal to the original implementation\n\ndef test_single_article():\n    # Test with a single article with tags\n    codeflash_output = find_common_tags([{"tags": ["python", "coding", "development"]}])\n    # Test with a single article with no tags\n    codeflash_output = find_common_tags([{"tags": []}])\n    # Outputs were verified to be equal to the original implementation\n\ndef test_multiple_articles_some_common_tags():\n    # Test with multiple articles having some common tags\n    articles = [\n        {"tags": ["python", "coding", "development"]},\n        {"tags": ["python", "development", "tutorial"]},\n        {"tags": ["python", "development", "guide"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n\n    articles = [\n        {"tags": ["tech", "news"]},\n        {"tags": ["tech", "gadgets"]},\n        {"tags": ["tech", "reviews"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_multiple_articles_no_common_tags():\n    # Test with multiple articles having no common tags\n    articles = [\n        {"tags": ["python", "coding"]},\n        {"tags": ["development", "tutorial"]},\n        {"tags": ["guide", "learning"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n\n    articles = [\n        {"tags": ["apple", "banana"]},\n        {"tags": ["orange", "grape"]},\n        {"tags": ["melon", "kiwi"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_articles_with_duplicate_tags():\n    # Test with articles having duplicate tags\n    articles = [\n        {"tags": ["python", "python", "coding"]},\n        {"tags": ["python", "development", "python"]},\n        {"tags": ["python", "guide", "python"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n\n    articles = [\n        {"tags": ["tech", "tech", "news"]},\n        {"tags": ["tech", "tech", "gadgets"]},\n        {"tags": ["tech", "tech", "reviews"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_articles_with_mixed_case_tags():\n    # Test with articles having mixed case tags\n    articles = [\n        {"tags": ["Python", "Coding"]},\n        {"tags": ["python", "Development"]},\n        {"tags": ["PYTHON", "Guide"]}\n    ]\n    codeflash_output = find_common_tags(articles)  # Assuming case sensitivity\n\n    articles = [\n        {"tags": ["Tech", "News"]},\n        {"tags": ["tech", "Gadgets"]},\n        {"tags": ["TECH", "Reviews"]}\n    ]\n    codeflash_output = find_common_tags(articles)  # Assuming case sensitivity\n    # Outputs were verified to be equal to the original implementation\n\ndef test_articles_with_non_string_tags():\n    # Test with articles having non-string tags\n    articles = [\n        {"tags": ["python", 123, "coding"]},\n        {"tags": ["python", "development", 123]},\n        {"tags": ["python", "guide", 123]}\n    ]\n    codeflash_output = find_common_tags(articles)\n\n    articles = [\n        {"tags": [None, "news"]},\n        {"tags": ["tech", None]},\n        {"tags": [None, "reviews"]}\n    ]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation\n\ndef test_large_scale_test_cases():\n    # Test with large scale input where all tags should be common\n    articles = [\n        {"tags": ["tag" + str(i) for i in range(1000)]} for _ in range(100)\n    ]\n    expected_output = {"tag" + str(i) for i in range(1000)}\n    codeflash_output = find_common_tags(articles)\n\n    # Test with large scale input where no tags should be common\n    articles = [\n        {"tags": ["tag" + str(i) for i in range(1000)]} for _ in range(50)\n    ] + [{"tags": ["unique_tag"]}]\n    codeflash_output = find_common_tags(articles)\n    # Outputs were verified to be equal to the original implementation'

    instrumented_behavior_tests, instrumented_perf_tests = instrument_tests(generated_test_source, data, python_version)
    await asyncio.sleep(5)
    return TestGenResponseSchema(
        generated_tests=generated_test_source,
        instrumented_behavior_tests=instrumented_behavior_tests or "",
        instrumented_perf_tests=instrumented_perf_tests or "",
    )


async def hack_for_demo_gsq(data: TestGenSchema, python_version: tuple[int, int, int]) -> TestGenResponseSchema:
    from core.shared.testgen_models import TestGenResponseSchema

    if data.test_index == 0:
        generated_test_source = "from datetime import datetime, timedelta\n# function to test\nfrom functools import reduce\n\nimport numpy as np\nimport pandas as pd\n# imports\nimport pytest  # used for our unit tests\nfrom gs_quant.timeseries.algebra import weighted_sum\n\n\nclass MqTypeError(TypeError): pass\nclass MqValueError(ValueError): pass\nfrom gs_quant.timeseries.algebra import weighted_sum\n\n# unit tests\n\n# --- Basic Test Cases ---\n\ndef test_weighted_sum_simple_two_series():\n    # Simple case: two series, same index, weights sum to 1\n    idx = pd.date_range('2024-01-01', periods=3)\n    s1 = pd.Series([1, 2, 3], index=idx)\n    s2 = pd.Series([4, 5, 6], index=idx)\n    weights = [0.7, 0.3]\n    codeflash_output = weighted_sum([s1, s2], weights); result = codeflash_output\n    expected = (s1 * 0.7 + s2 * 0.3) / (0.7 + 0.3)\n\ndef test_weighted_sum_three_series_equal_weights():\n    # Three series, equal weights\n    idx = pd.date_range('2024-01-01', periods=3)\n    s1 = pd.Series([1, 2, 3], index=idx)\n    s2 = pd.Series([4, 5, 6], index=idx)\n    s3 = pd.Series([7, 8, 9], index=idx)\n    weights = [1, 1, 1]\n    codeflash_output = weighted_sum([s1, s2, s3], weights); result = codeflash_output\n    expected = (s1 + s2 + s3) / 3\n\ndef test_weighted_sum_weights_not_normalized():\n    # Weights do not sum to 1\n    idx = pd.date_range('2024-01-01', periods=2)\n    s1 = pd.Series([10, 20], index=idx)\n    s2 = pd.Series([30, 40], index=idx)\n    weights = [2, 3]\n    codeflash_output = weighted_sum([s1, s2], weights); result = codeflash_output\n    expected = (s1 * 2 + s2 * 3) / (2 + 3)\n\n# --- Edge Test Cases ---\n\ndef test_weighted_sum_empty_series_list():\n    # Empty input series list\n    with pytest.raises(MqValueError):\n        weighted_sum([], [1, 2])\n\ndef test_weighted_sum_empty_weights_list():\n    # Empty weights list\n    idx = pd.date_range('2024-01-01', periods=2)\n    s1 = pd.Series([1, 2], index=idx)\n    with pytest.raises(MqValueError):\n        weighted_sum([s1], [])\n\ndef test_weighted_sum_weights_length_mismatch():\n    # Weights and series length mismatch\n    idx = pd.date_range('2024-01-01', periods=2)\n    s1 = pd.Series([1, 2], index=idx)\n    s2 = pd.Series([3, 4], index=idx)\n    with pytest.raises(MqValueError):\n        weighted_sum([s1, s2], [1])\n\ndef test_weighted_sum_non_series_input():\n    # Non-pandas Series input\n    idx = pd.date_range('2024-01-01', periods=2)\n    s1 = pd.Series([1, 2], index=idx)\n    with pytest.raises(MqTypeError):\n        weighted_sum([s1, [3, 4]], [0.5, 0.5])\n\ndef test_weighted_sum_non_numeric_weights():\n    # Non-numeric weights\n    idx = pd.date_range('2024-01-01', periods=2)\n    s1 = pd.Series([1, 2], index=idx)\n    s2 = pd.Series([3, 4], index=idx)\n    with pytest.raises(MqTypeError):\n        weighted_sum([s1, s2], [0.5, \"a\"])\n\ndef test_weighted_sum_disjoint_calendars():\n    # Series with disjoint calendars (no overlap)\n    idx1 = pd.date_range('2024-01-01', periods=2)\n    idx2 = pd.date_range('2024-02-01', periods=2)\n    s1 = pd.Series([1, 2], index=idx1)\n    s2 = pd.Series([3, 4], index=idx2)\n    weights = [0.5, 0.5]\n    codeflash_output = weighted_sum([s1, s2], weights); result = codeflash_output\n\ndef test_weighted_sum_partial_overlap_calendars():\n    # Series with partial overlap in calendar\n    idx1 = pd.date_range('2024-01-01', periods=3)\n    idx2 = pd.date_range('2024-01-02', periods=3)\n    s1 = pd.Series([1, 2, 3], index=idx1)\n    s2 = pd.Series([4, 5, 6], index=idx2)\n    weights = [1, 1]\n    codeflash_output = weighted_sum([s1, s2], weights); result = codeflash_output\n    # Only dates present in both should be included\n    expected_idx = idx1.intersection(idx2)\n    expected = (s1.reindex(expected_idx) + s2.reindex(expected_idx)) / 2\n\ndef test_weighted_sum_zero_weights():\n    # Zero weights (should return NaN for all dates)\n    idx = pd.date_range('2024-01-01', periods=2)\n    s1 = pd.Series([1, 2], index=idx)\n    s2 = pd.Series([3, 4], index=idx)\n    weights = [0, 0]\n    codeflash_output = weighted_sum([s1, s2], weights); result = codeflash_output\n\ndef test_weighted_sum_negative_weights():\n    # Negative weights\n    idx = pd.date_range('2024-01-01', periods=2)\n    s1 = pd.Series([10, 20], index=idx)\n    s2 = pd.Series([30, 40], index=idx)\n    weights = [1, -1]\n    codeflash_output = weighted_sum([s1, s2], weights); result = codeflash_output\n    expected = (s1 * 1 + s2 * -1) / (1 + -1)\n\ndef test_weighted_sum_nan_values_in_series():\n    # Series containing NaN values\n    idx = pd.date_range('2024-01-01', periods=3)\n    s1 = pd.Series([1, None, 3], index=idx)\n    s2 = pd.Series([4, 5, None], index=idx)\n    weights = [1, 1]\n    codeflash_output = weighted_sum([s1, s2], weights); result = codeflash_output\n    # NaN propagation: result should be NaN where any input is NaN\n    expected = (s1 * 1 + s2 * 1) / 2\n\ndef test_weighted_sum_mixed_int_float_weights():\n    # Mixed int and float weights\n    idx = pd.date_range('2024-01-01', periods=2)\n    s1 = pd.Series([1, 2], index=idx)\n    s2 = pd.Series([3, 4], index=idx)\n    weights = [1, 0.5]\n    codeflash_output = weighted_sum([s1, s2], weights); result = codeflash_output\n    expected = (s1 * 1 + s2 * 0.5) / (1 + 0.5)\n\n# --- Large Scale Test Cases ---\n\ndef test_weighted_sum_large_number_of_series():\n    # Large number of series (up to 1000)\n    n = 1000\n    idx = pd.date_range('2024-01-01', periods=10)\n    series_list = [pd.Series([i]*10, index=idx) for i in range(n)]\n    weights = [1 for _ in range(n)]\n    codeflash_output = weighted_sum(series_list, weights); result = codeflash_output\n    # Each date should be the average of 0...999, i.e. (sum(range(n)))/n\n    expected_value = sum(range(n)) / n\n\ndef test_weighted_sum_large_series_length():\n    # Large series length (up to 1000 dates)\n    idx = pd.date_range('2024-01-01', periods=1000)\n    s1 = pd.Series(range(1000), index=idx)\n    s2 = pd.Series(range(1000, 0, -1), index=idx)\n    weights = [0.25, 0.75]\n    codeflash_output = weighted_sum([s1, s2], weights); result = codeflash_output\n    expected = (s1 * 0.25 + s2 * 0.75) / (0.25 + 0.75)\n\ndef test_weighted_sum_large_partial_overlap():\n    # Large series with partial overlap\n    idx1 = pd.date_range('2024-01-01', periods=1000)\n    idx2 = pd.date_range('2024-01-500', periods=500)\n    s1 = pd.Series(range(1000), index=idx1)\n    s2 = pd.Series(range(500), index=idx2)\n    weights = [1, 2]\n    codeflash_output = weighted_sum([s1, s2], weights); result = codeflash_output\n    expected_idx = idx1.intersection(idx2)\n    expected = (s1.reindex(expected_idx) * 1 + s2.reindex(expected_idx) * 2) / (1 + 2)\n# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."
    else:
        generated_test_source = "from datetime import datetime, timedelta\n\nimport pandas as pd\n# imports\nimport pytest  # used for our unit tests\nfrom gs_quant.timeseries.algebra import weighted_sum\n\n# ---------- Unit Tests ----------\n\n# Helper to create pd.Series with consecutive dates\ndef make_series(start, count, value_func=lambda i: i):\n    dates = [datetime(2020, 1, 1) + timedelta(days=i) for i in range(count)]\n    values = [value_func(i) for i in range(count)]\n    return pd.Series(values, index=dates)\n\n# 1. Basic Test Cases\n\ndef test_weighted_sum_simple_two_series_equal_weights():\n    # Two series, equal weights\n    s1 = make_series(0, 5, lambda i: 2)\n    s2 = make_series(0, 5, lambda i: 4)\n    codeflash_output = weighted_sum([s1, s2], [1, 1]); result = codeflash_output\n    expected = pd.Series([3.0]*5, index=s1.index)\n\ndef test_weighted_sum_simple_two_series_unequal_weights():\n    # Two series, weights 2:1\n    s1 = make_series(0, 5, lambda i: 2)\n    s2 = make_series(0, 5, lambda i: 8)\n    codeflash_output = weighted_sum([s1, s2], [2, 1]); result = codeflash_output\n    expected = pd.Series([(2*2+1*8)/3]*5, index=s1.index)\n\ndef test_weighted_sum_three_series_varied_weights():\n    # Three series, varied weights\n    s1 = make_series(0, 3, lambda i: 1)\n    s2 = make_series(0, 3, lambda i: 3)\n    s3 = make_series(0, 3, lambda i: 5)\n    codeflash_output = weighted_sum([s1, s2, s3], [0.2, 0.3, 0.5]); result = codeflash_output\n    expected = pd.Series([1*0.2+3*0.3+5*0.5], index=s1.index)\n    expected = pd.Series([1*0.2+3*0.3+5*0.5]*3, index=s1.index)\n\ndef test_weighted_sum_one_series():\n    # Single series, should return the series itself\n    s1 = make_series(0, 4, lambda i: i+5)\n    codeflash_output = weighted_sum([s1], [1]); result = codeflash_output\n\n# 2. Edge Test Cases\n\ndef test_weighted_sum_empty_series_list():\n    # Empty series list should raise ValueError\n    with pytest.raises(ValueError):\n        weighted_sum([], [])\n\ndef test_weighted_sum_weights_length_mismatch():\n    # Series and weights length mismatch\n    s1 = make_series(0, 3)\n    s2 = make_series(0, 3)\n    with pytest.raises(ValueError):\n        weighted_sum([s1, s2], [1])\n\ndef test_weighted_sum_non_series_input():\n    # Non-Series in series list\n    s1 = make_series(0, 3)\n    with pytest.raises(TypeError):\n        weighted_sum([s1, [1, 2, 3]], [1, 1])\n\ndef test_weighted_sum_non_numeric_weights():\n    # Non-numeric in weights\n    s1 = make_series(0, 3)\n    s2 = make_series(0, 3)\n    with pytest.raises(TypeError):\n        weighted_sum([s1, s2], [1, 'a'])\n\ndef test_weighted_sum_zero_weights():\n    # All weights zero should return NaN series\n    s1 = make_series(0, 3, lambda i: 1)\n    s2 = make_series(0, 3, lambda i: 2)\n    codeflash_output = weighted_sum([s1, s2], [0, 0]); result = codeflash_output\n\ndef test_weighted_sum_negative_weights():\n    # Negative weights\n    s1 = make_series(0, 3, lambda i: 2)\n    s2 = make_series(0, 3, lambda i: 4)\n    codeflash_output = weighted_sum([s1, s2], [-1, 2]); result = codeflash_output\n    expected = pd.Series([(2*-1+4*2)/(-1+2)]*3, index=s1.index)\n\ndef test_weighted_sum_partial_overlap_indices():\n    # Series with partially overlapping indices\n    s1 = make_series(0, 5, lambda i: i)\n    s2 = make_series(2, 5, lambda i: i+10)\n    # Overlap is on days 2,3,4\n    overlap_dates = s1.index.intersection(s2.index)\n    codeflash_output = weighted_sum([s1, s2], [1, 1]); result = codeflash_output\n    expected = pd.Series([(s1[d]+s2[d])/2 for d in overlap_dates], index=overlap_dates)\n\ndef test_weighted_sum_series_with_nans():\n    # Series with NaNs, intersection should drop NaN dates\n    s1 = make_series(0, 4, lambda i: float('nan') if i==2 else i)\n    s2 = make_series(0, 4, lambda i: i)\n    # Only dates where both have non-NaN should be used\n    codeflash_output = weighted_sum([s1, s2], [1, 1]); result = codeflash_output\n\ndef test_weighted_sum_series_with_different_freq():\n    # Series with different frequencies\n    s1 = make_series(0, 5, lambda i: i)\n    s2 = pd.Series([10, 20], index=[datetime(2020, 1, 2), datetime(2020, 1, 4)])\n    # Only dates 2020-01-02 and 2020-01-04 are in both\n    overlap_dates = s1.index.intersection(s2.index)\n    codeflash_output = weighted_sum([s1, s2], [1, 3]); result = codeflash_output\n    expected = pd.Series([(s1[d]*1+s2[d]*3)/4 for d in overlap_dates], index=overlap_dates)\n\ndef test_weighted_sum_all_nan_series():\n    # All series are NaN\n    dates = [datetime(2020, 1, 1) + timedelta(days=i) for i in range(3)]\n    s1 = pd.Series([float('nan')] * 3, index=dates)\n    s2 = pd.Series([float('nan')] * 3, index=dates)\n    codeflash_output = weighted_sum([s1, s2], [1, 1]); result = codeflash_output\n\n# 3. Large Scale Test Cases\n\ndef test_weighted_sum_large_number_of_series():\n    # 100 series, each with 100 points\n    N = 100\n    M = 100\n    series_list = [make_series(0, M, lambda i: i+j) for j in range(N)]\n    weights = [j+1 for j in range(N)]\n    codeflash_output = weighted_sum(series_list, weights); result = codeflash_output\n    # For each date, expected value is weighted sum of values across series\n    dates = series_list[0].index\n    expected_values = []\n    total_weight = sum(weights)\n    for i in range(M):\n        # For date i, value in series j is i+j\n        weighted_sum_val = sum((i+j)*weights[j] for j in range(N)) / total_weight\n        expected_values.append(weighted_sum_val)\n    expected = pd.Series(expected_values, index=dates)\n\ndef test_weighted_sum_large_series_length():\n    # Two series, each with 1000 points\n    M = 1000\n    s1 = make_series(0, M, lambda i: i)\n    s2 = make_series(0, M, lambda i: 2*i)\n    codeflash_output = weighted_sum([s1, s2], [1, 2]); result = codeflash_output\n    expected = pd.Series([(i*1+2*i*2)/3 for i in range(M)], index=s1.index)\n\ndef test_weighted_sum_large_sparse_overlap():\n    # Series with sparse overlap\n    s1 = pd.Series([1]*1000, index=[datetime(2020, 1, 1)+timedelta(days=i) for i in range(0, 1000, 2)])\n    s2 = pd.Series([2]*500, index=[datetime(2020, 1, 1)+timedelta(days=i) for i in range(1, 1000, 2)])\n    # There should be no overlap, so result should be empty\n    codeflash_output = weighted_sum([s1, s2], [1, 1]); result = codeflash_output\n\ndef test_weighted_sum_large_all_nan():\n    # Large series, all NaN\n    M = 1000\n    dates = [datetime(2020, 1, 1) + timedelta(days=i) for i in range(M)]\n    s1 = pd.Series([float('nan')]*M, index=dates)\n    s2 = pd.Series([float('nan')]*M, index=dates)\n    codeflash_output = weighted_sum([s1, s2], [1, 2]); result = codeflash_output\n# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."
    instrumented_behavior_tests, instrumented_perf_tests = instrument_tests(generated_test_source, data, python_version)
    await asyncio.sleep(2)
    return TestGenResponseSchema(
        generated_tests=generated_test_source,
        instrumented_behavior_tests=instrumented_behavior_tests or "",
        instrumented_perf_tests=instrumented_perf_tests or "",
    )
