"""
Arduino Code Generator for MechanicsDSL

Generates Arduino-compatible C++ code for embedded physics simulations.
Optimized for microcontroller constraints (limited RAM/ROM).
"""

from typing import Dict, List

import sympy as sp
from sympy.printing.c import ccode

from ..utils import logger
from .base import CodeGenerator


class ArduinoGenerator(CodeGenerator):
    """
    Generates Arduino sketch (.ino) files for embedded simulations.

    Features:
    - Fixed-point arithmetic option for speed
    - RAM-optimized data structures
    - Serial plotter output
    - PWM/servo output for physical feedback

    Example:
        >>> gen = ArduinoGenerator(
        ...     system_name="pendulum",
        ...     coordinates=['theta'],
        ...     parameters={'g': 9.81, 'l': 1.0},
        ...     initial_conditions={'theta': 0.5, 'theta_dot': 0.0},
        ...     equations={'theta_ddot': -g/l * sin(theta)}
        ... )
        >>> gen.generate("pendulum.ino")
    """

    def __init__(
        self,
        system_name: str,
        coordinates: List[str],
        parameters: Dict[str, float],
        initial_conditions: Dict[str, float],
        equations: Dict[str, sp.Expr],
        use_serial_plotter: bool = True,
        servo_pin: int = None,
    ):

        super().__init__(system_name, coordinates, parameters, initial_conditions, equations)

        self.use_serial_plotter = use_serial_plotter
        self.servo_pin = servo_pin

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

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

    def generate(self, output_file: str) -> str:
        """Generate Arduino sketch file."""
        logger.info(f"Generating Arduino code for {self.system_name}")

        code = self._generate_source()

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

        logger.info(f"Generated {output_file}")
        return output_file

    def generate_equations(self) -> str:
        """Generate equations code."""
        lines = []
        idx = 0
        for coord in self.coordinates:
            accel_key = f"{coord}_ddot"
            lines.append(f"  dydt[{idx}] = state[{idx+1}];")
            if accel_key in self.equations:
                expr = self.equations[accel_key]
                c_expr = ccode(expr)
                lines.append(f"  dydt[{idx+1}] = {c_expr};")
            else:
                lines.append(f"  dydt[{idx+1}] = 0.0;")
            idx += 2
        return "\n".join(lines)

    def _generate_source(self) -> str:
        """Generate complete Arduino sketch."""
        state_dim = len(self.coordinates) * 2

        params = "\n".join(f"const float {name} = {val}f;" for name, val in self.parameters.items())

        init_vals = []
        for coord in self.coordinates:
            init_vals.append(str(self.initial_conditions.get(coord, 0.0)) + "f")
            init_vals.append(str(self.initial_conditions.get(f"{coord}_dot", 0.0)) + "f")
        init_str = ", ".join(init_vals)

        # State unpacking
        unpack = "\n".join(
            f"  float {c} = state[{2*i}]; float {c}_dot = state[{2*i+1}];"
            for i, c in enumerate(self.coordinates)
        )

        # Servo code
        servo_include = ""
        servo_init = ""
        servo_update = ""
        if self.servo_pin is not None:
            servo_include = "#include <Servo.h>\nServo outputServo;"
            servo_init = f"outputServo.attach({self.servo_pin});"
            servo_update = """
  // Map first coordinate to servo angle (0-180)
  int angle = constrain(map(state[0] * 100, -314, 314, 0, 180), 0, 180);
  outputServo.write(angle);"""

        # Serial output
        serial_output = ""
        if self.use_serial_plotter:
            serial_output = """
  // Serial Plotter format
  Serial.print(t);
  for (int i = 0; i < STATE_DIM; i++) {
    Serial.print(",");
    Serial.print(state[i], 4);
  }
  Serial.println();"""

        return f"""/*
 * Arduino Physics Simulation: {self.system_name}
 * Generated by MechanicsDSL
 * 
 * Upload to Arduino and open Serial Plotter (Tools > Serial Plotter)
 */

{servo_include}

// Physical Parameters
{params}

#define STATE_DIM {state_dim}

// Simulation state
float state[STATE_DIM] = {{ {init_str} }};
float t = 0.0f;
const float dt = 0.01f;  // 10ms timestep
unsigned long lastMicros = 0;

// Compute derivatives
void computeDerivatives(const float* state, float* dydt) {{
{unpack}

{self.generate_equations()}
}}

// RK4 integration step (using float for speed)
void rk4Step() {{
  float k1[STATE_DIM], k2[STATE_DIM], k3[STATE_DIM], k4[STATE_DIM];
  float temp[STATE_DIM];
  
  computeDerivatives(state, k1);
  
  for (int i = 0; i < STATE_DIM; i++) temp[i] = state[i] + 0.5f * dt * k1[i];
  computeDerivatives(temp, k2);
  
  for (int i = 0; i < STATE_DIM; i++) temp[i] = state[i] + 0.5f * dt * k2[i];
  computeDerivatives(temp, k3);
  
  for (int i = 0; i < STATE_DIM; i++) temp[i] = state[i] + dt * k3[i];
  computeDerivatives(temp, k4);
  
  for (int i = 0; i < STATE_DIM; i++) {{
    state[i] += dt * (k1[i] + 2.0f*k2[i] + 2.0f*k3[i] + k4[i]) / 6.0f;
  }}
  
  t += dt;
}}

void setup() {{
  Serial.begin(115200);
  while (!Serial) {{ ; }}  // Wait for serial connection
  
  {servo_init}
  
  Serial.println("MechanicsDSL: {self.system_name}");
  Serial.println("Time,{','.join([c + ',' + c + '_dot' for c in self.coordinates]).strip(',')}");
  
  lastMicros = micros();
}}

void loop() {{
  // Maintain consistent timing
  unsigned long now = micros();
  if (now - lastMicros >= (unsigned long)(dt * 1000000)) {{
    lastMicros = now;
    
    // Integration step
    rk4Step();
{servo_update}
{serial_output}
  }}
}}

// Reset simulation (call from Serial command if needed)
void resetSimulation() {{
  float initial[STATE_DIM] = {{ {init_str} }};
  for (int i = 0; i < STATE_DIM; i++) {{
    state[i] = initial[i];
  }}
  t = 0.0f;
}}
"""
