"""
Python Code Generator for MechanicsDSL

Generates standalone Python simulation scripts.
"""
from typing import Dict, List
import sympy as sp
from sympy.printing import pycode

from .base import CodeGenerator
from ..utils import logger


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
