incomp-flame-post/code/code_gen/post.py

"""
DNS Post-Processing Code Generator Core (post.py)
==================================================
이 모듈은 난류 및 연소 DNS(Direct Numerical Simulation) 데이터의 후처리를 위한 고성능 Fortran 코드를 생성하는 컴파일러의 코어입니다.
사용자가 정의한 DSL(Domain Specific Language) 입력 식을 읽어 파싱한 후, 다음과 같은 4단계 최적화 컴파일 과정을 거쳐 극도로 최적화된 3차원 루프 Fortran 모듈 코드를 자동 생성합니다.

[컴파일러 파이프라인 4단계 개요]
1. Stage 1 (AST 수집 및 변수 정의):
   - Lark 파서가 생성한 AST(Abstract Syntax Tree)를 순회하며 기본 입력 변수(Primary), 계산이 필요한 대입식(Derived), 
     그리고 통계 물리량 계산을 위한 평균화 변수(Averaged)를 추출하고 필드 메타데이터 객체를 구성합니다.
2. Stage 2 (수치 미분 및 변동량 확장):
   - 수식 내에 수치 미분자(ddx, d2dy 등)나 변동량(fluctuation, u')이 존재하면, 이를 물리적으로 차분 연산할 중간 
     미분 필드(DerivedField) 및 변동량 필드(FluctuationField)로 자동 변환하고 변수 테이블에 등록하여 확장합니다.
3. Stage 3 (의존성 분석 및 위상 정렬):
   - 필드 간의 선후 연산 관계를 분석하여 유향 의존성 그래프(Directed Dependency Graph)를 생성합니다.
   - 난류 통계의 특성상 평균 연산을 기준으로 "평균치 계산 전의 루프(Pass 1)"와 "평균치를 구한 후 변동량을 계산하는 루프(Pass 2)"로
     전체 연산 블록을 논리적으로 분할하고, 각각의 블록 내에서 올바른 순서로 계산되도록 위상 정렬(Topological Sort)을 수행합니다.
4. Stage 4 (수식 기호 최적화, Liveness 분석 및 Buffer Array Pooling):
   - SymPy 기호 수학 라이브러리를 이용하여 복잡한 3차원 수식을 대수적으로 간소화하고, 공통 부분 식 제거(CSE)를 적용하여 연산 비용(Flops)을 최적화합니다.
   - 격자 데이터가 거대하므로 모든 변수에 개별 3D 배열을 할당하면 메모리가 고갈됩니다. 이를 방지하기 위해 변수들의 
     생명 주기(Liveness Window)를 수학적으로 추적하고, 동적 메모리 풀을 구축하여 동시에 활성화되지 않는 임시 변수들이 
     공통의 제한된 버퍼 배열(xyzbuffer0, xyzbuffer1, ...)을 나누어 사용(Array Pooling)하도록 할당하여 메모리 사용량을 최소화합니다.
"""

import sys
from lark import Lark, Visitor, Transformer, v_args, Token
import warnings
from jinja2 import Template
import sympy
from sympy.printing.fortran import FCodePrinter

# --- Registries for SOLID (OCP) Compliance ---

class FunctionRegistry:
    """Registry to map mathematical DSL functions to SymPy constructors and LaTeX representations.

    This class enables SOLID Open-Closed Principle (OCP) compliance by decoupling core parsing
    logic from mathematical functions, allowing new functions to be added without modifying the parser.
    """

    def __init__(self):
        """Initializes FunctionRegistry with empty SymPy and LaTeX mappings."""
        self._sympy_registry = {}
        self._latex_registry = {}

    def register_sympy(self, name, sympy_builder):
        """Registers a handler to convert a DSL function to a SymPy representation.

        Args:
            name (str): The name of the DSL function.
            sympy_builder (callable): A callable mapping function arguments to a SymPy object.
        """
        self._sympy_registry[name] = sympy_builder

    def register_latex(self, name, latex_builder):
        """Registers a handler to convert a DSL function to a LaTeX math representation.

        Args:
            name (str): The name of the DSL function.
            latex_builder (callable): A callable mapping arguments to a LaTeX string.
        """
        self._latex_registry[name] = latex_builder

    def to_sympy(self, name, *args):
        """Converts a function call to its corresponding SymPy expression.

        Args:
            name (str): The name of the function.
            *args: Arguments to pass to the sympy builder.

        Returns:
            sympy.Expr: SymPy node representation of the function call.
        """
        if name in self._sympy_registry:
            return self._sympy_registry[name](*args)
        return sympy.Function(name)(*args)

    def to_latex(self, name, *args):
        """Converts a function call to its corresponding LaTeX representation.

        Args:
            name (str): The name of the function.
            *args (str): Already-formatted LaTeX strings of the arguments.

        Returns:
            str: The LaTeX math representation of the function call.
        """
        if name in self._latex_registry:
            return self._latex_registry[name](*args)
        b = ", ".join(args)
        if name.startswith("\\"):
            return r"{}{{({})}}".format(name, b)
        return r"\mathrm{{{}}}({})".format(name, b)


function_registry = FunctionRegistry()

# Register standard mathematical functions
function_registry.register_sympy("sqrt", lambda *args: sympy.sqrt(args[0]))
function_registry.register_sympy("abs", lambda *args: sympy.Abs(args[0]))
function_registry.register_sympy("log", lambda *args: sympy.log(args[0]))
function_registry.register_sympy("exp", lambda *args: sympy.exp(args[0]))
function_registry.register_sympy("rxn_rate", lambda *args: sympy.Function("rxn_rate")(args[0]))

function_registry.register_latex("sqrt", lambda *args: r"\sqrt{{{}}}".format(", ".join(args)))
function_registry.register_latex("abs", lambda *args: r"\left| {} \right|".format(", ".join(args)))
function_registry.register_latex(r"\log", lambda *args: r"\log{{({})}}".format(", ".join(args)))
function_registry.register_latex(r"\exp", lambda *args: r"\exp{{({})}}".format(", ".join(args)))
function_registry.register_latex(r"\omega", lambda *args: r"\omega{{({})}}".format(", ".join(args)))


class DifferentialOperatorRegistry:
    """Registry for spatial differential operators mapping operators to LaTeX representations.

    Handles default and custom differential operators (e.g., ddx, d2dy) used during derivation expansion.
    """

    def __init__(self):
        """Initializes DifferentialOperatorRegistry with empty operator mappings."""
        self._operators = {}

    def register(self, op_name, latex_symbol):
        """Registers a custom LaTeX representation for a given operator.

        Args:
            op_name (str): The differential operator name (e.g., 'ddx').
            latex_symbol (str): The corresponding LaTeX math symbol.
        """
        self._operators[op_name] = latex_symbol

    def get_latex_symbol(self, op_name):
        """Retrieves the LaTeX representation of a differential operator.

        Args:
            op_name (str): The name of the operator.

        Returns:
            str: The LaTeX code for the differential operator (e.g. '\\partial_{x}').
        """
        if op_name in self._operators:
            return self._operators[op_name]
        # Fallback to dynamic parsing matching original code
        fmt = r"\partial_{{{}}}"
        coord = op_name[-1] if op_name else ""
        return fmt.format(coord + coord if len(op_name) > 3 else coord)


differential_operator_registry = DifferentialOperatorRegistry()

# Register standard derivative operators
differential_operator_registry.register("ddx", r"\partial_{x}")
differential_operator_registry.register("d2dx", r"\partial_{xx}")
differential_operator_registry.register("ddy", r"\partial_{y}")
differential_operator_registry.register("d2dy", r"\partial_{yy}")
differential_operator_registry.register("ddz", r"\partial_{z}")
differential_operator_registry.register("d2dz", r"\partial_{zz}")


