"""
Fortran Code Generator for MechanicsDSL

Generates Fortran 90+ simulation code.
"""

from sympy.printing.fortran import fcode

from ..utils import logger
from .base import CodeGenerator


class FortranGenerator(CodeGenerator):
    """
    Generates Fortran 90+ simulation code.

    Produces standalone programs with explicit RK4 integration.
    """

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

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

    def generate(self, output_file: str = "simulation.f90") -> str:
        """Generate Fortran simulation code."""
        logger.info(f"Generating Fortran 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 Fortran code for equations."""
        lines = []
        idx = 1  # Fortran 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]
                f_expr = fcode(expr, source_format="free")
                lines.append(f"    dydt({idx+1}) = {f_expr}  ! d{coord}_dot/dt")
            else:
                lines.append(f"    dydt({idx+1}) = 0.0d0")
            idx += 2
        return "\n".join(lines)

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

        # State variable unpacking
        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)

        # Variable declarations for unpacking
        var_decls = []
        for coord in self.coordinates:
            var_decls.append(f"    real(8) :: {coord}, {coord}_dot")
        var_decl_str = "\n".join(var_decls)

        # 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([f"{pos}d0", f"{vel}d0"])

        state_dim = len(self.coordinates) * 2

        init_assignments = []
        for i, val in enumerate(init_vals):
            init_assignments.append(f"    y({i+1}) = {val}")
        init_str = "\n".join(init_assignments)

        template = f"""! Auto-generated simulation: {self.system_name}
! Generated by MechanicsDSL
! Compile with: gfortran -O3 -o simulation simulation.f90

program {self.system_name}_simulation
    implicit none
    
    integer, parameter :: DIM = {state_dim}
    real(8) :: y(DIM), dydt(DIM), k1(DIM), k2(DIM), k3(DIM), k4(DIM), tmp(DIM)
    real(8) :: t, dt, t_end
    integer :: i, unit_out
    
    ! Physical Parameters
{param_str}
    
    ! Initial conditions
{init_str}
    
    dt = 0.01d0
    t_end = 10.0d0
    t = 0.0d0
    
    ! Open output file
    unit_out = 20
    open(unit=unit_out, file='{self.system_name}.csv', status='replace')
    write(unit_out, '(A)') 't,{",".join(self.coordinates)}'
    
    print *, 'Running {self.system_name} simulation...'
    
    ! RK4 integration loop
    do while (t < t_end)
        write(unit_out, '({state_dim + 1}(ES14.6,:,","))') t, {", ".join(f"y({i*2+1})" for i in range(len(self.coordinates)))}
        
        call equations_of_motion(t, y, k1)
        tmp = y + 0.5d0 * dt * k1
        call equations_of_motion(t + 0.5d0*dt, tmp, k2)
        tmp = y + 0.5d0 * dt * k2
        call equations_of_motion(t + 0.5d0*dt, tmp, k3)
        tmp = y + dt * k3
        call equations_of_motion(t + dt, tmp, k4)
        
        y = y + dt/6.0d0 * (k1 + 2.0d0*k2 + 2.0d0*k3 + k4)
        t = t + dt
    end do
    
    close(unit_out)
    print *, 'Simulation complete. Results saved to {self.system_name}.csv'
    
contains

    subroutine equations_of_motion(t, y, dydt)
        real(8), intent(in) :: t, y(DIM)
        real(8), intent(out) :: dydt(DIM)
{var_decl_str}
        
        ! Unpack state
{unpack_str}
        
        ! Compute derivatives
{eq_str}
    end subroutine

end program
"""
        return template
