"""
Unity game engine code generator for MechanicsDSL.

Generates C# MonoBehaviour scripts for Unity physics simulation.
"""

from typing import Any, Dict, List, Optional

try:
    import sympy as sp
    from sympy.printing import ccode
    SYMPY_AVAILABLE = True
except ImportError:
    sp = None
    ccode = None
    SYMPY_AVAILABLE = False

try:
    from mechanics_dsl.codegen.base import CodeGenerator
except ImportError:
    CodeGenerator = object


def sympy_to_csharp(expr: Any, coord_map: Dict[str, int] = None) -> str:
    """
    Convert a sympy expression to C# code.
    
    Args:
        expr: Sympy expression to convert
        coord_map: Mapping of coordinate names to state array indices
        
    Returns:
        C# code string representing the expression
    """
    if expr is None:
        return "0.0f"
    
    if not SYMPY_AVAILABLE:
        return "0.0f /* sympy not available */"
    
    if coord_map is None:
        coord_map = {}
    
    # Convert sympy expression to C code (close to C#)
    try:
        # First get the C code representation
        c_code = ccode(expr)
        
        # Replace C math functions with C#/Unity equivalents
        replacements = [
            ('sin(', 'Mathf.Sin('),
            ('cos(', 'Mathf.Cos('),
            ('tan(', 'Mathf.Tan('),
            ('asin(', 'Mathf.Asin('),
            ('acos(', 'Mathf.Acos('),
            ('atan(', 'Mathf.Atan('),
            ('atan2(', 'Mathf.Atan2('),
            ('sqrt(', 'Mathf.Sqrt('),
            ('pow(', 'Mathf.Pow('),
            ('exp(', 'Mathf.Exp('),
            ('log(', 'Mathf.Log('),
            ('abs(', 'Mathf.Abs('),
            ('fabs(', 'Mathf.Abs('),
            ('floor(', 'Mathf.Floor('),
            ('ceil(', 'Mathf.Ceil('),
            ('M_PI', 'Mathf.PI'),
        ]
        
        for old, new in replacements:
            c_code = c_code.replace(old, new)
        
        # Ensure floating point literals
        # Add 'f' suffix to numeric literals that don't have it
        import re
        # Match numbers that are not already followed by 'f'
        c_code = re.sub(r'(\d+\.\d+)(?!f)', r'\1f', c_code)
        c_code = re.sub(r'(\d+)(?!\.\d)(?!f)(?![a-zA-Z_])', r'\1.0f', c_code)
        
        return c_code
        
    except Exception as e:
        return f"0.0f /* Error: {str(e)[:50]} */"