class FortranTemplateStore:
    """Static store containing Jinja2 code templates for generating Fortran source blocks.

    Attributes:
        REAL_ARRAY_LOOP (str): Standard 3D loop for calculating derived variables, integrating SymPy CSE blocks.
        FLUCTUATION_ARRAY_LOOP (str): Loop for calculating fluctuation fields (with averages subtracted).
        AVG_ARRAY_SUM (str): Accumulator loop to sum grid values for spatial averaging.
        AVG_ARRAY_DIVIDE (str): Global MPI MPI_ALLREDUCE and divide step to finalize statistics.
        FMT_DECL_SUBARRAY (str): MPI subarray-based I/O file handle and type variable declarations.
        FMT_INIT_SUBARRAY (str): Initialization block for creating MPI file handles and commit subarray types.
        FMT_FINAL_SUBARRAY (str): Cleanup code block to close and release MPI types and file handles.
        FMT_CALC_SUBARRAY (str): Actual high-performance parallel MPI file writing using subarray layout.
        FMT_DECL_LEGACY (str): Declarations for legacy, process-local buffer-based MPI exports.
        FMT_INIT_LEGACY (str): Allocation and initialization for legacy buffer arrays.
        FMT_FINAL_LEGACY (str): Finalization and deallocation for legacy buffer arrays.
        FMT_CALC_LEGACY (str): Data slicing and local buffer writing for legacy export routines.
        AVG_ARRAY_WRITE (str): Serial file writing structure for exporting 1D averaged data and derivatives.
    """
    REAL_ARRAY_LOOP = """
! {{ comment }}
{% if decls_str -%}
block
{{ decls_str | indent(4, True) }}
{%- endif %}
do k = 1, nzp
do j = 1, nyp
do i = 1, nxp
{% if assigns_str -%}
{{ assigns_str | indent(4, True) }}
    {{ array }}(i,j,k) = {{ rhs }}
{%- else -%}
    {{ array }}(i,j,k) = {{ rhs }}
{%- endif %}
end do
end do
end do
{% if decls_str -%}
end block
{%- endif %}
"""

    FLUCTUATION_ARRAY_LOOP = """
! {{ comment }}
do k = 1, nzp
do j = 1, nyp
do i = 1, nxp
{{ array }}(i,j,k) = {{ rhs }}
end do
end do
end do
"""

    AVG_ARRAY_SUM = """
do k = 1, nzp
do j = 1, nyp
do i = 1, nxp
{{ name }}(i) = {{ name }}(i) + {{ arrname }}
end do
end do
end do
"""

    AVG_ARRAY_DIVIDE = """
call MPI_ALLREDUCE(MPI_IN_PLACE, {{ name }}, nxp, MPI_REAL8, MPI_SUM, MPI_COMM_TASK, mpi_err)

{{ name }} = {{ name }} {{ dWeight }} / denum
"""

    FMT_DECL_SUBARRAY = """
! - file_handles and mpi_infos
integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_fh
integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_info
integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_filetype
"""

    FMT_INIT_SUBARRAY = """
! init subarray datatype for {{ field_name }}
block
    integer(4) :: sizes(3), subsizes(3), starts(3)
    call MPI_INFO_CREATE({{ field_name }}_info, mpi_err)
    call MPI_FILE_OPEN(MPI_COMM_TASK,'export-{{ field_name }}.dat',MPI_MODE_WRONLY+MPI_MODE_CREATE,{{ field_name }}_info,{{ field_name }}_fh,mpi_err)
    sizes = (/ nxp, nyp, nzp /)
    subsizes = (/ {{ len_xpts }}, {{ ye }} - {{ ys }} + 1, {{ ze }} - {{ zs }} + 1 /)
    starts = (/ {{ xs }} - 1, {{ ys }} - 1, {{ zs }} - 1 /)
    call MPI_TYPE_CREATE_SUBARRAY(3, sizes, subsizes, starts, MPI_ORDER_FORTRAN, MPI_REAL8, {{ field_name }}_filetype, mpi_err)
    call MPI_TYPE_COMMIT({{ field_name }}_filetype, mpi_err)
end block
"""

    FMT_FINAL_SUBARRAY = """
! finalize
call MPI_FILE_CLOSE({{ field_name }}_fh, mpi_err)
call MPI_INFO_FREE({{ field_name }}_info, mpi_err)
call MPI_TYPE_FREE({{ field_name }}_filetype, mpi_err)
"""

    FMT_CALC_SUBARRAY = """
! write to file via MPI Subarray
count = ({{ len_xpts }}) * ({{ ye }} - {{ ys }} + 1) * ({{ ze }} - {{ zs }} + 1)
offset = export_offset(fidx) * count * 8
call MPI_FILE_WRITE_AT({{ field_name }}_fh, offset, {{ work_array }}, 1, {{ field_name }}_filetype, mpi_status, mpi_err)
"""

    FMT_DECL_LEGACY = """
! - file_handles and mpi_infos
integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_fh
integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_info

! - buffer
real(real64), allocatable, dimension(:,:,:) :: {{ field_name }}_export_array
integer, allocatable, dimension(:) :: {{ field_name }}_xpts
"""

    FMT_INIT_LEGACY = """
! init
call MPI_INFO_CREATE({{ field_name }}_info, mpi_err)
call MPI_FILE_OPEN(MPI_COMM_TASK,'export-{{ field_name }}.dat',MPI_MODE_WRONLY+MPI_MODE_CREATE,{{ field_name }}_info,{{ field_name }}_fh,mpi_err)
allocate({{ field_name }}_export_array(1:{{ len_xpts }},{{ ys }}:{{ ye }},{{ zs }}:{{ ze }}), stat=ierr)
if (ierr /= 0) then
    write(0,*) 'Error: allocation of {{ field_name }}_export_array failed on process', myid
    call MPI_ABORT(MPI_COMM_TASK, 1, mpi_err)
end if
{{ field_name }}_export_array = 0.
allocate({{ field_name }}_xpts(1:{{ len_xpts }}), stat=ierr)
if (ierr /= 0) then
    write(0,*) 'Error: allocation of {{ field_name }}_xpts failed on process', myid
    call MPI_ABORT(MPI_COMM_TASK, 1, mpi_err)
end if
{{ xpts_init }}
"""

    FMT_FINAL_LEGACY = """
! finalize
call MPI_FILE_CLOSE({{ field_name }}_fh, mpi_err)
call MPI_INFO_FREE({{ field_name }}_info, mpi_err)
deallocate({{ field_name }}_export_array)
deallocate({{ field_name }}_xpts)
"""

    FMT_CALC_LEGACY = """
! copy to array for export
do k = {{ zs }}, {{ ze }}
do j = {{ ys }}, {{ ye }}
do i = 1, {{ len_xpts }}
{{ field_name }}_export_array(i,j,k) = {{ work_array }}({{ field_name }}_xpts(i),j,k)
end do
end do
end do

! write to file
count = ({{ len_xpts }}) * ({{ ye }} - {{ ys }} + 1) * ({{ ze }} - {{ zs }} + 1)
offset = export_offset(fidx) * count * 8
call MPI_FILE_WRITE_AT({{ field_name }}_fh, offset, {{ field_name }}_export_array, count, MPI_REAL8, mpi_status, mpi_err)
"""

    AVG_ARRAY_WRITE = """
real(real64), dimension(nxp) :: xbuffer
integer :: i

open (200, file="qEdge_X.dat")
write (200,*) output_header
do i=1,nxp
 write (200,'({{ num_args }}e20.10)') real(i)*hxp, {{ formatted_avglist }}
end do
close (200)

open (200, file="d1.dat")
{{ deriv1_lines }}
close (200)

open (200, file="d2.dat")
{{ deriv2_lines }}
close (200)
"""


class FortranCodeGenerator(object):
    """Visitor implementation that decouples code generation from domain AST node details.

    Implements a classic visitor dispatch pattern that matches class names of the Field objects
    to their code generation routines (e.g. generate_code, generate_decl, generate_alloc, generate_free).
    """

    def __init__(self, fdict):
        """Initializes FortranCodeGenerator with the variable registry dictionary.

        Args:
            fdict (dict): Dictionary mapping variable names to their corresponding FieldBase objects.
        """
        self.fdict = fdict

    def generate_code(self, field, alloc=None):
        """Dispatches visitor to generate actual calculation code for the field.

        Args:
            field (FieldBase): The field node to generate code for.
            alloc (dict, optional): Buffer allocation map for array pooling. Defaults to None.

        Returns:
            str: Generated Fortran statement(s) inside loop blocks.
        """
        method_name = 'visit_' + field.__class__.__name__ + '_code'
        visitor = getattr(self, method_name, self.generic_code)
        return visitor(field, alloc)

    def generate_decl(self, field):
        """Dispatches visitor to generate variable type declarations.

        Args:
            field (FieldBase): The field node to generate declaration for.

        Returns:
            str: Fortran variable declaration statement.
        """
        method_name = 'visit_' + field.__class__.__name__ + '_decl'
        visitor = getattr(self, method_name, self.generic_decl)
        return visitor(field)

    def generate_alloc(self, field):
        """Dispatches visitor to generate dynamic memory allocation statements.

        Args:
            field (FieldBase): The field node to allocate.

        Returns:
            str: Fortran allocate and initialization statements.
        """
        method_name = 'visit_' + field.__class__.__name__ + '_alloc'
        visitor = getattr(self, method_name, self.generic_alloc)
        return visitor(field)

    def generate_free(self, field):
        """Dispatches visitor to generate deallocation statements.

        Args:
            field (FieldBase): The field node to deallocate.

        Returns:
            str: Fortran deallocate statements.
        """
        method_name = 'visit_' + field.__class__.__name__ + '_free'
        visitor = getattr(self, method_name, self.generic_free)
        return visitor(field)

    def generate_avg(self, field):
        """Dispatches visitor to generate average accumulators.

        Args:
            field (FieldBase): The averaged field node.

        Returns:
            str: Fortran summation and normalization statements.
        """
        method_name = 'visit_' + field.__class__.__name__ + '_avg'
        visitor = getattr(self, method_name, self.generic_avg)
        return visitor(field)

    # --- Generic Fallbacks ---
    def generic_code(self, field, alloc=None):
        """Generic fallback for code generation.

        Args:
            field (FieldBase): The field node.
            alloc (dict, optional): Buffer allocation map. Defaults to None.

        Returns:
            str: Empty string.
        """
        return ""

    def generic_decl(self, field):
        """Generic fallback for declaration code generation.

        Args:
            field (FieldBase): The field node.

        Returns:
            str: Standard allocatable real 3D/1D array declaration.
        """
        real_array_decl = "real(real64), allocatable, dimension({1}) :: {0}"
        return real_array_decl.format(field.name, field.dim)

    def generic_alloc(self, field):
        """Generic fallback for dynamic allocation.

        Args:
            field (FieldBase): The field node.

        Returns:
            str: Fortran allocate statement.
        """
        return make_allocate(field.name, field.shape)

    def generic_free(self, field):
        """Generic fallback for deallocation.

        Args:
            field (FieldBase): The field node.

        Returns:
            str: Fortran deallocate statement.
        """
        real_array_free = "deallocate({})"
        return real_array_free.format(field.name)

    def generic_avg(self, field):
        """Generic fallback for average code generation.

        Args:
            field (FieldBase): The field node.

        Returns:
            str: Empty string.
        """
        return ""

    # --- Visit Methods ---

    def visit_FieldExporter_code(self, exporter, alloc=None):
        """Generates Fortran code for exporting fields using parallel MPI-IO.

        Args:
            exporter (FieldExporter): Exporter metadata.
            alloc (dict, optional): Buffer allocation map. Defaults to None.

        Returns:
            str: MPI writing code block.
        """
        exporter.params["work_array"] = exporter.parent.array
        if exporter.use_subarray:
            return Template(FortranTemplateStore.FMT_CALC_SUBARRAY).render(**exporter.params)
        else:
            return Template(FortranTemplateStore.FMT_CALC_LEGACY).render(**exporter.params)

    def visit_FieldExporter_decl(self, exporter):
        """Generates declarations for MPI-IO handles.

        Args:
            exporter (FieldExporter): Exporter metadata.

        Returns:
            str: Declarations block.
        """
        if exporter.use_subarray:
            return Template(FortranTemplateStore.FMT_DECL_SUBARRAY).render(**exporter.params)
        else:
            return Template(FortranTemplateStore.FMT_DECL_LEGACY).render(**exporter.params)

    def visit_FieldExporter_alloc(self, exporter):
        """Generates initialization code for MPI-IO types and file opens.

        Args:
            exporter (FieldExporter): Exporter metadata.

        Returns:
            str: Initialization block.
        """
        if exporter.use_subarray:
            return Template(FortranTemplateStore.FMT_INIT_SUBARRAY).render(**exporter.params)
        else:
            return Template(FortranTemplateStore.FMT_INIT_LEGACY).render(**exporter.params)

    def visit_FieldExporter_free(self, exporter):
        """Generates finalization code for MPI-IO types and file closes.

        Args:
            exporter (FieldExporter): Exporter metadata.

        Returns:
            str: MPI-IO finalization statements.
        """
        if exporter.use_subarray:
            return Template(FortranTemplateStore.FMT_FINAL_SUBARRAY).render(**exporter.params)
        else:
            return Template(FortranTemplateStore.FMT_FINAL_LEGACY).render(**exporter.params)

    def visit_Field_code(self, field, alloc=None):
        """Generates calculation code for normal derived fields with SymPy CSE optimization.

        Args:
            field (Field): The calculated field.
            alloc (dict, optional): Buffer allocation map. Defaults to None.

        Returns:
            str: Renders optimized 3D loop including CSE declarations.
        """
        field.array = alloc[field.name] if alloc else field.name

        opt = SympyOptimizer.get_instance(self.fdict)
        rhs, cse_decls, cse_assigns = opt.optimize_field(field.name, alloc)

        decls_str = "\n".join(cse_decls) if cse_decls else ""
        assigns_str = "\n".join(cse_assigns) if cse_assigns else ""

        calculation_code = Template(FortranTemplateStore.REAL_ARRAY_LOOP).render(
            comment=field.comment,
            decls_str=decls_str,
            assigns_str=assigns_str,
            array=field.array,
            rhs=rhs
        )

        export_code = self.generate_code(field.exporter) if field.export_on() else ""
        return calculation_code + export_code

    def visit_FluctuationField_code(self, field, alloc=None):
        """Generates calculation code for turbulence fluctuations.

        Args:
            field (FluctuationField): The fluctuation field.
            alloc (dict, optional): Buffer allocation map. Defaults to None.

        Returns:
            str: Renders 3D loop for calculating u' = u - <u_w>.
        """
        field.array = alloc[field.name] if alloc else field.name
        rhs = ExpToCode(self.fdict).transform(field.field.exp)

        if field.field.is_fluctuation():
            rhs = rhs.format(field.w)

        return Template(FortranTemplateStore.FLUCTUATION_ARRAY_LOOP).render(
            comment=field.comment,
            array=field.array,
            rhs=rhs
        )

    def visit_PrimaryField_code(self, field, alloc=None):
        """No code generation needed for primary inputs.

        Args:
            field (PrimaryField): The primary input.
            alloc (dict, optional): Buffer allocation map. Defaults to None.

        Returns:
            str: Comment statement.
        """
        return "! {} is read from file".format(field.name)

    def visit_PrimaryField_decl(self, field):
        """No declaration needed for primary inputs.

        Args:
            field (PrimaryField): The primary input.

        Returns:
            str: Comment statement.
        """
        return "! {} is read from file".format(field.name)

    def visit_PrimaryField_alloc(self, field):
        """No allocation needed for primary inputs.

        Args:
            field (PrimaryField): The primary input.

        Returns:
            str: Comment statement.
        """
        return "! {} is read from file".format(field.name)

    def visit_PrimaryField_free(self, field):
        """No deallocation needed for primary inputs.

        Args:
            field (PrimaryField): The primary input.

        Returns:
            str: Comment statement.
        """
        return "! {} is read from file".format(field.name)

    def visit_DerivedField_code(self, field, alloc=None):
        """Generates derivative calculation statements, calling Compact solver subroutines.

        Args:
            field (DerivedField): The derivative field.
            alloc (dict, optional): Buffer allocation map. Defaults to None.

        Returns:
            str: Fortran subroutine call string.
        """
        field.array = alloc[field.name] if alloc else field.name
        varray = alloc[field.v] if alloc else field.v
        return "call {0} ( {2}, {1} )".format(field.op, varray, field.array)

    def visit_AveragedField_code(self, field, alloc=None):
        """Generates local accumulation summation loop for spatial averaging.

        Args:
            field (AveragedField): The averaged variable field.
            alloc (dict, optional): Buffer allocation map. Defaults to None.

        Returns:
            str: Renders local sum code block.
        """
        arrname = self.fdict[field.tgt].array + "(i,j,k)"
        if field.weighted is not None:
            arrname = arrname + " * " + field.w.array + "(i,j,k)"

        return Template(FortranTemplateStore.AVG_ARRAY_SUM).render(name=field.name, arrname=arrname)

    def visit_AveragedField_avg(self, field):
        """Generates normalization and parallel reduction via MPI_ALLREDUCE.

        Args:
            field (AveragedField): The averaged variable field.

        Returns:
            str: Renders global reduce and normalize statement.
        """
        dWeight = (f"/ avg_{field.weighted}" if field.weighted else "")
        return Template(FortranTemplateStore.AVG_ARRAY_DIVIDE).render(name=field.name, dWeight=dWeight)

    def generate_write_avg(self, avglist):
        """Generates final module output writers for 1D spatial averages.

        Also generates first and second derivatives of the averaged data using
        ddx1d and d2dx1d subroutines.

        Args:
            avglist (list): List of averaged field names.

        Returns:
            str: Code block to write values to output files.
        """
        avgarr = "{}(i)"
        deriv1_avgarr = """call ddx1d ( xbuffer, {} ) ; write (200,*) xbuffer"""
        deriv2_avgarr = """call d2dx1d ( xbuffer, {} ) ; write (200,*) xbuffer"""

        num_args = len(avglist) + 1
        formatted_avglist = ", ".join(map(avgarr.format, avglist))
        deriv1_lines = "\n".join(map(deriv1_avgarr.format, avglist))
        deriv2_lines = "\n".join(map(deriv2_avgarr.format, avglist))

        write_avg = Template(FortranTemplateStore.AVG_ARRAY_WRITE).render(
            num_args=num_args,
            formatted_avglist=formatted_avglist,
            deriv1_lines=deriv1_lines,
            deriv2_lines=deriv2_lines
        )
        return write_avg


