import sympy as sp
import random
from typing import Dict, Any, Tuple

class TaskGenerationEngine:
    def __init__(self):
        self.x = sp.Symbol('x')
        # Components for generating random functions F(x)
        self.basic_functions = [
            lambda x, c: x**c,
            lambda x, c: sp.sin(c*x),
            lambda x, c: sp.cos(c*x),
            lambda x, c: sp.exp(c*x),
            lambda x, c: sp.ln(sp.Abs(c*x))
        ]

    def _score_difficulty(self, components: int, nesting: int) -> float:
        """D = num_components + degree_of_nesting * 2"""
        return float(components + nesting * 2.0)

    def generate_random_function(self, complexity: int) -> Tuple[Any, float]:
        """Generates a random F(x)."""
        num_components = max(1, int(complexity / 2))
        nesting = max(0, int(complexity / 4))
        
        f_expr = 0
        for _ in range(num_components):
            comp_func = random.choice(self.basic_functions)
            coeff = random.randint(1, 5)
            term = comp_func(self.x, coeff)
            
            # Apply nesting
            for _ in range(nesting):
                outer = random.choice(self.basic_functions)
                term = outer(term, 1)
            
            f_expr += random.randint(1, 10) * term
            
        return f_expr, self._score_difficulty(num_components, nesting)

    def generate_task(self, target_difficulty_band: float) -> Dict[str, Any]:
        """Provides an indefinite integral task."""
        complexity = max(1, int(target_difficulty_band))
        
        # 1. Generate F(x)
        F_expr, diff = self.generate_random_function(complexity)
        
        # 2. Differentiate to get the problem f(x)
        f_expr = sp.diff(F_expr, self.x)
        
        # 3. Format strings
        problem_text = f"Find the indefinite integral: \int ({sp.pretty(f_expr)}) dx"
        solution_text = f"{sp.simplify(F_expr)} + C"
        
        return {
            "problem": problem_text,
            "difficulty": diff,
            "solution": solution_text,
            "type": "integration",
            "sympy_F": F_expr,
            "sympy_f": f_expr
        }

    def generate_variants(self, task: Dict[str, Any], count: int = 2) -> list[Dict[str, Any]]:
        """
        LADDER Component: Recursive Decomposition for Integration.
        Breaks down sums or simplifies coefficients.
        """
        variants = []
        F_expr = task.get("sympy_F")
        
        if F_expr is None:
             # Fallback if task was not generated by us
             return [self.generate_task(max(1, task.get("difficulty", 2) - 2))]

        # Recursive Rule 1: Linearity (split sums)
        if isinstance(F_expr, sp.Add):
            args = F_expr.args
            for arg in args[:count]:
                sub_F = arg
                sub_f = sp.diff(sub_F, self.x)
                variants.append({
                    "problem": f"Integrate step-variant: \int ({sp.pretty(sub_f)}) dx",
                    "solution": f"{sub_F} + C",
                    "difficulty": task["difficulty"] - 1.0,
                    "type": "integration",
                    "sympy_F": sub_F,
                    "sympy_f": sub_f
                })

        # Recursive Rule 2: Constant simplification
        if not variants:
            # Just return a simpler integral by reducing difficulty
            variants.append(self.generate_task(max(1.0, task["difficulty"] - 2.0)))

        return variants[:count]