import json
import textwrap
import math
from pyvis.network import Network

def render_graph(nodes, edges, circuit_svg="", memory_limit_mb=500.0, output_html="tensor_network.html"):
    net = Network(height="100%", width="100%", directed=True, bgcolor="#1a1a1a", font_color="white")
    
    mem_map, lineage_map = {}, {}
    num_initial = len([n for n in nodes if n["group"] == "initial"])
    
    for n in nodes:
        node_id, group, lineage = n["id"], n["group"], n["lineage"]
        lineage_map[node_id] = lineage
        
        if group == "initial":
            color, shape, size = "#808080", "dot", 20
            wrapped_indices = textwrap.fill(n['indices'], width=60)
            wrapped_shape = textwrap.fill(n['shape_str'], width=60)
            tooltip = f"Indices:\n{wrapped_indices}\n\nShape:\n{wrapped_shape}"
        else:
            mem_mb = n["memory_mb"]
            mem_map[node_id] = mem_mb
            color = "#ff4d4d" if mem_mb > memory_limit_mb else "#4dff4d"
            shape, size = "box", 30
            
            leaf_lineage = [loc for loc in lineage if loc < num_initial]
            lineage_str = ", ".join([f"T{loc}" for loc in leaf_lineage])
            
            if mem_mb < 1.0:
                mem_str = f"{mem_mb * 1024:.2f} KB"
            else:
                mem_str = f"{mem_mb:.2f} MB"
                
            wrapped_lineage = textwrap.fill(lineage_str, width=60)
            wrapped_eq = textwrap.fill(n['equation_step'], width=60)
            wrapped_shape = textwrap.fill(n['shape_str'], width=60)
            
            tooltip = f"Original Tensors:\n{wrapped_lineage}\n\nEquation:\n{wrapped_eq}\n\nShape:\n{wrapped_shape}\n\nMemory: {mem_str}"
            
        net.add_node(node_id, label=n["label"], color=color, title=tooltip, shape=shape, size=size)
        
    for e in edges: net.add_edge(e["from"], e["to"])
        
    net.set_options('{"layout": {"hierarchical": {"enabled": true, "direction": "DU", "sortMethod": "directed"}}, "physics": {"hierarchicalRepulsion": {"centralGravity": 0.0, "springLength": 100, "springConstant": 0.01, "nodeDistance": 150, "damping": 0.09}, "solver": "hierarchicalRepulsion"}}')
    net.save_graph(output_html)

    with open(output_html, "r", encoding="utf-8") as f:
        html_content = f.read()

    mem_map_json, lineage_map_json = json.dumps(mem_map), json.dumps(lineage_map)
    max_mem_mb = max(mem_map.values()) if mem_map else 0.001
    
    true_max_mb = max(131072.0, max_mem_mb * 1.05)
    
    min_log = math.log10(0.01)
    max_log = math.log10(true_max_mb)
    
    safe_mem_limit = max(0.01, memory_limit_mb)
    initial_log = math.log10(safe_mem_limit)
    initial_slider_val = ((initial_log - min_log) / (max_log - min_log)) * 100

    ui_injection = f"""
    <style>
        body, html {{ margin: 0; padding: 0; height: 100%; overflow: hidden; background: #1a1a1a; }}
        #split-container {{ display: flex; width: 100%; height: 100%; }}
        #mynetwork {{ width: 50% !important; height: 100% !important; border: none !important; outline: none !important; border-right: 2px solid #333 !important; }}
        .vis-network {{ border: none !important; outline: none !important; }}
        
        #circuit-pane {{ width: 50%; height: 100%; position: relative; overflow: hidden; background: #1a1a1a; cursor: grab; }}
        #circuit-pane:active {{ cursor: grabbing; }}
        #circuit-container {{ transform-origin: 0 0; transition: transform 0.05s linear; }}
        
        .ram-btn {{ background: #333; color: white; border: 1px solid #555; border-radius: 4px; padding: 4px 8px; cursor: pointer; font-size: 12px; flex: 1; transition: background 0.2s; }}
        .ram-btn:hover {{ background: #555; }}
        
        #overlay-ui {{ position: absolute; top: 20px; left: 20px; z-index: 1000; background: rgba(0,0,0,0.8); padding: 15px; border-radius: 8px; color: white; font-family: sans-serif; border: 1px solid #444; width: 300px; }}

        #loadingBar {{
            position: absolute !important;
            top: 0 !important; left: 0 !important;
            width: 50% !important; height: 100% !important;
            background-color: #1a1a1a !important;
            z-index: 99999 !important;
            opacity: 1;
            transition: opacity 0.5s ease;
        }}
        .outerBorder {{
            position: absolute !important;
            top: 50% !important; left: 50% !important;
            transform: translate(-50%, -50%) !important;
            width: 400px !important; height: 44px !important;
            border-radius: 22px !important;
            background: #2a2a2a !important;
            border: 2px solid #444 !important;
            margin: 0 !important;
            box-shadow: 0px 4px 15px rgba(0,0,0,0.5) !important;
        }}
        #text {{
            position: absolute !important;
            top: 0 !important; left: 0 !important;
            width: 100% !important; height: 100% !important;
            line-height: 44px !important;
            text-align: center !important;
            color: #ffffff !important;
            font-size: 16px !important;
            font-family: sans-serif !important;
            font-weight: bold !important;
            margin: 0 !important;
            z-index: 10 !important;
            text-shadow: 1px 1px 2px rgba(0,0,0,0.8) !important;
        }}
        #border {{
            position: absolute !important;
            top: 4px !important; left: 4px !important;
            width: 388px !important; height: 32px !important;
            border-radius: 16px !important;
            background: #1a1a1a !important;
            box-shadow: none !important;
            border: none !important;
            overflow: hidden !important;
        }}
        #bar {{
            position: absolute !important;
            top: 0 !important; left: 0 !important;
            height: 100% !important;
            background: #4dff4d !important;
            border-radius: 16px !important;
            box-shadow: none !important;
            border: none !important;
            transition: width 0.1s ease !important;
        }}
    </style>
    
    <div id="split-container">
        <div id="circuit-pane">
            <div id="circuit-container">
                {circuit_svg}
            </div>
        </div>
    </div>
    
    <div id="overlay-ui">
        <label for="memSlider" style="display: block; margin-bottom: 8px;"><b>RAM Threshold:</b> <span id="memVal">{memory_limit_mb:.2f} MB</span></label>
        <input type="range" id="memSlider" min="0" max="100" value="{initial_slider_val:.4f}" step="0.1" oninput="onSliderInput(this.value)" style="width: 100%; margin-bottom: 10px;">
        <div style="display: flex; gap: 5px; justify-content: space-between;">
            <button class="ram-btn" onclick="setSlider(4096)">4GB</button>
            <button class="ram-btn" onclick="setSlider(16384)">16GB</button>
            <button class="ram-btn" onclick="setSlider(32768)">32GB</button>
            <button class="ram-btn" onclick="setSlider(65536)">64GB</button>
            <button class="ram-btn" onclick="setSlider(131072)">128GB</button>
        </div>
    </div>
    
    <script>
    document.addEventListener("DOMContentLoaded", function() {{
        var netDiv = document.getElementById('mynetwork');
        var splitCont = document.getElementById('split-container');
        splitCont.insertBefore(netDiv, splitCont.firstChild);
    }});

    var nodeMemoryMap = {mem_map_json};
    var nodeLineageMap = {lineage_map_json};
    
    var minLog = {min_log};
    var maxLog = {max_log};

    function setSlider(valMB) {{
        var slider = document.getElementById('memSlider');
        var safeVal = Math.max(0.01, valMB);
        
        if (safeVal > Math.pow(10, maxLog)) {{
            maxLog = Math.log10(safeVal * 1.05);
        }}
        
        var logVal = Math.log10(safeVal);
        slider.value = ((logVal - minLog) / (maxLog - minLog)) * 100;
        updateThreshold(safeVal);
    }}

    function onSliderInput(sliderPos) {{
        var logVal = minLog + (sliderPos / 100) * (maxLog - minLog);
        var actualMB = Math.pow(10, logVal);
        updateThreshold(actualMB);
    }}

    function updateThreshold(numVal) {{
        var displayStr = numVal < 1.0 ? (numVal * 1024).toFixed(2) + " KB" : numVal.toFixed(2) + " MB";
        document.getElementById('memVal').innerText = displayStr;
        
        if (typeof nodes === 'undefined') return; 
        var updates = [];
        nodes.get().forEach(function(node) {{
            if (nodeMemoryMap.hasOwnProperty(node.id)) {{
                var mem = nodeMemoryMap[node.id];
                updates.push({{ id: node.id, color: mem > numVal ? {{ background: "#ff4d4d", border: "#cc0000" }} : {{ background: "#4dff4d", border: "#00cc00" }}, borderWidth: 1 }});
            }} else {{
                updates.push({{ id: node.id, color: {{ background: "#808080", border: "#666666" }}, borderWidth: 1 }});
            }}
        }});
        nodes.update(updates); 
    }}

    function highlightOrigins(selectedNodeId) {{
        var origins = nodeLineageMap[selectedNodeId] || [];
        var updates = [];
        nodes.get().forEach(function(node) {{
            if (node.id === selectedNodeId) {{
                updates.push({{ id: node.id, borderWidth: 3, color: {{ border: "white" }} }});
            }} else if (origins.includes(node.id)) {{
                var isInitial = !nodeMemoryMap.hasOwnProperty(node.id);
                updates.push({{ id: node.id, color: isInitial ? {{ background: "#FFD700", border: "#FFA500" }} : {{ background: "#ffae42", border: "#ff9900" }}, borderWidth: isInitial ? 3 : 2 }});
            }} else {{
                updates.push({{ id: node.id, color: {{ background: "#2a2a2a", border: "#1a1a1a" }}, borderWidth: 1 }});
            }}
        }});
        nodes.update(updates);

        document.querySelectorAll('.svg-gate rect, .svg-gate line, .svg-gate circle').forEach(el => {{
            el.setAttribute('stroke', '#666'); 
            if (el.tagName === 'circle') el.setAttribute('fill', '#666');
        }});
        
        origins.forEach(function(origId) {{
            var gateGroup = document.getElementById('gate-T' + origId);
            if (gateGroup) {{
                gateGroup.querySelectorAll('rect, line').forEach(el => el.setAttribute('stroke', '#FFD700'));
                gateGroup.querySelectorAll('circle').forEach(el => el.setAttribute('fill', '#FFD700'));
            }}
        }});
    }}

    var checkExist = setInterval(function() {{
       if (typeof network !== 'undefined') {{
          clearInterval(checkExist);
          network.on("click", function (params) {{
              if (params.nodes.length > 0) {{
                  highlightOrigins(params.nodes[0]);
              }} else {{
                  // Re-evaluate using the inverse log function!
                  var sliderPos = document.getElementById('memSlider').value;
                  var logVal = minLog + (sliderPos / 100) * (maxLog - minLog);
                  updateThreshold(Math.pow(10, logVal));
                  
                  document.querySelectorAll('.svg-gate rect, .svg-gate line, .svg-gate circle').forEach(el => {{
                      el.setAttribute('stroke', '#666');
                      if (el.tagName === 'circle') el.setAttribute('fill', '#666');
                  }});
              }}
          }});
       }}
    }}, 100); 

    const circuitPane = document.getElementById('circuit-pane');
    const circuitCont = document.getElementById('circuit-container');
    
    let scale = 1;
    let pointX = 50; 
    let pointY = 50;
    let start = {{ x: 0, y: 0 }};
    let isPanned = false;

    circuitCont.style.transform = `translate(${{pointX}}px, ${{pointY}}px) scale(${{scale}})`;

    circuitPane.onmousedown = function(e) {{
        e.preventDefault();
        start = {{ x: e.clientX - pointX, y: e.clientY - pointY }};
        isPanned = true;
    }};

    circuitPane.onmouseup = function(e) {{
        isPanned = false;
    }};

    circuitPane.onmouseleave = function(e) {{
        isPanned = false;
    }};

    circuitPane.onmousemove = function(e) {{
        e.preventDefault();
        if (!isPanned) return;
        pointX = (e.clientX - start.x);
        pointY = (e.clientY - start.y);
        circuitCont.style.transform = `translate(${{pointX}}px, ${{pointY}}px) scale(${{scale}})`;
    }};

    circuitPane.onwheel = function(e) {{
        e.preventDefault();
        var rect = circuitPane.getBoundingClientRect();
        var mouseX = e.clientX - rect.left;
        var mouseY = e.clientY - rect.top;

        var xs = (mouseX - pointX) / scale;
        var ys = (mouseY - pointY) / scale;
        
        var delta = (e.wheelDelta ? e.wheelDelta : -e.deltaY);
        (delta > 0) ? (scale *= 1.1) : (scale /= 1.1);
        
        pointX = mouseX - xs * scale;
        pointY = mouseY - ys * scale;
        
        circuitCont.style.transform = `translate(${{pointX}}px, ${{pointY}}px) scale(${{scale}})`;
    }};
    </script>
    """
    html_content = html_content.replace('<body>', f'<body>\n{ui_injection}')
    with open(output_html, "w", encoding="utf-8") as f:
        f.write(html_content)