"""1D evaluator for the compiled backend.

All component functions live in ``trspecfit.functions.energy`` as the
single source of truth.  For 1D evaluation, parameters are plain
scalars (no ``(n_time, 1)`` broadcasting needed).
"""

from __future__ import annotations

import numpy as np

from trspecfit.graph_ir import (
    ScheduledPlan1D,
    _eval_expr_scalar,
    _evaluate_profile_expr_values,
    _evaluate_profile_sample_values,
    _evaluate_scheduled_op_1d,
)

# ---------------------------------------------------------------------------
# Core 1D evaluator
# ---------------------------------------------------------------------------


#
def evaluate_1d(plan: ScheduledPlan1D, theta: np.ndarray) -> np.ndarray:
    """Evaluate the compiled 1D model at optimizer parameters *theta*.

    Parameters
    ----------
    plan
        Immutable compiled execution schedule from ``schedule_1d``.
    theta
        ``(n_opt,)`` optimizer parameter vector.  Order must match
        ``plan.opt_param_names``.

    Returns
    -------
    ndarray
        ``(n_energy,)`` model spectrum.

    Raises
    ------
    ValueError
        If ``len(theta) != len(plan.opt_indices)``.
    """

    if len(theta) != len(plan.opt_indices):
        raise ValueError(
            f"theta length {len(theta)} does not match "
            f"plan.opt_indices length {len(plan.opt_indices)}"
        )

    # 1a. Copy parameter values -> scratch
    values = plan.param_values_init.copy()

    # 1b. Write optimizer params
    values[plan.opt_indices] = theta

    # 1c. Resolve expressions in topological order
    for i in range(plan.n_expressions):
        target = int(plan.expr_target_indices[i])
        values[target] = _eval_expr_scalar(
            plan.expr_instructions[plan.expr_indptr[i] : plan.expr_indptr[i + 1]],
            values,
        )

    profile_sample_values = _evaluate_profile_sample_values(
        plan.aux_axis,
        values,
        plan.profile_sample_base_indices,
        plan.profile_sample_component_indptr,
        plan.profile_component_func_ids,
        plan.profile_component_param_indptr,
        plan.profile_component_param_indices,
    )
    profile_expr_values = _evaluate_profile_expr_values(
        values,
        profile_sample_values,
        plan.n_params,
        plan.profile_expr_instructions,
        plan.profile_expr_indptr,
    )

    # 2. Component evaluation
    energy = plan.energy
    result = plan.cached_result.copy()
    peak_sum = plan.cached_peak_sum.copy()

    for op_idx in range(plan.n_ops):
        if plan.op_is_constant[op_idx]:
            continue
        kind = int(plan.op_kinds[op_idx])
        start = int(plan.op_param_indptr[op_idx])
        end = int(plan.op_param_indptr[op_idx + 1])
        needs_spectrum = bool(plan.op_needs_spectrum[op_idx])
        is_pre = bool(plan.op_is_pre_spectrum[op_idx])

        component = _evaluate_scheduled_op_1d(
            energy,
            kind,
            plan.op_param_source_kinds[start:end],
            plan.op_param_indices[start:end],
            values,
            profile_sample_values,
            profile_expr_values,
            peak_sum,
            needs_spectrum=needs_spectrum,
            is_profiled=bool(plan.op_is_profiled[op_idx]),
            n_aux=plan.n_aux,
        )

        result += component
        if is_pre:
            peak_sum += component

    return result