@v_args(inline=True)
class LarkToSympy(Transformer):
    """Transformer to convert Lark AST math nodes to SymPy symbolic expressions.

    Maps DSL expression trees into SymPy expressions to allow down-pipeline algebraic
    simplification, common subexpression elimination (CSE), and memory optimization.
    """

    def __init__(self, fdict):
        """Initializes LarkToSympy transformer.

        Args:
            fdict (dict): Dictionary mapping variable names to FieldBase objects.
        """
        self.fdict = fdict

    def number(self, numeral):
        """Converts number strings into SymPy Float objects.

        Args:
            numeral (Token/str): Numeric string from the AST.

        Returns:
            sympy.Float: SymPy floating point object.
        """
        return sympy.Float(float(numeral))

    def env(self, name):
        """Converts environment variable tokens to SymPy Symbol objects.

        Args:
            name (Token): Environment variable name token (prefixed with $).

        Returns:
            sympy.Symbol: SymPy symbol representation.
        """
        return sympy.Symbol(name.value)

    def paren(self, val):
        """Preserves precedence inside parentheses and returns the child expression.

        Args:
            val (sympy.Expr): Expression inside parentheses.

        Returns:
            sympy.Expr: The inner expression unchanged.
        """
        return val

    def var(self, name):
        """Maps variable name tokens to SymPy Symbol objects.

        Args:
            name (Token): Variable name token.

        Returns:
            sympy.Symbol: SymPy symbol representation.
        """
        return sympy.Symbol(name.value)

    def fluc(self, name):
        """Maps turbulence fluctuation variables (e.g., u') to a SymPy Symbol with '__prime' suffix.

        Args:
            name (Token): Variable name token representing fluctuation.

        Returns:
            sympy.Symbol: SymPy symbol with prime suffix identifier.
        """
        return sympy.Symbol(name.value + "__prime")

    def dnx(self, partial, b):
        """Maps spatial derivative operations (e.g., ddx(u)) to a single composite SymPy Symbol.

        This treats the derivative term (e.g., ddx_u) as an independent symbol.

        Args:
            partial (Token): Derivative operator token.
            b (Token): Variable name token being differentiated.

        Returns:
            sympy.Symbol: Composite derivative symbol representation.
        """
        signature = f"{partial.data}_{b.value}"
        return sympy.Symbol(signature)

    def icall(self, op, val):
        """Converts inline functions (e.g., sqr, pow3) to direct SymPy exponent expressions.

        Args:
            op (Token): Inline function operator token.
            val (sympy.Expr): The function argument expression.

        Returns:
            sympy.Expr: SymPy power expression.
        """
        if op.data == "sqr":
            return val**2
        elif op.data == "pow3":
            return val**3
        return val

    def fcall(self, *args):
        """Maps standard built-in functions or UDFs to their SymPy representations.

        Args:
            *args: Variable length argument list. The first argument is the function name Token,
                and the subsequent arguments are the function parameter expressions.

        Returns:
            sympy.Expr: SymPy function node or registered mathematical expression.
        """
        a = args[0]
        func_name = a.value if hasattr(a, 'value') else str(a)
        if func_name == "udf":
            return sympy.Function(a.value if hasattr(a, 'value') else str(a))(*args[1:])
        return function_registry.to_sympy(func_name, *args[1:])

    def neg(self, val):
        """Converts negation to SymPy unary negation.

        Args:
            val (sympy.Expr): The expression to negate.

        Returns:
            sympy.Expr: Negated SymPy expression.
        """
        return -val

    def add(self, a, b):
        """Converts addition to SymPy sum expression.

        Args:
            a (sympy.Expr): Left expression.
            b (sympy.Expr): Right expression.

        Returns:
            sympy.Expr: SymPy addition expression.
        """
        return a + b

    def sub(self, a, b):
        """Converts subtraction to SymPy difference expression.

        Args:
            a (sympy.Expr): Left expression.
            b (sympy.Expr): Right expression.

        Returns:
            sympy.Expr: SymPy subtraction expression.
        """
        return a - b

    def mul(self, a, b):
        """Converts multiplication to SymPy product expression.

        Args:
            a (sympy.Expr): Left expression.
            b (sympy.Expr): Right expression.

        Returns:
            sympy.Expr: SymPy product expression.
        """
        return a * b

    def div(self, a, b):
        """Converts division to SymPy division expression.

        Args:
            a (sympy.Expr): Left expression.
            b (sympy.Expr): Right expression.

        Returns:
            sympy.Expr: SymPy division expression.
        """
        return a / b

    def udf(self, a):
        """Maps user defined function names to string representations.

        Args:
            a (Token): Function token.

        Returns:
            str: Name of the user defined function.
        """
        return a.value

    log = lambda self: "log"
    exp = lambda self: "exp"
    sqrt = lambda self: "sqrt"
    abs = lambda self: "abs"
    rxn_rate = lambda self: "rxn_rate"


class ArrayFCodePrinter(FCodePrinter):
    """Custom SymPy printer that formats symbols as Fortran multidimensional array accesses.

    Transforms plain SymPy symbols into:
    - 3D grid accesses (e.g., var(i,j,k)) for spatial fields.
    - 1D array accesses (e.g., var(i)) for spatial averages.
    - Subtraction expressions (e.g., (u(i,j,k) - avg_u(i))) for fluctuation variables.
    """

    def __init__(self, settings=None, array_symbols=None, avg_symbols=None):
        """Initializes ArrayFCodePrinter with array settings and symbol catalogs.

        Args:
            settings (dict, optional): Printer settings configuration. Defaults to None.
            array_symbols (dict, optional): Maps 3D fields to their buffer array names. Defaults to None.
            avg_symbols (dict, optional): Maps averaged fields to their 1D arrays. Defaults to None.
        """
        settings = settings or {}
        settings.setdefault('source_format', 'free') # Default to free-form Fortran 95
        settings.setdefault('standard', 95)
        super().__init__(settings)
        self.array_symbols = array_symbols or {}
        self.avg_symbols = avg_symbols or {}

    def _print_Float(self, expr):
        """Prints Float constants with double precision (d0) suffix.

        Ensures precision is not degraded during Fortran compilation.

        Args:
            expr (sympy.Float): SymPy float node.

        Returns:
            str: Double precision literal string.
        """
        val = str(expr)
        if 'e' in val or 'E' in val:
            return val.replace('e', 'd').replace('E', 'd')
        if '.' not in val:
            return val + ".0d0"
        return val + "d0"

    def _print_Symbol(self, expr):
        """Maps a SymPy Symbol to a Fortran array expression.

        Args:
            expr (sympy.Symbol): SymPy Symbol to print.

        Returns:
            str: Formatted Fortran variable access or inline fluctuation calculation.
        """
        name = expr.name
        # 1. 3D grid array
        if name in self.array_symbols:
            return f"{self.array_symbols[name]}(i,j,k)"
        # 2. 1D averaged array
        if name in self.avg_symbols:
            return f"{self.avg_symbols[name]}(i)"
        # 3. Fluctuation substitution (e.g. u' -> u - <u_w>)
        if name.endswith("__prime"):
            base = name[:-7]
            arr = self.array_symbols.get(base, base)
            avg_name = f"avg_{base}"
            printed_avg = f"{self.avg_symbols.get(avg_name, avg_name)}(i)"
            return f"({arr}(i,j,k) - {printed_avg})"
        return name

    def _print_Function(self, expr):
        """Prints custom non-standard functions securely in Fortran output.

        Fallback logic for rxn_rate, udf, etc.

        Args:
            expr (sympy.Function): SymPy Function node.

        Returns:
            str: Standard Fortran call syntax for the function.
        """
        try:
            return super()._print_Function(expr)
        except Exception:
            args = ", ".join(self.doprint(arg) for arg in expr.args)
            return f"{expr.func.__name__}({args})"


