"""
MATLAB/Octave Code Generator for MechanicsDSL

Generates MATLAB/GNU Octave simulation scripts.
"""

from sympy.printing.octave import octave_code

from ..utils import logger
from .base import CodeGenerator


class MatlabGenerator(CodeGenerator):
    """
    Generates MATLAB/Octave simulation code.

    Produces scripts compatible with both MATLAB and GNU Octave.
    """

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

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

    def generate(self, output_file: str = "simulation.m") -> str:
        """Generate MATLAB/Octave simulation code."""
        logger.info(f"Generating MATLAB 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 MATLAB code for equations."""
        lines = []
        idx = 1  # MATLAB is 1-indexed
        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]
                m_expr = octave_code(expr)
                lines.append(f"    dydt({idx+1}) = {m_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 MATLAB/Octave simulation script."""
        # Parameter definitions (as global variables)
        param_lines = []
        param_globals = []
        for name, val in self.parameters.items():
            param_lines.append(f"{name} = {val};")
            param_globals.append(name)
        param_str = "\n".join(param_lines)
        global_decl = " ".join(param_globals)

        # State variable unpacking (MATLAB is 1-indexed)
        unpack_lines = []
        idx = 1
        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_vals = []
        for coord in self.coordinates:
            pos = self.initial_conditions.get(coord, 0.0)
            vel = self.initial_conditions.get(f"{coord}_dot", 0.0)
            init_vals.extend([str(pos), str(vel)])
        init_str = "; ".join(init_vals)

        state_dim = len(self.coordinates) * 2

        template = f"""%% Auto-generated simulation: {self.system_name}
%% Generated by MechanicsDSL
%% Compatible with MATLAB and GNU Octave

clear; clc; close all;

%% Physical Parameters
global {global_decl}
{param_str}

%% Initial conditions
y0 = [{init_str}];
tspan = [0 10];

%% Solve ODE
[t, y] = ode45(@equations_of_motion, tspan, y0);

%% Plot results
figure;
subplot(2,1,1);
plot(t, y(:,1), 'b-', 'LineWidth', 2);
xlabel('Time (s)');
ylabel('Position');
title('{self.system_name}');
grid on;

subplot(2,1,2);
plot(t, y(:,2), 'r-', 'LineWidth', 2);
xlabel('Time (s)');
ylabel('Velocity');
grid on;

saveas(gcf, '{self.system_name}.png');
fprintf('Simulation complete: %d points\\n', length(t));

%% Equations of motion function
function dydt = equations_of_motion(t, y)
    global {global_decl}
    
    dydt = zeros({state_dim}, 1);
    
    % Unpack state
{unpack_str}
    
    % Compute derivatives
{eq_str}
end
"""
        return template
