"""Number filtering task generator"""

from dataclasses import dataclass
from random import Random
from typing import Optional

from ..coaching import BaseCurriculum, RangeAttributeDefinition, ScalarAttributeDefinition
from ..factory import ProceduralDataset, register_dataset

DATASET_NAME = "number_filtering"


@dataclass
class NumberFilteringConfig:
    """Configuration for number filtering task generation"""

    min_numbers: int = 3  # Minimum numbers in list
    max_numbers: int = 10  # Maximum numbers in list
    min_decimals: int = 0  # Minimum decimal places
    max_decimals: int = 4  # Maximum decimal places
    min_value: float = -100.0  # Minimum number value
    max_value: float = 100.0  # Maximum number value
    seed: Optional[int] = None
    size: int = 500  # Virtual dataset size

    def validate(self) -> None:
        """Validate configuration parameters"""
        assert self.min_numbers > 0, "min_numbers must be positive"
        assert self.max_numbers >= self.min_numbers, "max_numbers must be >= min_numbers"
        assert self.min_decimals >= 0, "min_decimals must be non-negative"
        assert self.max_decimals >= self.min_decimals, "max_decimals must be >= min_decimals"
        assert self.max_value > self.min_value, "max_value must be > min_value"


class NumberFilteringDataset(ProceduralDataset):
    """Generates number filtering tasks"""

    def __init__(self, config: NumberFilteringConfig):
        super().__init__(config=config, seed=config.seed, size=config.size)

    def _format_number(self, num: float, decimals: int) -> str:
        """Format a number with specified decimal places"""
        return f"{num:.{decimals}f}"

    def _generate_numbers(self, rng: Random) -> tuple[list[float], list[str]]:
        """Generate list of numbers and their string representations"""
        count = rng.randint(self.config.min_numbers, self.config.max_numbers)
        numbers = []
        str_numbers = []

        for _ in range(count):
            num = rng.uniform(self.config.min_value, self.config.max_value)
            decimals = rng.randint(self.config.min_decimals, self.config.max_decimals)
            str_num = self._format_number(num, decimals)
            numbers.append(float(str_num))  # Convert back to simulate precision loss
            str_numbers.append(str_num)

        return numbers, str_numbers

    def __getitem__(self, idx: int) -> dict:
        """Generate a single number filtering task"""
        rng = Random(self.seed + idx)

        # Generate numbers and their string representations
        numbers, str_numbers = self._generate_numbers(rng)

        # Determine filter value between min and max of generated numbers
        min_val = min(numbers)
        max_val = max(numbers)
        filter_value = rng.uniform(min_val, max_val)
        decimals = rng.randint(self.config.min_decimals, self.config.max_decimals)
        filter_str = self._format_number(filter_value, decimals)
        filter_value = float(filter_str)  # Convert back to simulate precision loss

        # Randomly choose filter operation
        keep_larger = rng.choice([True, False])
        larger_smaller = "larger" if keep_larger else "smaller"
        keep_remove = "keep" if rng.choice([True, False]) else "remove"

        # Apply filter based on chosen operation
        if keep_remove == "keep":
            result = [n for n in numbers if (n > filter_value if keep_larger else n < filter_value)]
        else:  # remove
            result = [n for n in numbers if (n <= filter_value if keep_larger else n >= filter_value)]

        # Format results as strings with original precision
        result_strs = [str_numbers[numbers.index(n)] for n in result]

        return {
            "question": (
                f"{keep_remove.capitalize()} all numbers {larger_smaller} than {filter_str} "
                f"in this list: {str_numbers}\nReturn the new list in the same format."
            ),
            "answer": str(result_strs) if result_strs else "[]",
            "metadata": {
                "source_dataset": DATASET_NAME,
                "source_index": idx,
                "original_numbers": str_numbers,
                "filter_value": filter_str,
                "operation": f"{keep_remove}_{larger_smaller}",
                "result": result_strs,
                "numbers": len(numbers),
                "difficulty": {
                    "numbers": (self.config.min_numbers, self.config.max_numbers),
                    "decimals": (self.config.min_decimals, self.config.max_decimals),
                    "value": (self.config.min_value, self.config.max_value),
                },
            },
        }


class NumberFilteringCurriculum(BaseCurriculum):
    def __init__(self):
        super().__init__(NumberFilteringCurriculum.__name__, NumberFilteringConfig)

        # Define attributes
        self._define_attributes(
            RangeAttributeDefinition(
                name="numbers",
                levels=[10, 50, 100, 200],
                description="How many numbers to sort",
                lower_field_name="min_numbers",
                upper_field_name="max_numbers",
                ensure_interval=True,
            ),
            RangeAttributeDefinition(
                name="decimals",
                levels=[0, 2, 4, 6],
                description="Number of decimal places",
                lower_field_name="min_decimals",
                upper_field_name="max_decimals",
                ensure_interval=True,
            ),
            ScalarAttributeDefinition(
                name="min_value",
                field_name="min_value",
                levels=[-100, -500, -1000, -10000],
                description="Minimum number value",
            ),
            ScalarAttributeDefinition(
                name="max_value",
                field_name="max_value",
                levels=[100, 500, 1000, 10000],
                description="Maximum number value",
            ),
        )


register_dataset(DATASET_NAME, NumberFilteringDataset, NumberFilteringConfig, NumberFilteringCurriculum)