class SympyOptimizer:
    """Manages algebra optimization and Common Subexpression Elimination (CSE) via SymPy.

    This optimization engine:
    1. Expands intermediate temporary variables recursively.
    2. Simplifies arithmetic terms (fraction cancellation, trigonometric expansions).
    3. Runs CSE to pull duplicate sub-operations out into loop local scalar variables,
       significantly reducing FLOPS and memory bandwidth requirements.
    """
    _instance = None

    @classmethod
    def get_instance(cls, fdict):
        """Retrieves or instantiates the singleton optimizer.

        Args:
            fdict (dict): Current variable field registry mapping name -> FieldBase object.

        Returns:
            SympyOptimizer: The active singleton instance.
        """
        if cls._instance is None or cls._instance.fdict is not fdict:
            cls._instance = cls(fdict)
        return cls._instance

    def __init__(self, fdict):
        """Initializes SympyOptimizer registry caches.

        Args:
            fdict (dict): Variable field registry.
        """
        self.fdict = fdict
        self.sympy_cache = {}
        self.exported_fields = set(
            name for name, f in fdict.items()
            if hasattr(f, 'attr') and f.attr.get('export')
        )
        self.averaged_targets = set()
        self.avg_names = set()

    def set_averaged(self, averaged_dict):
        """Sets the targets and names of averaged variables.

        Args:
            averaged_dict (dict): Dictionary mapping average variable names to AveragedField objects.
        """
        self.averaged_targets = {a.target for a in averaged_dict.values()}
        self.avg_names = set(averaged_dict.keys())

    def get_sympy_expr(self, name):
        """Recursively builds and caches a fully substituted SymPy Expression for a variable.

        Avoids expanding boundary fields (PrimaryField, DerivedField representing spatial derivatives,
        AveragedField, and FluctuationField) to preserve the grid boundaries. Expands other temporary variables.

        Args:
            name (str): Variable name.

        Returns:
            sympy.Expr: SymPy node representing the fully-expanded mathematical expression.
        """
        if name in self.sympy_cache:
            return self.sympy_cache[name]

        field = self.fdict[name]

        if hasattr(field, 'prime') and field.prime:
            expr = sympy.Symbol(name)
            self.sympy_cache[name] = expr
            return expr

        if hasattr(field, 'op'): # DerivedField (ddx, etc.)
            expr = sympy.Symbol(name)
            self.sympy_cache[name] = expr
            return expr

        if hasattr(field, 'weighted'): # AveragedField
            expr = sympy.Symbol(name)
            self.sympy_cache[name] = expr
            return expr

        if hasattr(field, 'field') and hasattr(field, 'w'): # FluctuationField
            expr = sympy.Symbol(name)
            self.sympy_cache[name] = expr
            return expr

        transformer = LarkToSympy(self.fdict)
        expr = transformer.transform(field.exp)

        # Recursively substitute intermediate variables
        expanded_expr = expr
        changed = True
        while changed:
            changed = False
            free_syms = list(expanded_expr.free_symbols)
            sub_dict = {}
            for sym in free_syms:
                sym_name = sym.name
                if sym_name in self.fdict:
                    f = self.fdict[sym_name]
                    is_derived_field = hasattr(f, 'op')
                    is_averaged_field = hasattr(f, 'weighted')
                    is_primary_field = hasattr(f, 'prime') and f.prime
                    is_exported = sym_name in self.exported_fields
                    is_averaged_target = sym_name in self.averaged_targets

                    if not (is_derived_field or is_averaged_field or is_primary_field or is_exported or is_averaged_target):
                        sub_dict[sym] = self.get_sympy_expr(sym_name)
                        changed = True

            if sub_dict:
                expanded_expr = expanded_expr.subs(sub_dict)

        self.sympy_cache[name] = expanded_expr
        return expanded_expr

    def calculate_flops_and_heavy(self, expr):
        """Measures computational cost (FLOPS and heavy operators) for a SymPy expression.

        Useful for logging optimization statistics.

        Args:
            expr (sympy.Expr): The expression to evaluate.

        Returns:
            tuple: (flops_count, heavy_operators_count).
        """
        flops = 0
        heavy = 0
        for node in sympy.preorder_traversal(expr):
            if isinstance(node, sympy.Add):
                flops += len(node.args) - 1
            elif isinstance(node, sympy.Mul):
                flops += len(node.args) - 1
            elif isinstance(node, sympy.Pow):
                base, exp = node.args
                if exp == 0.5 or exp == -0.5:
                    flops += 10 # Square root/inv square root weight
                    heavy += 1
                elif exp == -1:
                    flops += 4  # Division weight
                    heavy += 1
                elif isinstance(exp, sympy.Integer):
                    val = abs(int(exp))
                    if val > 1:
                        flops += val - 1
                else:
                    flops += 10
                    heavy += 1
            elif isinstance(node, (sympy.Derivative, sympy.Function)):
                name = node.func.__name__
                if name == 'sqrt':
                    flops += 10
                    heavy += 1
                elif name in ('exp', 'log', 'sin', 'cos', 'tan', 'rxn_rate'):
                    flops += 10
                    heavy += 1
                elif name == 'Abs':
                    flops += 1
                else:
                    flops += 10
                    heavy += 1
        return flops, heavy

    def count_3d_loads(self, expr, three_d_arrays):
        """Counts the total 3D array memory read accesses in the expression.

        Args:
            expr (sympy.Expr): The expression to analyze.
            three_d_arrays (set): Names of active 3D arrays.

        Returns:
            int: The memory load count.
        """
        count = 0
        for node in sympy.preorder_traversal(expr):
            if isinstance(node, sympy.Symbol) and node.name in three_d_arrays:
                count += 1
        return count

    def optimize_field(self, name, alloc=None):
        """Optimizes a physical field expression and extracts Common Subexpressions.

        Applies SymPy simplification and CSE, printing optimization reports to stderr.

        Args:
            name (str): The name of the field to optimize.
            alloc (dict, optional): Buffer allocation mapping. Defaults to None.

        Returns:
            tuple: (rhs_code, cse_declarations, cse_assignments) where:
                rhs_code (str): The final right hand side Fortran expression.
                cse_declarations (list of str): Code strings to declare local CSE scalars.
                cse_assignments (list of str): Code strings to calculate CSE values.
        """
        expr = self.get_sympy_expr(name)

        three_d_arrays = {
            k for k, v in self.fdict.items()
            if hasattr(v, 'dim') and v.dim == ':,:,:'
        }

        # Optimization metrics before
        before_flops, before_heavy = self.calculate_flops_and_heavy(expr)
        before_loads = self.count_3d_loads(expr, three_d_arrays)

        # Simplify expression
        simplified_expr = sympy.simplify(expr)
        simplified_expr = sympy.cancel(simplified_expr)

        array_symbols = {}
        for k, v in self.fdict.items():
            if hasattr(v, 'array') and v.array:
                array_symbols[k] = v.array
            elif alloc and k in alloc:
                array_symbols[k] = alloc[k]
            else:
                array_symbols[k] = k

        avg_symbols = {k: k for k in getattr(self, 'avg_names', [])}
        printer = ArrayFCodePrinter(array_symbols=array_symbols, avg_symbols=avg_symbols)

        # Perform Common Subexpression Elimination
        replacements, reduced_exprs = sympy.cse(simplified_expr)
        reduced_expr = reduced_exprs[0]

        # Optimization metrics after
        after_flops = 0
        after_heavy = 0
        after_loads = 0

        for temp_var, temp_expr in replacements:
            f_val, h_val = self.calculate_flops_and_heavy(temp_expr)
            after_flops += f_val
            after_heavy += h_val
            after_loads += self.count_3d_loads(temp_expr, three_d_arrays)

        f_val, h_val = self.calculate_flops_and_heavy(reduced_expr)
        after_flops += f_val
        after_heavy += h_val
        after_loads += self.count_3d_loads(reduced_expr, three_d_arrays)

        def pct_str(before, after):
            if before == 0:
                return "0.0%" if after == 0 else "+inf%"
            diff = after - before
            pct = (diff / before) * 100
            return f"{pct:+.1f}%"

        flops_pct = pct_str(before_flops, after_flops)
        heavy_pct = pct_str(before_heavy, after_heavy)
        loads_pct = pct_str(before_loads, after_loads)

        if after_flops < before_flops * 0.5 or after_loads < before_loads * 0.5:
            est_speedup = "Highly significant"
        elif after_flops < before_flops or after_loads < before_loads:
            est_speedup = "Moderate"
        else:
            est_speedup = "Minimal / Already optimal"

        sys.stderr.write(f"\n[SymPy Optimizer Report: {name}]\n")
        sys.stderr.write(f"- Floating Point Ops : {before_flops} -> {after_flops} ({flops_pct})\n")
        sys.stderr.write(f"- Heavy Ops (Div/Sqrt): {before_heavy} -> {after_heavy} ({heavy_pct})\n")
        sys.stderr.write(f"- 3D Array Mem Reads  : {before_loads} -> {after_loads} ({loads_pct})\n")
        sys.stderr.write(f"=> Estimated Speedup in loop: {est_speedup}\n\n")

        cse_decls = []
        cse_assigns = []

        if replacements:
            for temp_var, temp_expr in replacements:
                cse_decls.append(f"real(real64) :: {temp_var}")
                cse_assigns.append(f"{temp_var} = {printer.doprint(temp_expr)}")

        rhs = printer.doprint(reduced_expr)

        return rhs, cse_decls, cse_assigns


