import json, sys

inst = json.load(open("instance.json"))
H = inst["hours"]
gens = inst["generators"]
G = len(gens)
reserve_pct = inst["reserve_pct"]
demand = inst["demand_mw"]
wind = inst["wind_mw"]
solar = inst["solar_mw"]

net_demand = [demand[h] - wind[h] - solar[h] for h in range(H)]
reserve_req = [reserve_pct * demand[h] for h in range(H)]

def onv(g,h): return f"on_{g}_{h}"
def pv(g,h): return f"p_{g}_{h}"
def suv(g,h): return f"su_{g}_{h}"

lines = []
lines.append("\\ grid unit commitment MILP")
lines.append("Minimize")
obj_terms = []
for g in range(G):
    gen = gens[g]
    for h in range(H):
        obj_terms.append(f"{gen['fuel_cost']:.10f} {pv(g,h)}")
        obj_terms.append(f"{gen['no_load_cost']:.10f} {onv(g,h)}")
        obj_terms.append(f"{gen['startup_cost']:.10f} {suv(g,h)}")
lines.append(" obj: " + " + ".join(obj_terms))

lines.append("Subject To")
cidx = 0

# demand equality
for h in range(H):
    terms = " + ".join(pv(g,h) for g in range(G))
    lines.append(f" dem_{h}: {terms} = {net_demand[h]:.10f}")

# reserve: sum(pmax*on) - sum(p) >= reserve_req
for h in range(H):
    terms = []
    for g in range(G):
        terms.append(f"{gens[g]['pmax']:.10f} {onv(g,h)}")
    for g in range(G):
        terms.append(f"-1 {pv(g,h)}")
    lines.append(f" res_{h}: " + " + ".join(terms).replace("+ -", "- ") + f" >= {reserve_req[h]:.10f}")

# p bounds linking on
for g in range(G):
    gen = gens[g]
    for h in range(H):
        # p <= pmax*on  => p - pmax*on <= 0
        lines.append(f" pub_{g}_{h}: {pv(g,h)} - {gen['pmax']:.10f} {onv(g,h)} <= 0")
        # p >= pmin*on => p - pmin*on >= 0
        lines.append(f" plb_{g}_{h}: {pv(g,h)} - {gen['pmin']:.10f} {onv(g,h)} >= 0")

# startup definition and min up/down, with init state
for g in range(G):
    gen = gens[g]
    init_on = gen["init_on"]
    init_hours = gen["init_hours"]
    min_up = gen["min_up"]
    min_down = gen["min_down"]

    # startup >= on[h] - on[h-1] (prev = init_on for h=0)
    for h in range(H):
        if h == 0:
            # su_0 >= on_0 - init_on  => su_0 - on_0 >= -init_on
            lines.append(f" sudef_{g}_{h}: {suv(g,h)} - {onv(g,h)} >= {-init_on}")
        else:
            lines.append(f" sudef_{g}_{h}: {suv(g,h)} - {onv(g,h)} + {onv(g,h-1)} >= 0")

    # forced initial state (mandatory continuation)
    if init_on == 1:
        R_on = min(H, max(0, min_up - init_hours))
        for h in range(R_on):
            lines.append(f" fix_{g}_{h}: {onv(g,h)} = 1")
    else:
        R_off = min(H, max(0, min_down - init_hours))
        for h in range(R_off):
            lines.append(f" fix_{g}_{h}: {onv(g,h)} = 0")

    # min-up constraints: for each h, for h2 in h..min(h+min_up-1,H-1):
    #   on[h2] >= on[h] - on[h-1]   (on[-1] = init_on)
    for h in range(H):
        prevterm = f"{onv(g,h-1)}" if h > 0 else None
        upper = min(h + min_up - 1, H - 1)
        for h2 in range(h, upper + 1):
            if h2 == h:
                continue  # trivial on[h]>=on[h]-on[h-1] always true-ish, skip self but keep others
            if h == 0:
                # on[h2] - on[0] >= -init_on
                lines.append(f" mu_{g}_{h}_{h2}: {onv(g,h2)} - {onv(g,0)} >= {-init_on}")
            else:
                lines.append(f" mu_{g}_{h}_{h2}: {onv(g,h2)} - {onv(g,h)} + {onv(g,h-1)} >= 0")

    # min-down constraints: for each h, for h2 in h..min(h+min_down-1,H-1):
    #   1 - on[h2] >= on[h-1] - on[h]   (on[-1] = init_on)
    for h in range(H):
        upper = min(h + min_down - 1, H - 1)
        for h2 in range(h, upper + 1):
            if h2 == h:
                continue
            if h == 0:
                # 1 - on[h2] >= init_on - on[0]  => -on[h2] + on[0] >= init_on - 1
                lines.append(f" md_{g}_{h}_{h2}: -{onv(g,h2)} + {onv(g,0)} >= {init_on - 1}")
            else:
                lines.append(f" md_{g}_{h}_{h2}: -{onv(g,h2)} + {onv(g,h-1)} - {onv(g,h)} >= -1")

lines.append("Bounds")
for g in range(G):
    gen = gens[g]
    for h in range(H):
        lines.append(f" 0 <= {pv(g,h)} <= {gen['pmax']:.10f}")
        lines.append(f" su_{g}_{h} >= 0")

lines.append("Binary")
bin_terms = []
for g in range(G):
    for h in range(H):
        bin_terms.append(onv(g,h))
# print multiple per line
for i in range(0, len(bin_terms), 10):
    lines.append(" " + " ".join(bin_terms[i:i+10]))

lines.append("End")

with open("model.lp", "w") as f:
    f.write("\n".join(lines))

print("wrote model.lp, lines:", len(lines))