import torch
import torch.nn as nn
from graphviz import Digraph
from IPython.display import Image, display
import tempfile


def visualize(model: nn.Module, sample_inputs=None, filename=None):
    def format_label(name):
        """Format label with subscripts for names containing underscores."""
        if '_' in name:
            parts = name.split('_', 1)
            return f"{parts[0]}<SUB>{parts[1]}</SUB>"
        return name

    # Create a mapping from parameter tensor id to parameter name
    param_names = {}
    for name, param in model.named_parameters():
        param_names[id(param)] = name

    # Create sample inputs if not provided
    if sample_inputs is None:
        # Try to infer from model attributes
        if hasattr(model, 'input_size') and hasattr(model, 'hidden_size'):
            # Get device from model parameters
            device = next(model.parameters()).device
            # RNN model - create single timestep inputs (input_size, seq_len, batch_size)
            x_t = torch.randn(model.input_size, 1, 1, requires_grad=True, device=device)
            h_prev = torch.randn(model.hidden_size, 1, requires_grad=True, device=device)
            sample_inputs = (x_t, h_prev)
        else:
            raise ValueError("Please provide sample_inputs for the forward pass")

    # Create mapping for input tensors
    input_names = {}
    if isinstance(sample_inputs, tuple):
        for i, inp in enumerate(sample_inputs):
            if i == 0:
                input_names[id(inp)] = 'x'
            elif i == 1:
                input_names[id(inp)] = 'h'
            else:
                input_names[id(inp)] = f'input_{i}'

    # Run forward pass to build autograd graph
    outputs = model(*sample_inputs)

    # Handle multiple outputs (e.g., RNN returns (outputs, h))
    if isinstance(outputs, tuple):
        # Trace from both outputs and final hidden to capture everything
        # The final hidden state captures the full recurrent chain
        output_scalar = outputs[0].sum()
        hidden_scalar = outputs[1].sum()
        output = output_scalar + hidden_scalar
    else:
        output = outputs.sum()

    # Traverse autograd graph
    seen_nodes = set()
    seen_edges = set()
    nodes = []
    edges = []

    def traverse(grad_fn, parent_id=None):
        if grad_fn is None:
            return

        node_id = id(grad_fn)
        op_name = type(grad_fn).__name__.replace('Backward', '')

        # Add node if not seen
        if node_id not in seen_nodes:
            nodes.append((node_id, op_name, grad_fn))
            seen_nodes.add(node_id)

        # Add edge if not seen
        if parent_id is not None:
            edge = (node_id, parent_id)
            if edge not in seen_edges:
                edges.append(edge)
                seen_edges.add(edge)

        # For AccumulateGrad, don't traverse further (leaf nodes)
        if op_name == 'AccumulateGrad':
            return

        # Traverse inputs
        if hasattr(grad_fn, 'next_functions'):
            for fn, _ in grad_fn.next_functions:
                if fn is not None:
                    traverse(fn, node_id)

    traverse(output.grad_fn)

    # Build Graphviz graph
    dot = Digraph(format="png")
    dot.attr(rankdir="LR", splines="spline", bgcolor="white")
    dot.attr('node', fontname='Helvetica')
    dot.attr('edge', fontname='Helvetica')
    dot.attr('graph', fontname='Helvetica')

    with dot.subgraph(name="cluster_cell") as c:
        c.attr(
            label="",
            style="rounded,filled",
            fillcolor="#fdfdfd",
            color="#bbbbbb",
            penwidth="2",
        )

        for node_id, op_name, grad_fn in nodes:
            # Determine label and style
            if op_name == 'AccumulateGrad':
                # Try to get parameter or input name
                if hasattr(grad_fn, 'variable'):
                    var_id = id(grad_fn.variable)
                    if var_id in param_names:
                        label = format_label(param_names[var_id])
                        fill = "#9ecae1"  # blue for parameters
                    elif var_id in input_names:
                        label = format_label(input_names[var_id])
                        fill = "#fdd0a2"  # orange for inputs
                    else:
                        label = 'param'
                        fill = "#9ecae1"
                else:
                    label = 'param'
                    fill = "#9ecae1"
                shape = "box"
            elif 'Tanh' in op_name:
                label = "tanh"
                fill = "#c7e9c0"
                shape = "circle"
            elif 'Add' in op_name or 'Sum0' in op_name:
                label = "+"
                fill = "#e0e0e0"
                shape = "circle"
            elif 'Mul' in op_name or 'Mm' in op_name:
                label = "·"
                fill = "#e0e0e0"
                shape = "circle"
            elif 'Stack0' in op_name:
                label = "Stack"
                fill = "#e0e0e0"
                shape = "circle"
            elif 'Select0' in op_name:
                label = "[ ]"
                fill = "#e0e0e0"
                shape = "circle"
            else:
                label = op_name
                fill = "#e0e0e0"
                shape = "circle"

            fontsize = "40" if label in {"·"} else "18"

            # Use HTML label if it contains SUB tags
            if '<SUB>' in label:
                c.node(
                    str(node_id),
                    label=f"<{label}>",
                    shape=shape,
                    style="filled",
                    fillcolor=fill,
                    fontsize=fontsize,
                    width=".8",
                    height=".8",
                    fixedsize="true",
                    penwidth="1.5",
                )
            else:
                c.node(
                    str(node_id),
                    label=label,
                    shape=shape,
                    style="filled",
                    fillcolor=fill,
                    fontsize=fontsize,
                    width=".8",
                    height=".8",
                    fixedsize="true",
                    penwidth="1.5",
                )

        for src, dst in edges:
            c.edge(str(src), str(dst), arrowsize="0.8")

    # Render to temp directory
    if filename is None:
        filename = "pytorch"
        tmpdir = tempfile.TemporaryDirectory()
        path = dot.render(f"{tmpdir}/{filename}", cleanup=True)
    else:
        path = dot.render(filename, cleanup=True)
    with open(path, 'rb') as f:
        img_data = f.read()
    display(Image(img_data))