class CollectDefinitions(Visitor):
    """Visitor that walks the Lark AST to collect variable and assignment declarations.

    Gathers the primary fields (direct file inputs), derived calculation definitions,
    and average directives (statistical averages) to build the compiler's initial model.
    """

    def __init__(self, primary, derived, averaged):
        """Initializes CollectDefinitions visitor.

        Args:
            primary (set): Set to accumulate primary input field names.
            derived (dict): Dictionary to accumulate derived Field objects.
            averaged (dict): Dictionary to accumulate average variable targets.
        """
        self.primary = primary
        self.derived = derived
        self.averaged = averaged

    def varlist(self, tree):
        """Parses list of primary inputs declared in bracket syntax (e.g. [u, v, w]).

        Args:
            tree (Tree): Lark AST node for varlist.
        """
        for v in tree.children:
            self.primary.add(v.value)
            self.derived[v.value] = PrimaryField(v.value, self.derived)

    def assign_var(self, tree):
        """Parses variable assignment statements and registers corresponding Field objects.

        Args:
            tree (Tree): Lark AST node for assignment.
        """
        if len(tree.children) > 2:
            lval, lattr, rval = tree.children
        else:
            lval, rval = tree.children
            lattr = None

        attr_dict = {}

        if lattr is not None:
            for t in lattr.children:
                k, v = t.children
                attr_dict[k.value] = v.value

        if lval.value in self.derived:
            raise ValueError("duplicate definition of " + lval)
        self.derived[lval.value] = Field(lval.value, attr_dict, rval, self.derived)

    def assign_avg_var(self, tree):
        """Parses average directives and registers target variables for spatial averaging.

        Args:
            tree (Tree): Lark AST node for average declaration.
        """
        w = tree.children[0]
        targets = tree.children[1:]

        if (not w.children) or (w.children[0] is None):
            self.averaged[None] = set([x.value for x in targets])
        else:
            self.averaged[w.children[0].value] = set([x.value for x in targets])


class ExpInspector(Visitor):
    """Visitor to inspect mathematical AST nodes to extract dependencies and attributes.

    Finds the list of dependent variables, checks if the expression references
    turbulence fluctuation variables (primed variables), and extracts derivative operators
    such as spatial first and second order derivatives (e.g. ddx, d2dy).
    """

    def __init__(self):
        """Initializes ExpInspector state."""
        self.fluctuation = False
        self.dep = set([])
        self.deriv = set([])

    @classmethod
    def inspect(cls, tree):
        """Inspects the given AST tree and returns extracted attributes.

        Args:
            tree (Tree): Lark AST subtree.

        Returns:
            tuple: (has_fluctuation, dependencies_set, derivatives_set).
        """
        self = cls()
        return self(tree)

    def __call__(self, tree):
        """Executes the inspection traversal.

        Args:
            tree (Tree): Lark AST subtree.

        Returns:
            tuple: (has_fluctuation, dependencies_set, derivatives_set).
        """
        self.visit(tree)
        return self.fluctuation, self.dep, self.deriv

    def fluc(self, tree):
        """Processes a fluctuation node.

        Args:
            tree (Tree): Fluctuation AST node.
        """
        self.fluctuation = True
        self.dep.add(tree.children[0].value)

    def var(self, tree):
        """Processes a variable reference node.

        Args:
            tree (Tree): Variable reference AST node.
        """
        self.dep.add(tree.children[0].value)

    def dnx(self, tree):
        """Processes a spatial derivative node, registering intermediate fields.

        Args:
            tree (Tree): Derivative operation AST node.
        """
        op, v = tree.children
        deriv = "{}_{}".format(op.data, v.value)
        self.dep.add(deriv)
        self.deriv.add((op.data, v.value))


@v_args(inline=True)
class ExpToLatex(Transformer):
    """Transformer to compile the DSL mathematical AST into a LaTeX math string.

    Converts operations, variables, derivatives, and functions into LaTeX syntax
    (e.g., partial derivative notation, fraction blocks, and bracket sizing)
    to generate mathematical report files.
    """

    def __init__(self, fdict):
        """Initializes ExpToLatex transformer.

        Args:
            fdict (dict): Dictionary mapping variable names to FieldBase objects.
        """
        self.fdict = fdict

    def arithmatic_rooted(self, name):
        """Checks if a variable's root operation is arithmetic.

        Args:
            name (str): Variable name.

        Returns:
            bool: True if arithmetic-rooted.
        """
        try:
            exproot = self.fdict[name].exp.data
        except AttributeError:
            exproot = "something_11fasq2afa3rfzsaerqw23"

        return ((exproot == "add") or (exproot == "sub") or
                (exproot == "mul") or (exproot == "div"))

    def parenthise(self, name):
        """Formats a variable name, wrapping it in parentheses if it has lower priority operators.

        Args:
            name (str): Variable name.

        Returns:
            str: Parenthesized LaTeX equation representation.
        """
        try:
            latex = self.fdict[name].latex
            latex_given = self.fdict[name].latex_given
        except KeyError:
            warnings.warn(name + " is not found")
            latex = r"\mathrm{{{}}}".format(name)
            latex_given = None

        if self.arithmatic_rooted(name) and (latex_given is None):
            latex = "(" + latex + ")"

        return latex

    def number(self, numeral):
        """Returns LaTeX literal for numbers.

        Args:
            numeral (Token): Number token.

        Returns:
            str: LaTeX representation.
        """
        return numeral

    def env(self, name):
        """Returns LaTeX literal for environment variables.

        Args:
            name (Token): Environment variable name token.

        Returns:
            str: LaTeX representation.
        """
        return r"\mathrm{{{}}}".format(name.value)

    def paren(self, name):
        """Formats parentheses.

        Args:
            name (str): Inner LaTeX content.

        Returns:
            str: Parenthesized string.
        """
        return "({})".format(str(name))

    def var(self, name):
        """Formats variables in LaTeX.

        Args:
            name (Token): Variable token.

        Returns:
            str: LaTeX variable representation.
        """
        return self.parenthise(name.value)

    def fluc(self, name):
        """Formats fluctuation prime (u'') variables in LaTeX.

        Args:
            name (Token): Variable token.

        Returns:
            str: LaTeX prime variable string.
        """
        return self.parenthise(name.value) + "''"

    def dnx(self, partial, b):
        """Formats spatial derivatives in LaTeX (e.g. \\partial_x(u)).

        Args:
            partial (Token): Derivative operator token.
            b (Token): Variable token.

        Returns:
            str: LaTeX partial derivative expression.
        """
        fmt = r"\partial_{{{}}}"
        coord = partial.data[-1]
        op = fmt.format(coord + coord if len(partial.data) > 3 else coord)

        signature = "{}_{}".format(partial.data, b.value)

        try:
            eq = self.fdict[signature].latex
        except KeyError:
            eq = op + self.parenthise(b.value)
            warnings.warn(signature + " is not found: " + eq)

        return eq

    def icall(self, a, b):
        """Formats inline functions (sqr, pow3) in LaTeX.

        Args:
            a (Token): Inline function operator.
            b (str): LaTeX representation of argument.

        Returns:
            str: LaTeX representation.
        """
        if a.data == "sqr":
            fcode = "({0})^2".format(b)
        elif a.data == "pow3":
            fcode = "({0})^3".format(b)
        else:
            fcode = "({0})".format(b)
        return fcode

    def fcall(self, *args):
        """Formats function calls in LaTeX.

        Args:
            *args: Variable length arguments. The first argument is function name token.

        Returns:
            str: LaTeX function representation.
        """
        a = args[0]
        func_name = a.value if hasattr(a, 'value') else str(a)
        return function_registry.to_latex(func_name, *args[1:])

    def neg(self, b):
        """Formats negation.

        Args:
            b (str): Inner LaTeX expression.

        Returns:
            str: LaTeX negated expression.
        """
        fcode = "(-{})".format(b)
        return fcode

    def add(self, a, b):
        """Formats addition.

        Args:
            a (str): Left operand.
            b (str): Right operand.

        Returns:
            str: LaTeX addition expression.
        """
        fcode = "{} + {}".format(a, b)
        return fcode

    def sub(self, a, b):
        """Formats subtraction.

        Args:
            a (str): Left operand.
            b (str): Right operand.

        Returns:
            str: LaTeX subtraction expression.
        """
        fcode = "{} - {}".format(a, b)
        return fcode

    def mul(self, a, b):
        """Formats multiplication.

        Args:
            a (str): Left operand.
            b (str): Right operand.

        Returns:
            str: LaTeX multiplication expression.
        """
        fcode = "{} {}".format(a, b)
        return fcode

    def div(self, a, b):
        """Formats division.

        Args:
            a (str): Left operand.
            b (str): Right operand.

        Returns:
            str: LaTeX division expression.
        """
        fcode = "{} / {}".format(a, b)
        return fcode

    log = lambda self: "\log"
    exp = lambda self: "\exp"
    sqrt = lambda self: "sqrt"
    abs = lambda self: "abs"
    rxn_rate = lambda self: "\omega"
    udf = lambda self, a: a.value


@v_args(inline=True)
class ExpToCode(Transformer):
    """Transformer to compile the DSL mathematical AST directly into Fortran array expressions.

    Formats variable accesses with spatial index wrappers (e.g., `(i,j,k)` or `(i)`) and
    converts inline functions or standard math calls.
    """

    def __init__(self, fdict):
        """Initializes ExpToCode transformer.

        Args:
            fdict (dict): Dictionary mapping variable names to FieldBase objects.
        """
        self.fdict = fdict

    def number(self, numeral):
        """Converts float literal numbers.

        Args:
            numeral (Token): Number token.

        Returns:
            str: Floating point string.
        """
        return str(float(numeral))

    def env(self, name):
        """Converts environment variables.

        Args:
            name (Token): Environment variable name token.

        Returns:
            str: Variable name.
        """
        return name.value

    def paren(self, name):
        """Wraps in parentheses.

        Args:
            name (str): Inner expression.

        Returns:
            str: Parenthesized string.
        """
        return "({})".format(str(name))

    def var(self, name):
        """Formats a variable access with spatial grid index (i,j,k).

        Args:
            name (Token): Variable token.

        Returns:
            str: Fortran array access.
        """
        try:
            arrname = self.fdict[name.value].array
        except KeyError:
            arrname = name.value

        return arrname + "(i,j,k)"

    def fluc(self, name):
        """Formats a fluctuation prime variable access.

        Saves index reference for inline replacement of averages.

        Args:
            name (Token): Fluctuation variable token.

        Returns:
            str: Fortran inline fluctuation subtraction formula.
        """
        try:
            arrname = self.fdict[name.value].array
        except KeyError:
            arrname = name.value

        fmt = "({0}(i,j,k) - {{0}}avg_{1}(i))"
        return fmt.format(arrname, name.value)

    def dnx(self, partial, b):
        """Formats spatial derivatives as array accesses.

        Args:
            partial (Token): Derivative operator.
            b (Token): Variable token.

        Returns:
            str: Fortran array access for derivative.
        """
        signature = "{}_{}".format(partial.data, b.value)
        try:
            arrname = self.fdict[signature].array
        except KeyError:
            arrname = signature

        return arrname + "(i,j,k)"

    def icall(self, a, b):
        """Formats inline functions (sqr, pow3) to Fortran multiplication.

        Args:
            a (Token): Inline function operator.
            b (str): Argument code.

        Returns:
            str: Fortran math expression.
        """
        if a.data == "sqr":
            fcode = "(({0})*({0}))".format(b)
        elif a.data == "pow3":
            fcode = "(({0})*({0})*({0}))".format(b)
        else:
            fcode = "({0})".format(b)
        return fcode

    def fcall(self, *args):
        """Formats standard math function calls.

        Args:
            *args: Variable length argument list. First argument is the function token.

        Returns:
            str: Fortran function call string.
        """
        a = args[0]
        b = ", ".join(args[1:])
        fcode = "( {} ( {} ) )".format(a, b)
        return fcode

    def neg(self, b):
        """Formats negation.

        Args:
            b (str): Code string to negate.

        Returns:
            str: Negated code string.
        """
        fcode = "( - {} )".format(b)
        return fcode

    def add(self, a, b):
        """Formats addition.

        Args:
            a (str): Left operand code.
            b (str): Right operand code.

        Returns:
            str: Fortran addition code.
        """
        fcode = "( {} + {} )".format(a, b)
        return fcode

    def sub(self, a, b):
        """Formats subtraction.

        Args:
            a (str): Left operand code.
            b (str): Right operand code.

        Returns:
            str: Fortran subtraction code.
        """
        fcode = "( {} - {} )".format(a, b)
        return fcode

    def mul(self, a, b):
        """Formats multiplication.

        Args:
            a (str): Left operand code.
            b (str): Right operand code.

        Returns:
            str: Fortran multiplication code.
        """
        fcode = "( {} * {} )".format(a, b)
        return fcode

    def div(self, a, b):
        """Formats division.

        Args:
            a (str): Left operand code.
            b (str): Right operand code.

        Returns:
            str: Fortran division code.
        """
        fcode = "( {} / {} )".format(a, b)
        return fcode

    log = lambda self: "log"
    exp = lambda self: "exp"
    sqrt = lambda self: "sqrt"
    abs = lambda self: "abs"
    rxn_rate = lambda self: "rxn_rate"
    udf = lambda self, a: a.value


