diff --git a/codegen.py b/codegen.py
index 15c8437..4be4d96 100644
--- a/codegen.py
+++ b/codegen.py
@@ -129,9 +129,10 @@ def generate(officer=None):
 		ind=lookup_piece(n[0],n[1])
 		rank,cap=min(rank,pro[n[2]][ind][0]),min(cap,pro[n[2]][ind][1])
 		out.append(ind)
-	cap=(cap-(cap<10)+(cap>10))//3 # Derp
 	if gun: out.insert(1,lookup_piece(gun,tables.guns,True))
 	else: out.insert(1,tables.gunmaps_wrapped[out[0]])
+	rank,cap=min(rank,pro['guns'][out[1]][0]),min(cap,pro['guns'][out[1]][1]) # Derp2, forgot to actually cap on gun
+	cap=(cap-(cap<10)+(cap>10))//3 # Derp
 	out[-1]=tables.clas[out[0]].index(out[-1])
 	muxed=shared.muxhero(out[0],out[3],out[4])
 	meds=[]