import numpy as np
NODE_TYPES = {"A": dict(cap=6.0, mem=4.0, p_idle=4.0, p_peak=15.0),
              "B": dict(cap=24.0, mem=16.0, p_idle=12.0, p_peak=65.0)}
DT = 0.1

def make_nodes(n_nodes, rng):
    nodes = []
    for i in range(n_nodes):
        t = "A" if rng.random() < 0.5 else "B"
        base = NODE_TYPES[t]; jit = rng.uniform(0.85, 1.15)
        nodes.append(dict(id=i, type=t, cap=base["cap"]*jit, mem=base["mem"],
                          p_idle=base["p_idle"], p_peak=base["p_peak"]))
    return nodes

def node_arrays(nodes):
    return (np.array([n["cap"] for n in nodes], dtype=float),
            np.array([n["mem"] for n in nodes], dtype=float),
            np.array([n["p_idle"] for n in nodes], dtype=float),
            np.array([n["p_peak"] for n in nodes], dtype=float))

def make_tasks(n_tasks, rng):
    if n_tasks == 0:
        return (np.zeros(0), np.zeros(0), np.zeros(0), np.zeros(0, dtype=int))
    cls = (rng.random(n_tasks) < 0.40).astype(int)
    demand = np.where(cls == 1, rng.uniform(0.05, 0.30, n_tasks),
                      rng.uniform(0.40, 1.50, n_tasks))
    mem = np.where(cls == 1, rng.uniform(0.05, 0.4, n_tasks),
                   rng.uniform(0.2, 1.2, n_tasks))
    deadline = np.where(cls == 1, 0.050, 0.500)
    return demand, mem, deadline, cls

def make_link_latency(n_nodes, rng):
    return rng.uniform(0.002, 0.020, n_nodes)

def evaluate(assign, demand, mem, deadline, cls, backlog, cap, link_lat,
             p_idle, p_peak, alpha=0.5, beta=0.3, gamma=0.2, dt=DT,
             return_metrics=False, node_mem=None):
    n_nodes = cap.shape[0]; m = assign.shape[0]
    work = np.zeros(n_nodes); memload = np.zeros(n_nodes); lat = np.zeros(m)
    order = np.lexsort((demand, -cls)); ahead = backlog.copy()
    for idx in order:
        n = assign[idx]
        proc = demand[idx] / cap[n]; wait = ahead[n] / cap[n]
        lat[idx] = link_lat[n] + wait + proc
        ahead[n] += demand[idx]; work[n] += demand[idx]; memload[n] += mem[idx]
    served = np.minimum(ahead, cap * dt)
    util = np.clip(served / (cap * dt), 0.0, 1.0)
    avg_lat = float(lat.mean()) if m > 0 else 0.0
    energy = float(np.sum((p_idle + (p_peak - p_idle) * util) * dt))
    sigma_u = float(np.std(util))
    cap_over = np.maximum(work - cap * dt, 0.0).sum()
    mem_over = np.maximum(memload - node_mem, 0.0).sum() if node_mem is not None else 0.0
    penalty = 50.0 * cap_over + 50.0 * mem_over
    lat_n = avg_lat / 0.1; en_n = energy / (n_nodes * 0.065 + 1e-9)
    fit = alpha * lat_n + beta * en_n + gamma * sigma_u + penalty
    if not return_metrics:
        return fit
    classA = cls == 1
    a_met = float(np.mean(lat[classA] <= deadline[classA])) if classA.any() else 1.0
    met_all = float(np.mean(lat <= deadline)) if m > 0 else 1.0
    jain = float((util.sum() ** 2) / (n_nodes * np.sum(util ** 2) + 1e-12)) if util.sum() > 0 else 1.0
    return dict(fitness=fit, avg_latency=avg_lat, energy=energy, sigma_u=sigma_u,
                utilization=float(util.mean()), jain=jain,
                classA_deadline_met=a_met, all_deadline_met=met_all, served=served, work=work)