import random from app.schemas.grade_1.sum_with_image_reference import ( AddendGroup, FigureGroup, ImageOption, PictureAsset, SumWithImageReferenceProblem, ) def sum_with_image_reference( first_quantity: int, second_quantity: int, first_description: str, second_description: str, common_object: str, first_picture: dict, second_picture: dict, option_count: int = 3, min_option_quantity: int = 1, max_option_quantity: int = 20, seed: int | None = None, ) -> dict: """Generate an addition question with image answer options.""" if option_count < 2: raise ValueError("option_count must be at least 2") total = first_quantity + second_quantity option_values = list(range(max(2, min_option_quantity), max_option_quantity + 1)) if total not in option_values: raise ValueError("option quantity range must include the correct total") distractor_values = [value for value in option_values if value != total] if len(distractor_values) < option_count - 1: raise ValueError("option quantity range must contain enough distractor values") rng = random.Random(seed) selected_values = rng.sample(distractor_values, option_count - 1) + [total] rng.shuffle(selected_values) first_option_picture = PictureAsset.model_validate(first_picture) second_option_picture = PictureAsset.model_validate(second_picture) question = ( f"Hay {first_quantity} {first_description} y " f"{second_quantity} {second_description}, " f"¿cuántas {common_object} hay en total?" ) options: list[ImageOption] = [] for index, quantity in enumerate(selected_values): if quantity == total: option_first_quantity = first_quantity option_second_quantity = second_quantity else: option_first_quantity = rng.randint(1, quantity - 1) option_second_quantity = quantity - option_first_quantity options.append( ImageOption( position=index + 1, quantity=quantity, first_group=FigureGroup( quantity=option_first_quantity, description=first_description, picture=first_option_picture, ), second_group=FigureGroup( quantity=option_second_quantity, description=second_description, picture=second_option_picture, ), is_correct=quantity == total, ) ) problem = SumWithImageReferenceProblem( question=question, first_group=AddendGroup( quantity=first_quantity, description=first_description, ), second_group=AddendGroup( quantity=second_quantity, description=second_description, ), common_object=common_object, total=total, options=options, ) return problem.model_dump()