import pandas as pd

df = pd.read_csv("./processed_test.csv")


func_lab = []
stat_lab = []
cwe_id = []
func = []

df_vul = df[df["function_label"]==1][:50]
df_vul = df_vul.reset_index()

df_non_vul = df[df["function_label"]==0][:50]
df_non_vul = df_non_vul.reset_index()

for i in range(len(df_vul)):
    func_lab.append(df_vul["function_label"][i])
    stat_lab.append(df_vul["statement_label"][i])
    
    id = df_vul["cwe_id"][i]
    if df_vul["function_label"][i] == 0:
        id = None
    cwe_id.append(id)
    func.append(df_vul["func_before"][i])

    func_lab.append(df_non_vul["function_label"][i])
    stat_lab.append(df_non_vul["statement_label"][i])
    
    id = df_non_vul["cwe_id"][i]
    if df_non_vul["function_label"][i] == 0:
        id = None
    
    cwe_id.append(id)
    func.append(df_non_vul["func_before"][i])
    
    
pd.DataFrame({"function": func, "function_label": func_lab, "cwe_id": cwe_id, "statement_label": stat_lab}).to_csv("./test.csv", index=False)