def make_allocate(name, shape, init_zero=True):
    """Fortran 배열을 동적 할당하고 예외 발생 시 프로세스를 안전하게 폭파시키는 할당 코드를 작성해 줍니다.

    Args:
        name (str): 할당할 배열의 이름.
        shape (str): 할당 크기 형태 (예: 'nxp,nyp,nzp').
        init_zero (bool, optional): True인 경우 0.0d0으로 초기화 구문을 덧붙입니다. Defaults to True.

    Returns:
        str: Fortran dynamic allocation code block.
    """
    alloc_str = f"allocate({name}({shape}), stat=ierr)\n"
    alloc_str += f"if (ierr /= 0) then\n"
    alloc_str += f"    write(0,*) 'Error: allocation of {name} failed on process', myid\n"
    alloc_str += f"    call MPI_ABORT(MPI_COMM_TASK, 1, mpi_err)\n"
    alloc_str += f"end if"
    if init_zero:
        alloc_str += f"\n{name} = 0."
    return alloc_str


class DependencyNode(object):
    """의존성 관계 분석을 담당하는 추상 인터페이스 및 노드 클래스입니다.

    Compiler pipeline stages use this graph to resolve topological sort order.
    """

    def __init__(self, name, fdict):
        """Initializes DependencyNode.

        Args:
            name (str): Variable name.
            fdict (dict): Global variable registry dictionary.
        """
        self.name = name
        self.fdict = fdict
        self.dep = set([])
        self.fluc = False

    def depends_on(self, a):
        """Checks if this node depends directly on the given variable.

        Args:
            a (str): Target variable name.

        Returns:
            bool: True if directly dependent.
        """
        return a in self.dep

    def is_fluctuation(self):
        """Checks if this node is a fluctuation variable.

        Returns:
            bool: True if fluctuation variable.
        """
        return self.fluc

    def checkFluctuation(self):
        """본 변수 혹은 의존하고 있는 하위 기호들 중에 변동량(Fluctuation) 관련 계산이 개입되어 있는지
        상향식으로 전파 추적하는 재귀 메서드입니다.

        Returns:
            set: Set of variable names that require fluctuation calculations.
        """
        fset = set([])
        for d in map(self.fdict.get, self.dep):
            fset.update(d.checkFluctuation())
        if self.is_fluctuation() or len(fset) > 0:
            fset.add(self.name)
        return fset

    def depClosure(self):
        """해당 변수를 계산하기 위해 선행 계산되어야 하는 모든 하위 변수 노드들을 
        재귀적으로 타고 내려가 총합 폐쇄 집합(Closure Set)으로 묶어 반환합니다.

        Returns:
            set: Complete transitive dependency set.
        """
        fset = set(self.dep)
        for d in self.dep:
            fset.update(self.fdict[d].depClosure())
        return fset


class FieldBase(DependencyNode):
    """모든 물리 필드 객체의 최상위 기본 클래스로서 공통 데이터 구조를 정의합니다."""

    def __init__(self, name, fdict):
        """Initializes FieldBase properties.

        Args:
            name (str): Variable name.
            fdict (dict): Global variable registry dictionary.
        """
        super(FieldBase, self).__init__(name, fdict)
        self.array = name
        self.prime = False
        self.shape = "nxp,nyp,nzp"
        self.dim = ":,:,:"

    def export_on(self):
        """Checks if disk exporting is enabled for this field.

        Returns:
            bool: Always False for the base field class.
        """
        return False

    def __repr__(self):
        """String representation of the field (returns variable name)."""
        return self.name


class FieldExporter(object):
    """물리 필드 데이터를 병렬 분산 디스크 시스템으로 직접 추출(Export)하는 고성능 MPI-IO 서브루틴 블록을 
    정의하는 도메인 데이터 클래스입니다.
    """
    mpi_io_decl = """
! field exporter common
integer(kind=MPI_OFFSET_KIND)  :: offset
"""

    def __init__(self, name, attr, parent):
        """Initializes FieldExporter with MPI configurations.

        Args:
            name (str): Field exporter name.
            attr (dict): Attributes dictionary containing slice coordinates (e.g. xs, xe).
            parent (Field): Parent field object owning this exporter.
        """
        self.name = name
        self.attr = attr
        self.parent = parent

        self.params = dict(attr)

        self.params.setdefault("xs", 1)
        self.params.setdefault("xe", "nxp")
        self.params.setdefault("ys", 1)
        self.params.setdefault("ye", "nyp")
        self.params.setdefault("zs", 1)
        self.params.setdefault("ze", "nzp")

        self.params.setdefault("field_name", self.name)

        self.params.setdefault("len_xpts", f"({self.params['xe']} - {self.params['xs']} + 1)")
        
        fmt_xpts_init = '''
do i = {{ xs }}, {{ xe }}
{{ field_name }}_xpts(i-{{ xs }}+1) = i
end do
'''
        self.params.setdefault("xpts_init", Template(fmt_xpts_init).render(**self.params))

        try:
            # Sampling at listed x coordinates
            fmt_xpts_init_list = "{{ field_name }}_xpts = (/ {{ list_xpts }} /)"

            raw_xpts = self.params["xpts"]
            int_xpts = list(map(int, raw_xpts.split()))
            len_xpts = len(int_xpts)
            self.params["len_xpts"] = len_xpts
            self.params["list_xpts"] = ",".join(map(str, int_xpts))
            self.params["xpts_init"] = Template(fmt_xpts_init_list).render(**self.params)
        except KeyError:
            pass

        self.use_subarray = ("xpts" not in self.params)


class Field(FieldBase):
    """일반 대입 계산 변수를 관리하며, 3차원 격자점 루프 코드 생성을 담당하는 핵심 클래스입니다.
    SymPy 최적화 및 CSE 적용 코드를 루프 본문에 결합합니다.
    """

    def __init__(self, name, attr, exp, fdict):
        """Initializes Field properties and inspects expression details.

        Args:
            name (str): Calculated variable name.
            attr (dict): Attribute configuration tags.
            exp (Tree): Mathematical equation AST subtree.
            fdict (dict): Global variable registry dictionary.
        """
        super(Field, self).__init__(name, fdict)
        self.attr = attr
        self.exp = exp
        self.fluc, self.dep, self.derivs = ExpInspector.inspect(exp)
        self.comment = self.name + " = " + ExpToCode(self.fdict).transform(self.exp)

        self.latex_given = self.attr.get("latex")
        if self.latex_given is None:
            self.latex = ExpToLatex(self.fdict).transform(self.exp)
        else:
            self.latex = self.latex_given

        self.exporter = None
        try:
            if self.attr["export"]:
                self.exporter = FieldExporter(self.name, self.attr, self)
        except KeyError:
            pass

    def export_on(self):
        """Checks if file exporting is enabled.

        Returns:
            bool: True if an exporter is configured.
        """
        return self.exporter is not None


class FluctuationField(FieldBase):
    """물리 필드의 난류 변동 성분(Fluctuation, u' = u - <u_w>)을 계산하기 위한 변수 클래스입니다.
    수식 내의 u' 기호를 평균량과의 차이 수식으로 팽창하여 할당합니다.
    """

    def __init__(self, w, field, fset, fdict):
        """Initializes FluctuationField.

        Args:
            w (str/None): Weight variable name or average identifier.
            field (str): Base variable name (e.g. 'u').
            fset (set): Active fluctuation dependency names.
            fdict (dict): Global variable registry.
        """
        super(FluctuationField, self).__init__(self.id(w, field), fdict)

        if w is not None:
            self.w = w + "_"
        else:
            self.w = ""

        self.field = fdict[field]
        self.dep = self.field.dep - fset
        for df in self.field.dep & fset:
            self.dep.add(self.id(w, df))

        self.comment = ExpToCode(self.fdict).transform(self.field.exp)

        if self.field.is_fluctuation():
            self.comment = self.comment.format(self.w)

        self.comment = self.name + " = " + self.comment

    @classmethod
    def id(cls, w, field):
        """Generates dynamic fluctuation field naming key.

        Args:
            w (str/None): Average weight suffix.
            field (str): Base field name.

        Returns:
            str: Generated composite fluctuation variable name.
        """
        if w:
            name = "{}____{}_avg".format(field, w)
        else:
            name = "{}____avg".format(field)
        return name


class PrimaryField(FieldBase):
    """격자 정보나 외부 물리계 수치 데이터(u, v, w, T 등) 파일에서 사전에 로드하여 
    메모리에 상주하는 기본 원본 입력 필드 클래스입니다. 
    자체 계산 루프나 동적 할당 코드를 직접 생성하지 않습니다.
    """

    def __init__(self, name, fdict):
        """Initializes PrimaryField.

        Args:
            name (str): Input field name.
            fdict (dict): Variable registry dictionary.
        """
        super(PrimaryField, self).__init__(name, fdict)
        self.derivs = set([])
        self.prime = True
        self.latex = name
        self.latex_given = None


