"""
Python Code Generator for MechanicsDSL

Generates standalone Python simulation scripts.
"""

from sympy.printing import pycode

from ..utils import logger
from .base import CodeGenerator


class PythonGenerator(CodeGenerator):
    """
    Generates Python simulation code with NumPy/SciPy.

    Produces standalone scripts that can run without MechanicsDSL installed.
    """

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

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

    def generate(self, output_file: str = "simulation.py") -> str:
        """Generate Python simulation code."""
        logger.info(f"Generating Python code for {self.system_name}")

        code = self._generate_code()

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

        logger.info(f"Successfully wrote {output_file}")
        return output_file

    def generate_equations(self) -> str:
        """Generate Python code for equations."""
        lines = []
        idx = 0
        for coord in self.coordinates:
            accel_key = f"{coord}_ddot"
            lines.append(f"    dydt[{idx}] = y[{idx+1}]  # d{coord}/dt = {coord}_dot")
            if accel_key in self.equations:
                expr = self.equations[accel_key]
                py_expr = pycode(expr)
                lines.append(f"    dydt[{idx+1}] = {py_expr}  # d{coord}_dot/dt")
            else:
                lines.append(f"    dydt[{idx+1}] = 0.0")
            idx += 2
        return "\n".join(lines)

    def _generate_code(self) -> str:
        """Generate complete Python simulation script."""
        # Parameter definitions
        param_str = "\n".join(f"{name} = {val}" for name, val in self.parameters.items())

        # State variable unpacking
        unpack_lines = []
        idx = 0
        for coord in self.coordinates:
            unpack_lines.append(f"    {coord} = y[{idx}]")
            unpack_lines.append(f"    {coord}_dot = y[{idx+1}]")
            idx += 2
        unpack_str = "\n".join(unpack_lines)

        # Equations
        eq_str = self.generate_equations()

        # Initial conditions
        init_str = self.generate_initial_conditions()

        # Full template
        template = f'''"""
Auto-generated simulation: {self.system_name}
Generated by MechanicsDSL
"""
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# Physical Parameters
{param_str}

def equations_of_motion(t, y):
    """Equations of motion for {self.system_name}."""
    dydt = np.zeros({len(self.coordinates) * 2})
    
    # Unpack state
{unpack_str}
    
    # Compute derivatives
{eq_str}
    
    return dydt

def simulate(t_span=(0, 10), num_points=1000):
    """Run simulation."""
    y0 = [{init_str}]
    t_eval = np.linspace(t_span[0], t_span[1], num_points)
    
    solution = solve_ivp(
        equations_of_motion,
        t_span,
        y0,
        t_eval=t_eval,
        method='RK45'
    )
    
    return solution

def plot_results(solution):
    """Plot simulation results."""
    fig, axes = plt.subplots(2, 1, figsize=(10, 8), sharex=True)
    
    # Position
    axes[0].plot(solution.t, solution.y[0], label='Position')
    axes[0].set_ylabel('Position')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # Velocity
    axes[1].plot(solution.t, solution.y[1], label='Velocity')
    axes[1].set_xlabel('Time (s)')
    axes[1].set_ylabel('Velocity')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    
    plt.suptitle('{self.system_name}')
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    print("Running {self.system_name} simulation...")
    sol = simulate()
    print(f"Simulation complete: {{len(sol.t)}} points")
    plot_results(sol)
'''
        return template
