"""
Julia Code Generator for MechanicsDSL

Generates standalone Julia simulation scripts using DifferentialEquations.jl.
"""

from sympy.printing.julia import julia_code

from ..utils import logger
from .base import CodeGenerator


class JuliaGenerator(CodeGenerator):
    """
    Generates Julia simulation code with DifferentialEquations.jl.

    Produces standalone scripts that can run with Julia 1.6+.
    """

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

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

    def generate(self, output_file: str = "simulation.jl") -> str:
        """Generate Julia simulation code."""
        logger.info(f"Generating Julia 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 Julia code for equations."""
        lines = []
        idx = 1  # Julia is 1-indexed
        for coord in self.coordinates:
            accel_key = f"{coord}_ddot"
            lines.append(f"    du[{idx}] = u[{idx+1}]  # d{coord}/dt = {coord}_dot")
            if accel_key in self.equations:
                expr = self.equations[accel_key]
                jl_expr = julia_code(expr)
                lines.append(f"    du[{idx+1}] = {jl_expr}  # d{coord}_dot/dt")
            else:
                lines.append(f"    du[{idx+1}] = 0.0")
            idx += 2
        return "\n".join(lines)

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

        # State variable unpacking (Julia is 1-indexed)
        unpack_lines = []
        idx = 1
        for coord in self.coordinates:
            unpack_lines.append(f"    {coord} = u[{idx}]")
            unpack_lines.append(f"    {coord}_dot = u[{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)

        len(self.coordinates) * 2

        template = f"""#=
Auto-generated simulation: {self.system_name}
Generated by MechanicsDSL
=#

using DifferentialEquations
using Plots

# Physical Parameters
{param_str}

function equations_of_motion!(du, u, p, t)
    # Unpack state
{unpack_str}
    
    # Compute derivatives
{eq_str}
end

function simulate(tspan=(0.0, 10.0))
    u0 = [{init_str}]
    prob = ODEProblem(equations_of_motion!, u0, tspan)
    sol = solve(prob, Tsit5(), saveat=0.01)
    return sol
end

function plot_results(sol)
    p1 = plot(sol.t, sol[1, :], label="Position", xlabel="Time (s)", ylabel="Position")
    p2 = plot(sol.t, sol[2, :], label="Velocity", xlabel="Time (s)", ylabel="Velocity")
    plot(p1, p2, layout=(2, 1), title="{self.system_name}")
end

# Run simulation
println("Running {self.system_name} simulation...")
sol = simulate()
println("Simulation complete: $(length(sol.t)) points")
plot_results(sol)
savefig("{self.system_name}.png")
"""
        return template
