import torch
import torch._dynamo
import trme_torch

def trme_compiler_backend(gm: torch.fx.GraphModule, example_inputs):
    print("⚡ TRME Compiler: Optimizing Graph...")

    # Simple substitution pass
    new_graph = gm.graph

    for node in new_graph.nodes:
        if node.op == 'call_function':
            if node.target == torch.nn.functional.linear:
                print(f"   -> Replacing Linear Node: {node.name}")
                # We should replace with TRME op
                # node.target = trme_torch.matmul
                # Logic: get args (input, weight, bias)
                # bias is optional. TRME matmul supports input, weight, bias.

                # Swap always (assuming forward logic handles bias)
                node.target = trme_torch.matmul

    gm.recompile()
    return gm

def compile_model(model):
    return torch.compile(model, backend=trme_compiler_backend)
