#2 입력 feature 생성 #

%load_ext autoreload
%autoreload 2

import sys
import os

import pickle
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import pandas as pd

sys.path.append('/data3/projects/2025_Antibiotics/YSH/bin')
from sc import *

os.chdir('/data3/projects/2025_Antibiotics/YSH')

seqdir = 'data/res_dict'
seq_list = os.listdir(seqdir)
print(len(seq_list))

169

med_dir = 'res'
seqdir = 'data/res_dict'
outdir = 'res/feature1'

with open(f"{med_dir}/all_meds.txt", 'r') as f:
    all_meds = [line.strip() for line in f if line.strip()]
    all_meds = [s.replace("/", "_") for s in all_meds]

for med in all_meds:
    print(med)
    cur_path = f'{seqdir}/{med}.pkl'

    with open(cur_path, 'rb') as f:
        res_dict = pickle.load(f)
        
    feature1_list = []

    for pid, df in res_dict.items():
        news_bf = df.iloc[2]['NEWS']  # 3번째 행 (0-indexed)
        news_af = df.iloc[3:]['NEWS'].max()  # 4번째 행부터 마지막까지 중 최댓값

        if news_af < news_bf:  # "작은" 경우만 (같은 건 포함하지 않음)
            feature1_list.append(pid)
    
    print(len(feature1_list))
    filtered_res_dict = {pid: res_dict[pid] for pid in feature1_list if pid in res_dict}
    
    with open(f"{outdir}/{med}.pkl", 'wb') as f:
        pickle.dump(filtered_res_dict, f)

strain_dic = {}

for med in all_meds:
    cur_path = f'{outdir}/{med}.pkl'

    with open(cur_path, 'rb') as f:
        filtered_res_dict = pickle.load(f)

    for pid, df in filtered_res_dict.items():
        if len(df) < 3:
            continue  # 3번째 행이 없는 경우 skip

        try:
            cur_strain = df.iloc[2]['strain']
            if isinstance(cur_strain, list):
                strains = cur_strain
            else:
                strains = [cur_strain]
        except Exception as e:
            continue

        for strain in strains:
            if strain in strain_dic:
                strain_dic[strain].append(med)
            else:
                strain_dic[strain] = [med]

for strain in strain_dic:
    strain_dic[strain] = list(set(strain_dic[strain]))

strains = list(strain_dic.keys())
print(len(strains))
print(strains[0], '\n', strain_dic[strains[0]])

98
Pseudomonas aeruginosa 
 ['Imicil Kit', 'Prepenem', 'Sevatrim', 'Meropen', 'Finibax', 'Ciprobay', 'Hanomycin', 'Citopcin', 'Tazocin', 'Dexamethasone', 'Pospenem', 'Tygacil', 'Pentamidine Isethionate', 'Lagevrio', 'Plunazol', 'Vancomycin HCl']

ourdir = 'res'

with open(f"{outdir}/Feature1.pkl", 'wb') as f:
    pickle.dump(strain_dic, f)