import tsnet
from tsnet.network import TransientModel
from tsnet.simulation import Initializer, MOCSimulator
import matplotlib.pyplot as plt
import io

def simulate_valve_closure(inp_path, node_id, start_time, duration, end_open=0, closure_constant=2):
    tm = TransientModel(inp_path)
    tm.set_wavespeed(1200.0)
    tm.set_time(20)
    tm.valve_closure(node_id, [duration, start_time, end_open, closure_constant])
    tm = Initializer(tm, 0)
    tm = MOCSimulator(tm)

    fig, ax = plt.subplots()
    for node in ['N2', 'N3']:
        head = tm.get_node_head(node)
        ax.plot(tm.Time, head, label=f"Node {node}")
    ax.set_title("Head vs Time")
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Head (m)")
    ax.legend()
    buf = io.BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)
    return buf