src/cost_model.py

"""
FusionOps Cost Model
Computes subgraph execution latency using the roofline model.
Handles tiling, split-K, data reuse, and memory transfers.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional

from .models import (
    Action, Config, Graph, OpType, ScheduleState, Tensor, TensorRole,
)


@dataclass
class TensorClassification:
    """How each tensor relates to a subgraph."""
    boundary_inputs: dict[int, str]   # tensor_id -> role_in_op ("pointwise", "lhs", "rhs")
    boundary_outputs: set[int]        # tensor_ids that leave the subgraph
    ephemeral: set[int]               # tensor_ids internal to subgraph
    resident_inputs: set[int]         # tensor_ids already in fast memory


@dataclass
class TileGeometry:
    """Tiling structure for a subgraph execution."""
    num_tiles_w: int
    num_tiles_h: int
    total_spatial_tiles: int
    num_k_steps: int
    total_iterations: int
    output_w: int  # actual output tensor width
    output_h: int  # actual output tensor height
    reduction_dim: int  # K for MatMul, 0 for pure Pointwise


@dataclass
class SubgraphLatencyResult:
    """Full result of latency computation."""
    total_latency: float
    tile_latencies: list[float]
    working_set: float
    is_valid: bool
    error: Optional[str] = None


def classify_tensors(
    graph: Graph,
    op_ids: list[int],
    state: ScheduleState,
) -> TensorClassification:
    """Classify every tensor touched by this subgraph."""
    op_id_set = set(op_ids)

    # All tensors produced within this subgraph
    produced_internally = set()
    for oid in op_ids:
        op = graph.get_op(oid)
        produced_internally.update(op.output_tensor_ids)

    # All tensors consumed within this subgraph
    consumed_internally = set()
    for oid in op_ids:
        op = graph.get_op(oid)
        consumed_internally.update(op.input_tensor_ids)

    # Boundary inputs: consumed but not produced internally
    boundary_input_ids = consumed_internally - produced_internally

    # Classify boundary inputs by their role
    boundary_inputs: dict[int, str] = {}
    for oid in op_ids:
        op = graph.get_op(oid)
        for tid in op.input_tensor_ids:
            if tid in boundary_input_ids:
                if op.op_type == OpType.MATMUL:
                    # First input is LHS, second is RHS
                    idx = op.input_tensor_ids.index(tid)
                    role = "lhs" if idx == 0 else "rhs"
                else:
                    role = "pointwise"
                boundary_inputs[tid] = role

    # Separate resident inputs (already in fast memory)
    resident_inputs = set()
    non_resident_inputs: dict[int, str] = {}
    for tid, role in boundary_inputs.items():
        if tid in state.tensors_in_fast_memory:
            resident_inputs.add(tid)
        else:
            non_resident_inputs[tid] = role

    # Classify produced tensors as ephemeral or boundary output.
    # A tensor is ephemeral if it is produced AND consumed within the subgraph
    # (even if it also has external consumers - those are served by recomputation).
    # A tensor is a boundary output if it is produced by the subgraph and:
    #   - NOT consumed by any op within the subgraph, OR
    #   - Is a graph output (must eventually reach slow memory)
    # But for cost purposes, even graph outputs that are retained don't get evicted now.
    boundary_outputs = set()
    ephemeral = set()

    for tid in produced_internally:
        consumed_by_internal_op = tid in consumed_internally

        if consumed_by_internal_op:
            # This tensor flows between ops inside the subgraph = ephemeral
            # It may also have external consumers, but those will be served
            # by recomputation or a different subgraph
            ephemeral.add(tid)
        else:
            # Not consumed internally at all = pure output of the subgraph
            boundary_outputs.add(tid)

    # Special case: tensors that are ephemeral but also graph outputs
    # or have external consumers still need to be trackable.
    # But for cost purposes, they are NOT evicted from this subgraph
    # unless they are in the boundary_outputs set.
    # If caller wants to retain an ephemeral tensor, it must appear 
    # in boundary_outputs. Check if any retained tensor is ephemeral
    # and promote it to boundary output (it needs to materialize).
    # Actually, retained tensors that are ephemeral don't make sense
    # in the current model - you can only retain outputs.

    return TensorClassification(
        boundary_inputs=non_resident_inputs,
        boundary_outputs=boundary_outputs,
        ephemeral=ephemeral,
        resident_inputs=resident_inputs,
    )


def compute_tile_geometry(
    graph: Graph,
    op_ids: list[int],
    config: Config,
) -> TileGeometry:
    """Compute the tiling structure for a subgraph."""
    # Determine output tensor dimensions
    # For a subgraph, all ops share the same spatial tiling
    # Use the output of the last op (or the subgraph's boundary output)
    output_w = 0
    output_h = 0
    reduction_dim = 0

    for oid in op_ids:
        op = graph.get_op(oid)
        for tid in op.output_tensor_ids:
            t = graph.get_tensor(tid)
            output_w = max(output_w, t.width)
            output_h = max(output_h, t.height)

        if op.op_type == OpType.MATMUL:
            # Reduction dim is the shared dimension between LHS cols and RHS rows
            # LHS: input[0], RHS: input[1]
            # For MatMul: LHS is (H x K), RHS is (K x W), output is (H x W)
            # The reduction dimension K = width of LHS = height of RHS
            lhs_tensor = graph.get_tensor(op.input_tensor_ids[0])
            reduction_dim = max(reduction_dim, lhs_tensor.width)

    num_tiles_w = math.ceil(output_w / config.w)
    num_tiles_h = math.ceil(output_h / config.h)
    total_spatial_tiles = num_tiles_w * num_tiles_h

    if reduction_dim > 0:
        num_k_steps = math.ceil(reduction_dim / config.k)
    else:
        num_k_steps = 1

    return TileGeometry(
        num_tiles_w=num_tiles_w,
        num_tiles_h=num_tiles_h,
        total_spatial_tiles=total_spatial_tiles,
        num_k_steps=num_k_steps,
        total_iterations=total_spatial_tiles * num_k_steps,
        output_w=output_w,
        output_h=output_h,
        reduction_dim=reduction_dim,
    )


def compute_working_set(
    graph: Graph,
    tensor_class: TensorClassification,
    config: Config,
    op_ids: list[int],
    state: ScheduleState,
) -> float:
    """Compute peak working set for memory validation."""
    ws = 0.0

    # Determine if we have split-K
    has_matmul = any(
        graph.get_op(oid).op_type == OpType.MATMUL for oid in op_ids
    )

    # Compute reduction dimension for k-step count
    reduction_dim = 0
    for oid in op_ids:
        op = graph.get_op(oid)
        if op.op_type == OpType.MATMUL:
            lhs = graph.get_tensor(op.input_tensor_ids[0])
            reduction_dim = max(reduction_dim, lhs.width)

    num_k_steps = math.ceil(reduction_dim / config.k) if reduction_dim > 0 else 1
    is_split_k = num_k_steps > 1

    # Input slices
    for tid, role in tensor_class.boundary_inputs.items():
        t = graph.get_tensor(tid)
        if role == "lhs":
            if is_split_k:
                # Split-K: LHS stored fully
                ws += t.size
            else:
                ws += config.k * config.h  # LHS slice
        elif role == "rhs":
            ws += config.w * config.k  # RHS slice (streamed)
        else:  # pointwise
            ws += config.w * config.h

    # Resident input slices (already in fast memory, still consume capacity)
    for tid in tensor_class.resident_inputs:
        # Resident tensors occupy their full size, not sliced
        t = graph.get_tensor(tid)
        ws += t.size

    # Output slices (accumulator for MatMul, output for Pointwise)
    # For the output, we hold one spatial tile at a time
    has_matmul = any(
        graph.get_op(oid).op_type == OpType.MATMUL for oid in op_ids
    )

    output_slice_size = config.w * config.h
    # Count unique boundary output tensors (usually 1)
    for tid in tensor_class.boundary_outputs:
        ws += output_slice_size

    # Retained tensors from previous steps that are still in fast memory
    # (but not already counted as resident inputs)
    for tid in state.tensors_in_fast_memory:
        if tid not in tensor_class.resident_inputs:
            # This tensor is just sitting in fast memory, occupying space
            t = graph.get_tensor(tid)
            ws += t.size

    return ws


def compute_subgraph_latency(
    graph: Graph,
    action: Action,
    state: ScheduleState,
) -> SubgraphLatencyResult:
    """
    Compute the total latency for executing a subgraph.
    This is the core physics engine of the environment.
    """
    config = action.config
    op_ids = action.operation_ids
    bw = graph.hardware.slow_memory_bandwidth
    native_w, native_h = graph.hardware.native_granularity

    # Step 1: Classify tensors
    tensor_class = classify_tensors(graph, op_ids, state)

    # Step 2: Compute tile geometry
    geom = compute_tile_geometry(graph, op_ids, config)

    # Step 3: Check working set
    ws = compute_working_set(graph, tensor_class, config, op_ids, state)
    if ws > graph.hardware.fast_memory_capacity:
        return SubgraphLatencyResult(
            total_latency=0.0,
            tile_latencies=[],
            working_set=ws,
            is_valid=False,
            error=f"OOM: working set {ws:.0f} > capacity {graph.hardware.fast_memory_capacity}",
        )

    # Step 4: Compute per-op costs
    pointwise_cost = sum(
        graph.get_op(oid).base_cost
        for oid in op_ids
        if graph.get_op(oid).op_type == OpType.POINTWISE
    )
    matmul_cost_per_k_step = sum(
        graph.get_op(oid).base_cost / geom.num_k_steps
        for oid in op_ids
        if graph.get_op(oid).op_type == OpType.MATMUL
    )

    # Step 5: Determine traversal order
    traversal = action.traversal_order
    if traversal is None:
        traversal = list(range(geom.total_spatial_tiles))

    # Step 6: Iterate tiles and compute latencies
    tile_latencies = []

    for tile_idx_in_order, tile_flat in enumerate(traversal):
        # Convert flat index to (tile_w_idx, tile_h_idx)
        tile_w_idx = tile_flat % geom.num_tiles_w
        tile_h_idx = tile_flat // geom.num_tiles_w

        # Previous tile for reuse detection
        if tile_idx_in_order > 0:
            prev_flat = traversal[tile_idx_in_order - 1]
            prev_tile_w = prev_flat % geom.num_tiles_w
            prev_tile_h = prev_flat // geom.num_tiles_w
        else:
            prev_tile_w = -1
            prev_tile_h = -1

        for k_step in range(geom.num_k_steps):
            # --- Memory In ---
            mem_in = 0.0
            is_first_k_step = (k_step == 0)
            is_first_tile = (tile_idx_in_order == 0)
            is_first_iteration = is_first_tile and is_first_k_step

            for tid, role in tensor_class.boundary_inputs.items():
                t = graph.get_tensor(tid)

                if role == "pointwise":
                    # Pointwise: load slice every new spatial tile, not on k-steps
                    if is_first_k_step:
                        slice_size = config.w * config.h
                        mem_in += slice_size / bw

                elif role == "lhs":
                    # LHS behavior depends on whether we have split-K
                    if geom.num_k_steps > 1:
                        # Split-K: load FULL tensor on first iteration, reuse across k-steps
                        if is_first_iteration:
                            mem_in += t.size / bw
                        # On subsequent k-steps AND tiles: already resident, reused
                    else:
                        # No split-K: load LHS strip per spatial tile
                        # LHS strip: [k x h], reused when tile_h is same
                        lhs_slice_size = config.k * config.h
                        if is_first_iteration:
                            mem_in += lhs_slice_size / bw
                        elif tile_h_idx == prev_tile_h:
                            pass  # Reused: same row
                        else:
                            mem_in += lhs_slice_size / bw

                elif role == "rhs":
                    # RHS is always streamed along k-dimension
                    rhs_slice_size = config.w * config.k
                    if geom.num_k_steps > 1:
                        # Split-K: load new RHS strip every k-step
                        if is_first_iteration:
                            mem_in += rhs_slice_size / bw
                        elif not is_first_k_step:
                            # New k-step: new RHS strip
                            mem_in += rhs_slice_size / bw
                        # New spatial tile, first k-step: reload RHS strip
                        elif is_first_k_step and not is_first_tile:
                            mem_in += rhs_slice_size / bw
                    else:
                        # No split-K: load RHS strip per spatial tile
                        # RHS strip: [w x k], reused when tile_w is same
                        if is_first_iteration:
                            mem_in += rhs_slice_size / bw
                        elif tile_w_idx == prev_tile_w:
                            pass  # Reused: same column
                        else:
                            mem_in += rhs_slice_size / bw

            # Resident inputs: zero load cost (already in fast memory)
            # But for MatMul, resident LHS/RHS may have different behavior
            # For now, resident = full tensor in fast memory, always available
            # (This handles the T0 reuse in Example 5 k-steps 2-4)

            # --- Memory Out ---
            mem_out = 0.0

            is_last_k_step = (k_step == geom.num_k_steps - 1)
            is_last_tile = (tile_idx_in_order == len(traversal) - 1)

            if is_last_k_step:
                # Output slice can be evicted (or retained)
                for tid in tensor_class.boundary_outputs:
                    if tid in action.tensors_to_retain:
                        # Retained: no eviction cost
                        # But only skip if this is NOT a graph output that must go to slow mem
                        # Actually, retain means keep in fast memory.
                        # Graph outputs must eventually go to slow memory,
                        # but that happens when the episode ends or a later subgraph evicts.
                        # For now, if retained, no eviction cost this step.
                        pass
                    else:
                        output_slice_size = config.w * config.h
                        mem_out += output_slice_size / bw

            # --- Compute ---
            compute = pointwise_cost + matmul_cost_per_k_step

            # --- Tile Latency (Roofline) ---
            tile_lat = max(compute, mem_in + mem_out)
            tile_latencies.append(tile_lat)

    total_latency = sum(tile_latencies)

    return SubgraphLatencyResult(
        total_latency=total_latency,
        tile_latencies=tile_latencies,
        working_set=ws,
        is_valid=True,
    )


def compute_naive_latency(graph: Graph) -> float:
    """
    Compute the total latency if every op is scheduled individually
    with native granularity, no fusion, no retention.
    This is the worst-case baseline for grading.
    """
    bw = graph.hardware.slow_memory_bandwidth
    native_w, native_h = graph.hardware.native_granularity
    total = 0.0

    for op in graph.operations:
        # Each op in its own subgraph, config = native granularity
        # All inputs loaded from slow memory, all outputs evicted

        # Determine output dimensions
        out_w = max(graph.get_tensor(tid).width for tid in op.output_tensor_ids)
        out_h = max(graph.get_tensor(tid).height for tid in op.output_tensor_ids)

        num_tiles_w = math.ceil(out_w / native_w)
        num_tiles_h = math.ceil(out_h / native_h)
        total_tiles = num_tiles_w * num_tiles_h

        if op.op_type == OpType.MATMUL:
            lhs = graph.get_tensor(op.input_tensor_ids[0])
            rhs = graph.get_tensor(op.input_tensor_ids[1])
            K = lhs.width
            # Use full K (native), so 1 k-step
            num_k = 1
            k_val = K

            for tile_i in range(total_tiles):
                tile_w_idx = tile_i % num_tiles_w
                tile_h_idx = tile_i // num_tiles_w

                # LHS slice: k x native_h
                lhs_size = K * native_h
                # RHS slice: native_w x k
                rhs_size = native_w * K
                # Output slice
                out_size = native_w * native_h

                # Reuse logic for naive (raster order)
                mem_in = 0.0
                if tile_i == 0:
                    mem_in = (lhs_size + rhs_size) / bw
                elif tile_w_idx == 0:
                    # New row: reload both
                    mem_in = (lhs_size + rhs_size) / bw
                else:
                    # Same row: LHS reused, reload RHS
                    mem_in = rhs_size / bw

                mem_out = out_size / bw
                compute = op.base_cost
                total += max(compute, mem_in + mem_out)
        else:
            # Pointwise
            for tile_i in range(total_tiles):
                mem_in = sum(
                    native_w * native_h / bw
                    for tid in op.input_tensor_ids
                )
                mem_out = sum(
                    native_w * native_h / bw
                    for tid in op.output_tensor_ids
                )
                compute = op.base_cost
                total += max(compute, mem_in + mem_out)

    return total