"""
Inline Code Generator - Phase 2.8

Auto-generates simple Python code for mathematical operations that don't require
external documentation or research.

This handles the "inline_calculations" from Phase 2.7 LLM analysis.

Examples:
- Calculate average: avg = sum(values) / len(values)
- Find minimum: min_val = min(values)
- Normalize: norm_val = value / divisor
- Calculate percentage: pct = (value / baseline) * 100

Author: Atomizer Development Team
Version: 0.1.0 (Phase 2.8)
Last Updated: 2025-01-16
"""

from typing import Dict, Any, List, Optional
from dataclasses import dataclass


@dataclass
class GeneratedCode:
    """Result of code generation."""
    code: str
    variables_used: List[str]
    variables_created: List[str]
    imports_needed: List[str]
    description: str


class InlineCodeGenerator:
    """
    Generates Python code for simple mathematical operations.

    This class takes structured calculation descriptions (from LLM Phase 2.7)
    and generates clean, executable Python code.
    """

    def __init__(self):
        """Initialize the code generator."""
        self.supported_operations = {
            'mean', 'average', 'avg',
            'min', 'minimum',
            'max', 'maximum',
            'sum', 'total',
            'count', 'length',
            'normalize', 'norm',
            'percentage', 'percent', 'pct',
            'ratio',
            'difference', 'diff',
            'add', 'subtract', 'multiply', 'divide',
            'abs', 'absolute',
            'sqrt', 'square_root',
            'power', 'pow'
        }

    def generate_from_llm_output(self, calculation: Dict[str, Any]) -> GeneratedCode:
        """
        Generate code from LLM-analyzed calculation.

        Args:
            calculation: Dictionary from LLM with keys:
                - action: str (e.g., "calculate_average")
                - description: str
                - params: dict with input/operation/etc.
                - code_hint: str (optional, from LLM)

        Returns:
            GeneratedCode with executable Python code
        """
        action = calculation.get('action', '')
        params = calculation.get('params', {})
        description = calculation.get('description', '')
        code_hint = calculation.get('code_hint', '')

        # If LLM provided a code hint, validate and use it
        if code_hint:
            return self._from_code_hint(code_hint, params, description)

        # Otherwise, generate from action/params
        return self._from_action_params(action, params, description)

    def _from_code_hint(self, code_hint: str, params: Dict[str, Any],
                        description: str) -> GeneratedCode:
        """Generate from LLM-provided code hint."""
        # Extract variable names from code hint
        variables_used = self._extract_input_variables(code_hint, params)
        variables_created = self._extract_output_variables(code_hint)
        imports_needed = self._extract_imports_needed(code_hint)

        return GeneratedCode(
            code=code_hint.strip(),
            variables_used=variables_used,
            variables_created=variables_created,
            imports_needed=imports_needed,
            description=description
        )

    def _from_action_params(self, action: str, params: Dict[str, Any],
                           description: str) -> GeneratedCode:
        """Generate code from action name and parameters."""
        operation = params.get('operation', '').lower()
        input_var = params.get('input', 'values')
        divisor = params.get('divisor')
        baseline = params.get('baseline')
        current = params.get('current')

        # Detect operation type
        if any(op in action.lower() or op in operation for op in ['avg', 'average', 'mean']):
            return self._generate_average(input_var, description)

        elif any(op in action.lower() or op in operation for op in ['min', 'minimum']):
            return self._generate_min(input_var, description)

        elif any(op in action.lower() or op in operation for op in ['max', 'maximum']):
            return self._generate_max(input_var, description)

        elif any(op in action.lower() for op in ['normalize', 'norm']) and divisor:
            return self._generate_normalization(input_var, divisor, description)

        elif any(op in action.lower() for op in ['percentage', 'percent', 'pct', 'increase']):
            current = params.get('current')
            baseline = params.get('baseline')
            if current and baseline:
                return self._generate_percentage_change(current, baseline, description)
            elif divisor:
                return self._generate_percentage(input_var, divisor, description)

        elif 'sum' in action.lower() or 'total' in action.lower():
            return self._generate_sum(input_var, description)

        elif 'ratio' in action.lower():
            inputs = params.get('inputs', [])
            if len(inputs) >= 2:
                return self._generate_ratio(inputs[0], inputs[1], description)

        # Fallback: generic operation
        return self._generate_generic(action, params, description)

    def _generate_average(self, input_var: str, description: str) -> GeneratedCode:
        """Generate code to calculate average."""
        output_var = f"avg_{input_var}" if not input_var.startswith('avg') else input_var.replace('input', 'avg')
        code = f"{output_var} = sum({input_var}) / len({input_var})"

        return GeneratedCode(
            code=code,
            variables_used=[input_var],
            variables_created=[output_var],
            imports_needed=[],
            description=description or f"Calculate average of {input_var}"
        )

    def _generate_min(self, input_var: str, description: str) -> GeneratedCode:
        """Generate code to find minimum."""
        output_var = f"min_{input_var}" if not input_var.startswith('min') else input_var.replace('input', 'min')
        code = f"{output_var} = min({input_var})"

        return GeneratedCode(
            code=code,
            variables_used=[input_var],
            variables_created=[output_var],
            imports_needed=[],
            description=description or f"Find minimum of {input_var}"
        )

    def _generate_max(self, input_var: str, description: str) -> GeneratedCode:
        """Generate code to find maximum."""
        output_var = f"max_{input_var}" if not input_var.startswith('max') else input_var.replace('input', 'max')
        code = f"{output_var} = max({input_var})"

        return GeneratedCode(
            code=code,
            variables_used=[input_var],
            variables_created=[output_var],
            imports_needed=[],
            description=description or f"Find maximum of {input_var}"
        )

    def _generate_normalization(self, input_var: str, divisor: float,
                                description: str) -> GeneratedCode:
        """Generate code to normalize a value."""
        output_var = f"norm_{input_var}" if not input_var.startswith('norm') else input_var
        code = f"{output_var} = {input_var} / {divisor}"

        return GeneratedCode(
            code=code,
            variables_used=[input_var],
            variables_created=[output_var],
            imports_needed=[],
            description=description or f"Normalize {input_var} by {divisor}"
        )

    def _generate_percentage_change(self, current: str, baseline: str,
                                    description: str) -> GeneratedCode:
        """Generate code to calculate percentage change."""
        # Infer output variable name from inputs
        if 'mass' in current.lower() or 'mass' in baseline.lower():
            output_var = "mass_increase_pct"
        else:
            output_var = f"{current}_vs_{baseline}_pct"

        code = f"{output_var} = (({current} - {baseline}) / {baseline}) * 100.0"

        return GeneratedCode(
            code=code,
            variables_used=[current, baseline],
            variables_created=[output_var],
            imports_needed=[],
            description=description or f"Calculate percentage change from {baseline} to {current}"
        )

    def _generate_percentage(self, input_var: str, divisor: float,
                            description: str) -> GeneratedCode:
        """Generate code to calculate percentage."""
        output_var = f"pct_{input_var}"
        code = f"{output_var} = ({input_var} / {divisor}) * 100.0"

        return GeneratedCode(
            code=code,
            variables_used=[input_var],
            variables_created=[output_var],
            imports_needed=[],
            description=description or f"Calculate percentage of {input_var} vs {divisor}"
        )

    def _generate_sum(self, input_var: str, description: str) -> GeneratedCode:
        """Generate code to calculate sum."""
        output_var = f"total_{input_var}" if not input_var.startswith('total') else input_var
        code = f"{output_var} = sum({input_var})"

        return GeneratedCode(
            code=code,
            variables_used=[input_var],
            variables_created=[output_var],
            imports_needed=[],
            description=description or f"Calculate sum of {input_var}"
        )

    def _generate_ratio(self, numerator: str, denominator: str,
                       description: str) -> GeneratedCode:
        """Generate code to calculate ratio."""
        output_var = f"{numerator}_to_{denominator}_ratio"
        code = f"{output_var} = {numerator} / {denominator}"

        return GeneratedCode(
            code=code,
            variables_used=[numerator, denominator],
            variables_created=[output_var],
            imports_needed=[],
            description=description or f"Calculate ratio of {numerator} to {denominator}"
        )

    def _generate_generic(self, action: str, params: Dict[str, Any],
                         description: str) -> GeneratedCode:
        """Generate generic calculation code."""
        # Extract operation from action name
        operation = action.lower().replace('calculate_', '').replace('find_', '').replace('get_', '')
        input_var = params.get('input', 'value')
        output_var = f"{operation}_result"

        # Try to infer code from parameters
        if 'formula' in params:
            code = f"{output_var} = {params['formula']}"
        else:
            code = f"{output_var} = {input_var}  # TODO: Implement {action}"

        return GeneratedCode(
            code=code,
            variables_used=[input_var] if input_var != 'value' else [],
            variables_created=[output_var],
            imports_needed=[],
            description=description or f"Generic calculation: {action}"
        )

    def _extract_input_variables(self, code: str, params: Dict[str, Any]) -> List[str]:
        """Extract input variable names from code."""
        variables = []

        # Get from params if available
        if 'input' in params:
            variables.append(params['input'])
        if 'inputs' in params:
            variables.extend(params.get('inputs', []))
        if 'current' in params:
            variables.append(params['current'])
        if 'baseline' in params:
            variables.append(params['baseline'])

        # Extract from code (variables on right side of =)
        if '=' in code:
            rhs = code.split('=', 1)[1]
            # Simple extraction of variable names (alphanumeric + underscore)
            import re
            found_vars = re.findall(r'\b[a-zA-Z_][a-zA-Z0-9_]*\b', rhs)
            variables.extend([v for v in found_vars if v not in ['sum', 'min', 'max', 'len', 'abs']])

        return list(set(variables))  # Remove duplicates

    def _extract_output_variables(self, code: str) -> List[str]:
        """Extract output variable names from code."""
        # Variables on left side of =
        if '=' in code:
            lhs = code.split('=', 1)[0].strip()
            return [lhs]
        return []

    def _extract_imports_needed(self, code: str) -> List[str]:
        """Extract required imports from code."""
        imports = []

        # Check for math functions
        if any(func in code for func in ['sqrt', 'pow', 'log', 'exp', 'sin', 'cos']):
            imports.append('import math')

        # Check for numpy functions
        if any(func in code for func in ['np.', 'numpy.']):
            imports.append('import numpy as np')

        return imports

    def generate_batch(self, calculations: List[Dict[str, Any]]) -> List[GeneratedCode]:
        """
        Generate code for multiple calculations.

        Args:
            calculations: List of calculation dictionaries from LLM

        Returns:
            List of GeneratedCode objects
        """
        return [self.generate_from_llm_output(calc) for calc in calculations]

    def generate_executable_script(self, calculations: List[Dict[str, Any]],
                                   inputs: Dict[str, Any] = None) -> str:
        """
        Generate a complete executable Python script with all calculations.

        Args:
            calculations: List of calculations
            inputs: Optional input values for testing

        Returns:
            Complete Python script as string
        """
        generated = self.generate_batch(calculations)

        # Collect all imports
        all_imports = []
        for code in generated:
            all_imports.extend(code.imports_needed)
        all_imports = list(set(all_imports))  # Remove duplicates

        # Build script
        lines = []

        # Header
        lines.append('"""')
        lines.append('Auto-generated inline calculations')
        lines.append('Generated by Atomizer Phase 2.8 Inline Code Generator')
        lines.append('"""')
        lines.append('')

        # Imports
        if all_imports:
            lines.extend(all_imports)
            lines.append('')

        # Input values (if provided for testing)
        if inputs:
            lines.append('# Input values')
            for var_name, value in inputs.items():
                lines.append(f'{var_name} = {repr(value)}')
            lines.append('')

        # Calculations
        lines.append('# Inline calculations')
        for code_obj in generated:
            lines.append(f'# {code_obj.description}')
            lines.append(code_obj.code)
            lines.append('')

        return '\n'.join(lines)