class UnityGenerator(CodeGenerator if CodeGenerator != object else object):
    """
    Generate Unity C# scripts from MechanicsDSL.

    Creates MonoBehaviour components that simulate physics
    using RK4 integration in FixedUpdate.

    Example:
        gen = UnityGenerator(compiler)
        gen.generate('PendulumSimulator.cs')
    """

    @property
    def target_name(self) -> str:
        return "unity"

    @property
    def file_extension(self) -> str:
        return ".cs"

    def __init__(self, compiler=None):
        super().__init__()
        self.compiler = compiler
        self.equations: Dict[str, Any] = {}  # Sympy equations for each coordinate
        
        if compiler:
            self.system_name = getattr(compiler, "system_name", "Physics")
            self.coordinates: List[str] = getattr(compiler.simulator, "coordinates", [])
            self.parameters: Dict[str, float] = dict(getattr(compiler.simulator, "parameters", {}))
            
            # Extract acceleration equations from compiler
            self._extract_equations(compiler)
        else:
            self.system_name = "Physics"
            self.coordinates = []
            self.parameters = {}
    
    def _extract_equations(self, compiler) -> None:
        """
        Extract acceleration equations from the compiler.
        
        Attempts to get equations from multiple sources:
        1. compiler.accelerations (direct dict)
        2. compiler.simulator.equations
        3. compiler.equations_of_motion
        """
        # Try accelerations dict first
        if hasattr(compiler, 'accelerations') and compiler.accelerations:
            self.equations = dict(compiler.accelerations)
            return
            
        # Try simulator equations
        if hasattr(compiler, 'simulator') and hasattr(compiler.simulator, 'equations'):
            eqs = compiler.simulator.equations
            if isinstance(eqs, dict):
                self.equations = eqs
                return
            elif isinstance(eqs, list) and SYMPY_AVAILABLE:
                # Convert list of equations to dict keyed by coordinate
                for i, eq in enumerate(eqs):
                    if i < len(self.coordinates):
                        self.equations[self.coordinates[i]] = eq
                return
        
        # Try equations_of_motion
        if hasattr(compiler, 'equations_of_motion') and compiler.equations_of_motion:
            eom = compiler.equations_of_motion
            if isinstance(eom, dict):
                self.equations = eom
            elif isinstance(eom, list) and SYMPY_AVAILABLE:
                for i, eq in enumerate(eom):
                    if i < len(self.coordinates):
                        self.equations[self.coordinates[i]] = eq

    def generate(self, output_file: Optional[str] = None) -> str:
        """Generate Unity C# code."""
        class_name = self._to_pascal_case(self.system_name)

        code = f"""using UnityEngine;
using System;

/// <summary>
/// {class_name} physics simulation.
/// Generated by MechanicsDSL.
/// </summary>
public class {class_name}Simulator : MonoBehaviour
{{
    [Header("Physical Parameters")]
{self._generate_parameters()}

    [Header("State Variables")]
{self._generate_state_fields()}

    [Header("Visualization")]
    public bool showTrail = true;
    public LineRenderer trailRenderer;
    public int maxTrailPoints = 100;
    
    [Header("Simulation")]
    public float timeScale = 1.0f;
    public bool useFixedDeltaTime = true;
    
    private float[] state;
    private float[] k1, k2, k3, k4, temp;
    
    void Start()
    {{
        InitializeState();
        if (trailRenderer == null && showTrail)
        {{
            trailRenderer = gameObject.AddComponent<LineRenderer>();
            trailRenderer.positionCount = 0;
            trailRenderer.startWidth = 0.05f;
            trailRenderer.endWidth = 0.02f;
        }}
    }}
    
    void InitializeState()
    {{
        int n = {len(self.coordinates) * 2};
        state = new float[n];
        k1 = new float[n];
        k2 = new float[n];
        k3 = new float[n];
        k4 = new float[n];
        temp = new float[n];
        
        // Initial conditions
{self._generate_initial_conditions()}
    }}
    
    void FixedUpdate()
    {{
        float dt = useFixedDeltaTime ? Time.fixedDeltaTime : Time.deltaTime;
        dt *= timeScale;
        
        // RK4 integration step
        RK4Step(dt);
        
        // Update transform
        UpdateTransform();
        
        // Update trail
        if (showTrail && trailRenderer != null)
        {{
            UpdateTrail();
        }}
    }}
    
    void RK4Step(float dt)
    {{
        float t = Time.time;
        
        // k1
        ComputeDerivatives(t, state, k1);
        
        // k2
        for (int i = 0; i < state.Length; i++)
            temp[i] = state[i] + 0.5f * dt * k1[i];
        ComputeDerivatives(t + 0.5f * dt, temp, k2);
        
        // k3
        for (int i = 0; i < state.Length; i++)
            temp[i] = state[i] + 0.5f * dt * k2[i];
        ComputeDerivatives(t + 0.5f * dt, temp, k3);
        
        // k4
        for (int i = 0; i < state.Length; i++)
            temp[i] = state[i] + dt * k3[i];
        ComputeDerivatives(t + dt, temp, k4);
        
        // Update state
        for (int i = 0; i < state.Length; i++)
            state[i] += dt * (k1[i] + 2*k2[i] + 2*k3[i] + k4[i]) / 6.0f;
    }}
    
    void ComputeDerivatives(float t, float[] y, float[] dydt)
    {{
        // Unpack state
{self._generate_state_unpack()}
        
        // Compute accelerations
{self._generate_accelerations()}
        
        // Pack derivatives
{self._generate_derivatives_pack()}
    }}
    
    void UpdateTransform()
    {{
        // Override in subclass for specific visualization
{self._generate_transform_update()}
    }}
    
    void UpdateTrail()
    {{
        if (trailRenderer.positionCount < maxTrailPoints)
        {{
            trailRenderer.positionCount++;
        }}
        else
        {{
            // Shift positions
            for (int i = 0; i < maxTrailPoints - 1; i++)
                trailRenderer.SetPosition(i, trailRenderer.GetPosition(i + 1));
        }}
        trailRenderer.SetPosition(trailRenderer.positionCount - 1, transform.position);
    }}
    
    /// <summary>
    /// Reset simulation to initial conditions.
    /// </summary>
    public void ResetSimulation()
    {{
        InitializeState();
        if (trailRenderer != null)
            trailRenderer.positionCount = 0;
    }}
    
    /// <summary>
    /// Get current state as array.
    /// </summary>
    public float[] GetState()
    {{
        return (float[])state.Clone();
    }}
    
    /// <summary>
    /// Set state from array.
    /// </summary>
    public void SetState(float[] newState)
    {{
        if (newState.Length == state.Length)
            Array.Copy(newState, state, state.Length);
    }}
}}
"""

        if output_file:
            with open(output_file, "w") as f:
                f.write(code)

        return code

    def _to_pascal_case(self, name: str) -> str:
        """Convert snake_case to PascalCase."""
        return "".join(word.capitalize() for word in name.split("_"))

    def _generate_parameters(self) -> str:
        """Generate parameter fields."""
        lines = []
        for name, value in self.parameters.items():
            lines.append(f"    public float {name} = {value}f;")
        return "\n".join(lines) if lines else "    // No parameters"

    def _generate_state_fields(self) -> str:
        """Generate state variable fields."""
        lines = []
        for coord in self.coordinates:
            lines.append(f"    [SerializeField] private float {coord};")
            lines.append(f"    [SerializeField] private float {coord}_dot;")
        return "\n".join(lines) if lines else "    // No state variables"

    def _generate_initial_conditions(self) -> str:
        """Generate initial condition assignments."""
        lines = []
        for i, coord in enumerate(self.coordinates):
            lines.append(f"        state[{2*i}] = {coord};")
            lines.append(f"        state[{2*i + 1}] = {coord}_dot;")
        return "\n".join(lines) if lines else "        // No initial conditions"

    def _generate_state_unpack(self) -> str:
        """Generate state unpacking code."""
        lines = []
        for i, coord in enumerate(self.coordinates):
            lines.append(f"        float {coord} = y[{2*i}];")
            lines.append(f"        float {coord}_dot = y[{2*i + 1}];")
        return "\n".join(lines) if lines else "        // No state to unpack"

    def _generate_accelerations(self) -> str:
        """
        Generate acceleration computations from sympy equations.
        
        Converts symbolic acceleration equations to C# code using
        the sympy_to_csharp converter function.
        """
        lines = []
        
        for coord in self.coordinates:
            # Check if we have a real equation for this coordinate
            if coord in self.equations and self.equations[coord] is not None:
                # Convert sympy expression to C# code
                accel_expr = self.equations[coord]
                csharp_code = sympy_to_csharp(accel_expr)
                lines.append(f"        float {coord}_ddot = {csharp_code};")
            elif f"{coord}_ddot" in self.equations:
                # Try with _ddot suffix
                accel_expr = self.equations[f"{coord}_ddot"]
                csharp_code = sympy_to_csharp(accel_expr)
                lines.append(f"        float {coord}_ddot = {csharp_code};")
            else:
                # No equation available - generate placeholder with warning
                lines.append(f"        float {coord}_ddot = 0.0f; // Warning: No equation extracted for {coord}")
        
        return "\n".join(lines) if lines else "        // No accelerations defined"

    def _generate_derivatives_pack(self) -> str:
        """Generate derivative packing."""
        lines = []
        for i, coord in enumerate(self.coordinates):
            lines.append(f"        dydt[{2*i}] = {coord}_dot;")
            lines.append(f"        dydt[{2*i + 1}] = {coord}_ddot;")
        return "\n".join(lines) if lines else "        // No derivatives"

    def _generate_transform_update(self) -> str:
        """Generate transform update (basic rotation for angles)."""
        if not self.coordinates:
            return "        // No coordinates"

        self.coordinates[0]
        return f"""        // Example: rotate based on first coordinate
        float angle = state[0] * Mathf.Rad2Deg;
        transform.localRotation = Quaternion.Euler(0, 0, -angle);"""


__all__ = ["UnityGenerator"]