class DerivedField(FieldBase):
    """수치 공간 미분(ddx, d2dy 등)을 수행하여 계산되는 유도 필드 클래스입니다.
    Fortran 수치 차분 패키지 서브루틴(Compact.f90 에 구현된 dfnonp, dfp 등)의 
    동적 호출 코드를 출력합니다.
    """

    def __init__(self, op, v, fdict):
        """Initializes DerivedField derivative node.

        Args:
            op (str): Differential operator name (e.g. 'ddx').
            v (str): Variable being differentiated (e.g. 'u').
            fdict (dict): Global variable registry dictionary.
        """
        name = "{}_{}".format(op, v)
        super(DerivedField, self).__init__(name, fdict)
        self.op = op
        self.v = v
        self.dep = set([v])

        partial = differential_operator_registry.get_latex_symbol(op)
        self.latex = partial + "(" + fdict[v].latex + ")"


class AveragedField(FieldBase):
    """특정 물리 필드를 격자의 동질 차원(예: X 방향 1D선상 평균)에 대해 
    공간 통계 평균(Average) 연산을 수행하는 1차원 필드 클래스입니다.
    """

    @classmethod
    def id(cls, w, tgt):
        """Generates average field naming key.

        Args:
            w (str/None): Weight parameter.
            tgt (str): Target field variable name.

        Returns:
            str: Generated averaged variable key name.
        """
        if w:
            return "{}_avg_{}".format(w, tgt)
        else:
            return "avg_{}".format(tgt)

    def __init__(self, w, tgt, fdict):
        """Initializes AveragedField and populates dependency closure.

        Args:
            w (str/None): Average weights.
            tgt (str): Target variable to average.
            fdict (dict): Global variable registry dictionary.
        """
        name = self.id(w, tgt)
        super(AveragedField, self).__init__(name, fdict)
        self.shape = "nxp" # Y, Z dimensions averaged out, leaving X-dimension array of size nxp
        self.dim = ":"
        self.target = tgt

        tfield = fdict[tgt]
        self.fset = tfield.checkFluctuation()

        self.latex = r"\left\langle {} \right\rangle".format(tfield.latex)

        if not self.fset:
            self.tgt = tgt
            self.dep.add(tgt)
        else:
            ftgt = FluctuationField.id(w, tgt)
            self.tgt = ftgt
            self.dep.add(ftgt)

        self.weighted = w
        if w:
            self.w = fdict[w]
            self.dep.add(w)
            self.latex += ("_{{{}}}".format(w))

    def isWeighted(self):
        """Checks if average has a weighted density variable.

        Returns:
            bool: True if weighted.
        """
        return self.weighted is not None

    def pass1(self):
        """Checks if average variable can be computed in loop Pass 1.

        Returns:
            bool: True if Pass 1 average.
        """
        return not self.pass2()

    def pass2(self):
        """Checks if average variable requires Pass 2 (dependent on fluctuation).

        Returns:
            bool: True if Pass 2 average.
        """
        return len(self.fset) > 0


class CompilationContext(object):
    """Holds compilation pipeline state across parsing, resolution, and optimization stages.

    Attributes:
        primary (set of str): Names of primary input fields.
        derived (dict of str -> FieldBase): Calculated, derivative, and fluctuation fields.
        averaged (dict of str -> AveragedField): Spatial averaged variables.
        dependency (dict of str -> set of str): DAG representing direct dependencies.
        pass1 (list of str): Topologically sorted variables for average-precursor calculation.
        pass2 (list of str): Topologically sorted variables for post-average calculations.
        avg1 (set of AveragedField): Averaged variables calculated in Pass 1.
        avg2 (set of AveragedField): Averaged variables calculated in Pass 2.
        alloc1 (dict of str -> str): Buffer pooling maps for Pass 1.
        alloc2 (dict of str -> str): Buffer pooling maps for Pass 2.
        narr (int): Maximum number of shared XYZ buffer arrays needed.
    """

    def __init__(self):
        """Initializes an empty CompilationContext."""
        self.primary = set()
        self.derived = {}
        self.averaged = {}
        self.dependency = {}
        self.pass1 = []
        self.pass2 = []
        self.avg1 = set()
        self.avg2 = set()
        self.alloc1 = {}
        self.alloc2 = {}
        self.narr = 0


calc_grammar = """
    ?varlist: "[" [NAME ("," NAME)*] "]"

    ?start: statement*

    ?statement: NAME [ attr_list ] "=" sum -> assign_var
              | avg "{" [NAME ("," NAME)*] "}" -> assign_avg_var
              | varlist

    attr_list: "(" [attr_pair ("," attr_pair)*] ")"

    ?attr_pair: NAME "=" BOOL
              | NAME "=" INT
              | NAME "=" ESCAPED_STRING

    ?sum: product
        | sum "+" product   -> add
        | sum "-" product   -> sub

    ?product: atom
        | product "*" atom  -> mul
        | product "/" atom  -> div

    ?atom: NUMBER           -> number
         | "-" atom         -> neg
         | NAME             -> var
         | NAME "'"         -> fluc
         | "$" NAME         -> env
         | "(" sum ")"      -> paren
         | inlinefunc "(" sum ")" -> icall
         | mathfunc "(" sum ("," sum)* ")" -> fcall
         | derivative "(" NAME ")" -> dnx

    avg: "avg" [NAME]

    ?inlinefunc: "sqr" -> sqr
               | "pow3" -> pow3

    ?mathfunc: "log" -> log
             | "exp" -> exp
             | "sqrt" -> sqrt
             | "abs" -> abs
             | "rxn_rate" -> rxn_rate
             | "$" NAME   -> udf

    ?derivative: "ddx"  -> ddx
               | "d2dx" -> d2dx
               | "ddy" -> ddy
               | "d2dy" -> d2dy
               | "ddz" -> ddz
               | "d2dz" -> d2dz

    %import common.CNAME -> NAME
    %import common.NUMBER
    %import common.ESCAPED_STRING
    %import common.INT
    %import common.WS

    BOOL: "true" | "false"

    COMMENT: /#.*/

    %ignore COMMENT
    %ignore WS
"""

def tok_to_bool(tok):
    "Convert the value of `tok` from string to bool, while maintaining line number & column."
    return Token.new_borrow_pos(tok.type, tok.value == "true", tok)

def tok_to_int(tok):
    "Convert the value of `tok` from string to int, while maintaining line number & column."
    return Token.new_borrow_pos(tok.type, int(tok), tok)

def tok_to_str(tok):
    "Convert the value of `tok` from string to string, while maintaining line number & column."
    return Token.new_borrow_pos(tok.type, tok.value.strip('"'), tok)


class ParserStage(object):
    """Compiler pipeline Stage 1: Parses DSL specifications and extracts initial definitions.

    Reads raw DSL specifications and uses Lark parser to build the AST.
    Populates primary inputs, derived variables, and averaged variable lists.
    """

    def __init__(self):
        """Initializes ParserStage with calc_grammar."""
        self.parser = Lark(calc_grammar,
                           parser='lalr',
                           lexer_callbacks={
                               'ESCAPED_STRING': tok_to_str,
                               'INT': tok_to_int,
                               'BOOL': tok_to_bool
                           })

    def execute(self, terms_raw, ctx):
        """Executes Stage 1 parser.

        Args:
            terms_raw (str): Raw DSL term specification string.
            ctx (CompilationContext): Active compilation context.
        """
        tree = self.parser.parse(terms_raw)
        CollectDefinitions(ctx.primary, ctx.derived, ctx.averaged).visit(tree)


class DerivativeExpansionStage(object):
    """Compiler pipeline Stage 2: Expands differential operators and fluctuation terms.

    Finds derivative expressions (e.g., ddx, d2dy) and fluctuation identifiers (u'),
    instantiates DerivedField and FluctuationField objects, registers them in
    the variable registry, and builds initial DAG dependencies.
    """

    def execute(self, ctx):
        """Executes Stage 2 derivative and fluctuation expansion.

        Args:
            ctx (CompilationContext): Active compilation context.
        """
        # 1. 계산식 내에 존재하는 고차 차분 미분항(ddx, d2dy 등)을 찾아 중간 DerivedField로 등록
        dset = set()
        for k, v in ctx.derived.items():
            dset.update(v.derivs)

        for tup in dset:
            a = DerivedField(tup[0], tup[1], ctx.derived)
            ctx.derived[a.name] = a

        # 2. 통계 물리량 평균화 대상 변수들을 AveragedField 구조체로 구성하고, 
        #    변동량 계산이 필요한 경우 FluctuationField로 등록
        averaged_raw = ctx.averaged
        ctx.averaged = {}
        for w, tgts in averaged_raw.items():
            for t in tgts:
                a = AveragedField(w, t, ctx.derived)
                ctx.averaged[a.name] = a
                # 평균 편차가 동반된 항들에 대해 FluctuationField 생성
                for ff in a.fset:
                    b = FluctuationField(w, ff, a.fset, ctx.derived)
                    ctx.derived[b.name] = b

        # 3. 프로그램 내 모든 필드 간의 1차 의존 관계 그래프(Dependency Graph) 추출
        ctx.dependency = {}
        for k, v in ctx.derived.items():
            ctx.dependency[k] = v.dep
        for k, v in ctx.averaged.items():
            ctx.dependency[k] = v.dep


class DependencyResolutionStage(object):
    """Compiler pipeline Stage 3: Resolves data dependencies and orders calculations.

    Splits loops into:
    - Pass 1 (calculating variables required prior to averages).
    - Pass 2 (calculating fluctuations and statistics after averages are resolved).
    Performs topological sorting to ensure correctness.
    """

    def execute(self, ctx):
        """Executes Stage 3 dependency resolution and topological sort.

        Args:
            ctx (CompilationContext): Active compilation context.
        """
        # 이름 충돌 방지 검증
        assert set(ctx.derived.keys()).isdisjoint(ctx.averaged.keys())

        # Pass 1과 Pass 2 대상 평균화 변수 논리 분할
        ctx.avg1 = set(filter(AveragedField.pass1, ctx.averaged.values()))
        ctx.avg2 = set(filter(AveragedField.pass2, ctx.averaged.values()))

        # Pass 1 위상 정렬: Pass 1 평균 변수들의 연산에 관여하는 모든 종속 관계를 수집하여 정렬
        pass1calc = set(map(repr, ctx.avg1))
        for x in ctx.avg1:
            pass1calc.update(x.depClosure())
        ctx.pass1 = self.sort_vars_new(ctx.dependency, pass1calc - ctx.primary)

        # Pass 2 위상 정렬: Pass 2 평균 변수(변동 연산 연계)들의 연산에 관여하는 종속성을 정렬
        pass2calc = set(map(repr, ctx.avg2))
        for x in ctx.avg2:
            pass2calc.update(x.depClosure())
        ctx.pass2 = self.sort_vars_new(ctx.dependency, pass2calc - ctx.primary)

    def calc_size(self, dependency, ordered, remaining):
        """Calculates topological dependency metrics for ordering weight sorting.

        Args:
            dependency (dict): Graph mapping variable names to dependency sets.
            ordered (set): Topologically ordered variable names.
            remaining (set): Unsorted remaining variable names.

        Returns:
            int: The degree of active connection impact.
        """
        count = 0
        dep_union = set()
        for v in remaining:
            dep_union |= set(dependency[v])
        for v in ordered:
            if v in dep_union:
                count += 1
        return count

    def sort_vars_new(self, dependency, group):
        """Performs topological sort using an impact-weight heuristic.

        Minimizes variable lifetime durations to optimize array reuse.

        Args:
            dependency (dict): Graph mapping variable names to dependency sets.
            group (set): Set of variable names to sort.

        Returns:
            list: Topologically sorted list of variable names.
        """
        order = []
        remain = list(group)
        remain.sort()

        while len(remain) > 0:
            candidate = []
            for v in remain:
                if set(dependency[v]).isdisjoint(remain):
                    candidate.append(v)

            impact = {}
            size0 = self.calc_size(dependency, set(order), set(remain))

            for v in candidate:
                impact[v] = self.calc_size(dependency, set(order) | set([v]), set(remain) - set([v])) - size0
                
            candidate.sort(key=impact.get)
            order.append(candidate[0])
            remain.remove(candidate[0])

        return order