def main():
    """Test the inline code generator."""
    print("=" * 80)
    print("Phase 2.8: Inline Code Generator Test")
    print("=" * 80)
    print()

    generator = InlineCodeGenerator()

    # Test cases from Phase 2.7 LLM output
    test_calculations = [
        {
            "action": "normalize_stress",
            "description": "Normalize stress by 200 MPa",
            "params": {
                "input": "max_stress",
                "divisor": 200.0,
                "units": "MPa"
            }
        },
        {
            "action": "normalize_displacement",
            "description": "Normalize displacement by 5 mm",
            "params": {
                "input": "max_disp_y",
                "divisor": 5.0,
                "units": "mm"
            }
        },
        {
            "action": "calculate_mass_increase",
            "description": "Calculate mass increase percentage vs baseline",
            "params": {
                "current": "panel_total_mass",
                "baseline": "baseline_mass"
            }
        },
        {
            "action": "calculate_average",
            "description": "Calculate average of extracted forces",
            "params": {
                "input": "forces_z",
                "operation": "mean"
            }
        },
        {
            "action": "find_minimum",
            "description": "Find minimum force value",
            "params": {
                "input": "forces_z",
                "operation": "min"
            }
        }
    ]

    print("Test Calculations:")
    print()

    for i, calc in enumerate(test_calculations, 1):
        print(f"{i}. {calc['description']}")
        code_obj = generator.generate_from_llm_output(calc)
        print(f"   Generated Code: {code_obj.code}")
        print(f"   Inputs:  {', '.join(code_obj.variables_used)}")
        print(f"   Outputs: {', '.join(code_obj.variables_created)}")
        print()

    # Generate complete script
    print("=" * 80)
    print("Complete Executable Script:")
    print("=" * 80)
    print()

    test_inputs = {
        'max_stress': 150.5,
        'max_disp_y': 3.2,
        'panel_total_mass': 2.8,
        'baseline_mass': 2.5,
        'forces_z': [10.5, 12.3, 8.9, 11.2, 9.8]
    }

    script = generator.generate_executable_script(test_calculations, test_inputs)
    print(script)


if __name__ == '__main__':
    main()
