import pandas as pd
import numpy as np
import os

import matplotlib.pyplot as plt
import random

os.sys.path.append("/data/home/ysh980101/2410/Mutclust2") 
from Bin.sc import *

/Influenza
└── Raw/
     ├── A-H1N1.fasta
     ├── A-H1N1.fasta
     ├── ...
     └── B.fasta

gene_list = ['NP', 'HA', 'PB2', 'PA', 'MP', 'NS', 'PB1', 'NA']
subtype_list = ['A-H1N1', 'A-H3N2', 'A-H5N1', 'B']

for subtype in subtype_list:
    for gene in gene_list:
        # 입력 파일 경로와 출력 파일 경로
        input_file = f'/data3/projects/2020_MUTCLUST/Data/Projects/GISAID_revision/Influenza/Raw/{subtype}.fasta'
        output_dir = f'/data3/projects/2020_MUTCLUST/Data/Projects/GISAID_revision/Influenza/Preprocessed/{gene}'
        output_file = os.path.join(output_dir, f'{subtype}.fasta')

        # 출력 디렉토리가 없으면 생성
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # 파일을 열고 편집된 결과를 출력 파일에 쓰기
        with open(input_file, 'r') as infile, open(output_file, 'w') as outfile:
            cur_meta = ""
            cur_seq = []
            write_seq = False

            for line in infile:
                if line.startswith('>'):  # 메타 정보가 시작되는 행
                    # 이전 시퀀스를 처리
                    if cur_meta and write_seq:
                        # 메타와 시퀀스를 출력 파일에 기록
                        outfile.write(cur_meta)
                        outfile.write(''.join(cur_seq))

                    # 새로운 메타 정보 처리
                    cur_meta = line
                    cur_seq = []
                    split_meta = cur_meta.split('|')
                    if len(split_meta) > 6 and split_meta[6] == gene:  # 7번째 값이 gene인지 확인
                        write_seq = True  # 시퀀스를 출력할지 여부 결정
                    else:
                        write_seq = False

                else:
                    if write_seq:
                        cur_seq.append(line)  # 시퀀스 데이터를 리스트에 추가

            # 마지막 시퀀스를 처리
            if cur_meta and write_seq:
                outfile.write(cur_meta)
                outfile.write(''.join(cur_seq))

/Influenza
└── Raw/
     ├── A-H1N1.fasta
     ├── A-H1N1.fasta
     ├──  ...
     └── B.fasta
└── Preprocessed/
     ├── HA/
     │    ├── A-H1N1.fasta
     │    ├── A-H1N1.fasta
     │    ├── ...
     │    └── B.fasta
     └── ...
          └── (HA와 동일 구조)