class SympyOptimizationStage(object):
    """Compiler pipeline Stage 4: Optimizes mathematical equations and shares memory buffers.

    This stage:
    1. Invokes the SympyOptimizer to apply algebraic simplification and CSE.
    2. Updates variable dependencies to reflect optimized equations.
    3. Runs liveness window analysis to map multiple non-overlapping temporary variables
       to a limited set of shared XYZ buffers (array pooling) to prevent RAM exhaustion.
    """

    def execute(self, ctx):
        """Executes Stage 4 mathematical optimization and array buffer sharing.

        Args:
            ctx (CompilationContext): Active compilation context.
        """
        # 1. SymPy 수식 최적화 엔진 초기화
        opt = SympyOptimizer.get_instance(ctx.derived)
        opt.set_averaged(ctx.averaged)

        # 2. SymPy가 대입식 치환(Substitution) 과정에서 제거한 불필요한 의존성 관계를 
        #    의존성 그래프에 즉각 반영하여 실제 계산을 위한 의존성 체인을 슬림하게 정리합니다.
        updated_dependency = {}
        for name, dep_set in ctx.dependency.items():
            if name in ctx.derived and isinstance(ctx.derived[name], Field):
                expr = opt.get_sympy_expr(name)
                # 실제 정리된 SymPy 식에 잔존하는 자유 기호 명칭들만 추출
                free_sym_names = {sym.name for sym in expr.free_symbols}
                valid_deps = {dep for dep in free_sym_names if dep in ctx.derived or dep in ctx.primary}
                updated_dependency[name] = valid_deps
            else:
                updated_dependency[name] = dep_set
        ctx.dependency = updated_dependency

        self.array_name = "xyzbuffer{}"

        # 3. Pass 1 및 Pass 2 연산 순서 배열들에 대해 각각 버퍼 공유 매핑(Pooling) 수행
        narr1, alloc1 = (self.allocate_arr(ctx, ctx.pass1))
        narr2, alloc2 = (self.allocate_arr(ctx, ctx.pass2))

        # 전체 프로그램에서 필요한 동적 공유 3D 버퍼 배열의 최대 크기 설정
        ctx.narr = max(narr1, narr2)

        ctx.alloc1 = alloc1
        ctx.alloc2 = alloc2

    def liveness(self, ctx, l1, g):
        """Analyzes variable lifetimes to identify overlap intervals.

        Constructs a liveness matrix where matrix[i, j] is True if variable i
        is still live (in memory) at step j.

        Args:
            ctx (CompilationContext): Active compilation context.
            l1 (list): Topologically sorted variable names.
            g (dict): Graph mapping variable names to dependency sets.

        Returns:
            np.ndarray: Boolean liveness matrix of shape (len(l1), len(l1)).
        """
        import numpy as np
        img = np.zeros((len(l1), len(l1)))
        for i, v in enumerate(l1):
            for j in range(i, len(l1)):
                img[i,i:j+1] = img[i,i:j+1] + (1 if v in g[l1[j]] else 0)
        return img > 0

    def allocate_arr(self, ctx, l):
        """Performs liveness-based memory buffer pooling.

        Assigns variables with disjoint active lifetimes to share the same
        allocated `xyzbufferN` 3D arrays.

        Args:
            ctx (CompilationContext): Active compilation context.
            l (list): Topologically sorted list of variable names.

        Returns:
            tuple: (max_buffers_needed, var_to_buffer_mapping) where:
                max_buffers_needed (int): Maximum buffers required concurrently.
                var_to_buffer_mapping (dict): Map of variable names to their pooled buffer arrays.
        """
        import numpy as np
        dg = ctx.dependency
        mask = self.liveness(ctx, l, dg)
        try:
            narr = mask.astype(int).sum(axis=0).max()
        except ValueError:
            narr = 0

        array_pool = set([self.array_name.format(i) for i in range(narr)])
        livesets = [set([])] + [set(np.asarray(l)[row]) for row in mask.T]
        var2arr = { p : p for p in ctx.primary }

        for i, (s0, s1) in enumerate(zip(livesets[:-1], livesets[1:])):
            array_pool.update(map(var2arr.get, s0 - s1))
            for new in s1 - s0:
                var2arr[new] = array_pool.pop()

        return narr, var2arr


class FortranProgramWriter(object):
    """Renders the final, compiled Fortran 95 post-processing modules.

    Uses Jinja2 templates to combine calculations, averages, declarations,
    reductions, subarray parallel IO writers, and pooled dynamic array buffers.
    """

    def write(self, ctx):
        """Generates the code module and prints it directly to standard output.

        Args:
            ctx (CompilationContext): The compiled context containing sorted equations and pooling allocations.
        """
        from resources.m_template import mod_form

        allvar = dict(ctx.derived)
        allvar.update(ctx.averaged)

        generator = FortranCodeGenerator(allvar)

        # 평균 변수들을 순서대로 분배
        set1 = sorted([a.name for a in filter(AveragedField.pass1, ctx.averaged.values())])
        set2 = sorted([a.name for a in filter(AveragedField.pass2, ctx.averaged.values())])

        # 외부 디스크 파일 익스포트 활성화 여부 확인
        set_export_on = list(filter(lambda x: x.export_on(), ctx.derived.values()))

        ffmt = 'logical, parameter :: pass2_required={}'
        declf = ffmt.format('.true.' if len(set2) > 0 else '.false.')

        hfmt = 'character (len = *), parameter :: output_header="{}"'
        declh = hfmt.format(" ".join(["x"] + set1 + set2))

        # 공용 Pooling 3D 버퍼 배열의 선언/할당 코드 생성
        declarr, allocarr, freearr = self.work_array_codes(ctx.narr)

        # 1차원 평균 물리량 배열들의 선언/할당 코드 생성
        declavg  = "\n".join(generator.generate_decl(ctx.averaged[v]) for v in sorted(ctx.averaged))
        allocavg = "\n".join(generator.generate_alloc(ctx.averaged[v]) for v in sorted(ctx.averaged))
        freeavg  = "\n".join(generator.generate_free(ctx.averaged[v]) for v in sorted(ctx.averaged))

        # 병렬 파일 쓰기(MPI Subarray)를 위한 MPI 리소스 선언/할당 코드 생성
        decl_export  = "\n".join(generator.generate_decl(v.exporter)  for v in set_export_on)
        alloc_export = "\n".join(generator.generate_alloc(v.exporter) for v in set_export_on)
        free_export  = "\n".join(generator.generate_free(v.exporter)  for v in set_export_on)

        # Pass 1과 Pass 2 루프 내부 본문 연산 코드들을 생성 (각자 버퍼 맵 alloc1, alloc2 적용)
        sub_calc1 = "\n".join(generator.generate_code(allvar[v], ctx.alloc1) for v in ctx.pass1 if v in ctx.averaged or v in ctx.alloc1)
        sub_calc2 = "\n".join(generator.generate_code(allvar[v], ctx.alloc2) for v in ctx.pass2 if v in ctx.averaged or v in ctx.alloc2)

        # 평균 누적 연산 코드 생성
        sub_avg1 = "\n".join(generator.generate_avg(allvar[v]) for v in set1)
        sub_avg2 = "\n".join(generator.generate_avg(allvar[v]) for v in set2)

        # 파일 쓰기 루틴 코드 생성
        sub_write_avg = generator.generate_write_avg(set1+set2)

        md = {}
        md["module_name"] = "terms"
        md["module_data"] = "\n".join((declf, declh, declavg, FieldExporter.mpi_io_decl, decl_export, declarr))
        md["module_init"] = "\n".join((allocavg, alloc_export, allocarr))
        md["module_finalize"] = "\n".join((freeavg, free_export, freearr))
        md["module_pass1"]     = sub_calc1
        md["module_pass1_avg"] = sub_avg1
        md["module_pass2"]     = sub_calc2
        md["module_pass2_avg"] = sub_avg2
        md["module_write_result"] = sub_write_avg

        print(Template(mod_form).render(**md))

    def work_array_codes(self, narr):
        """Generates dynamic memory helper statements for shared xyzbuffer buffers.

        Args:
            narr (int): Number of buffers needed.

        Returns:
            tuple: (declarations_string, allocations_string, deallocations_string).
        """
        array_name = "xyzbuffer{}"
        array_names = [array_name.format(i) for i in range(narr)]

        real_array_decl  = "real(real64), allocatable, dimension(:,:,:) :: {0}"
        decl  = "\n".join([real_array_decl.format(v) for v in array_names])
        alloc = "\n".join([make_allocate(v, "nxp,nyp,nzp") for v in array_names])
        free  = "\n".join(["deallocate({})".format(v) for v in array_names])

        return decl, alloc, free


class ReportWriter(object):
    """Outputs a compilation analysis summary in JSON format.

    Creates `ir2.py` containing topological sort order and graph relationships
    to verify pipeline execution properties.
    """

    def write(self, ctx):
        """Writes the IR details to `ir2.py`.

        Args:
            ctx (CompilationContext): The compiled context.
        """
        import json
        dg = {k:list(v) for k,v in ctx.dependency.items()}

        with open("ir2.py", "w") as irf:
            print("g = ", json.dumps(dg, indent=4), file=irf)
            print("l1 = ", json.dumps(ctx.pass1, indent=4), file=irf)
            print("l2 = ", json.dumps(ctx.pass2, indent=4), file=irf)
            print("avg1 = ", json.dumps(list(map(repr, ctx.avg1)), indent=4), file=irf)
            print("avg2 = ", json.dumps(list(map(repr, ctx.avg2)), indent=4), file=irf)


class LatexWriter(object):
    """Outputs the compiled average physical quantities LaTeX equations as a Python dictionary.

    Prints the mapped string representation of the dictionary directly to stdout.
    """

    def write(self, ctx):
        """Writes the LaTeX equation definitions dictionary to standard output.

        Args:
            ctx (CompilationContext): The compiled context.
        """
        latex_lines = ["{"]
        for avg in ctx.averaged.values():
            latex_lines.append('  "{}" : r"${}$",'.format(avg.name, avg.latex))
        latex_lines.append("}")
        print("\n".join(latex_lines))