app.py

import os,sys

# install required packages
os.system('pip install plotly')  # plotly 설치
os.system('pip install matplotlib')  # matplotlib 설치
os.system('pip install dgl==1.0.2+cu116 -f https://data.dgl.ai/wheels/cu116/repo.html')
os.environ["DGLBACKEND"] = "pytorch"
print('Modules installed')

# 기본 args 설정
if not os.path.exists('./tmp'):
    os.makedirs('./tmp')

if not os.path.exists('./tmp/args.json'):
    default_args = {
        'checkpoint': None,
        'dump_trb': False,
        'dump_args': True,
        'save_best_plddt': True,
        'T': 25,
        'strand_bias': 0.0,
        'loop_bias': 0.0,
        'helix_bias': 0.0,
        'd_t1d': 24,
        'potentials': None,
        'potential_scale': None,
        'aa_composition': None
    }
    with open('./tmp/args.json', 'w') as f:
        json.dump(default_args, f)

# 체크포인트 파일 다운로드
if not os.path.exists('./SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'):
    print('Downloading model weights 1')
    os.system('wget http://files.ipd.uw.edu/pub/sequence_diffusion/checkpoints/SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt')
    print('Successfully Downloaded')

if not os.path.exists('./SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'):
    print('Downloading model weights 2')
    os.system('wget http://files.ipd.uw.edu/pub/sequence_diffusion/checkpoints/SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt')
    print('Successfully Downloaded')

from openai import OpenAI
import gradio as gr
import json  # json 모듈 추가
from datasets import load_dataset
import plotly.graph_objects as go
import numpy as np
import py3Dmol
from io import StringIO
import json
import secrets
import copy
import matplotlib.pyplot as plt
from utils.sampler import HuggingFace_sampler
from utils.parsers_inference import parse_pdb
from model.util import writepdb
from utils.inpainting_util import *
import os

# args 로드
with open('./tmp/args.json', 'r') as f:
    args = json.load(f)

plt.rcParams.update({'font.size': 13})

# manually set checkpoint to load
args['checkpoint'] = None
args['dump_trb'] = False
args['dump_args'] = True
args['save_best_plddt'] = True
args['T'] = 25
args['strand_bias'] = 0.0
args['loop_bias'] = 0.0
args['helix_bias'] = 0.0

# Hugging Face 토큰 설정
ACCESS_TOKEN = os.getenv("HF_TOKEN")
if not ACCESS_TOKEN:
    raise ValueError("HF_TOKEN not found in environment variables")

# OpenAI 클라이언트 설정 (Hugging Face 엔드포인트 사용)
client = OpenAI(
    base_url="https://api-inference.huggingface.co/v1/",
    api_key=ACCESS_TOKEN,
)


# 데이터셋 로드 및 구조 확인
try:
    ds = load_dataset("lamm-mit/protein_secondary_structure_from_PDB",
                     token=ACCESS_TOKEN)
    print("Dataset structure:", ds)
    print("First entry example:", next(iter(ds['train'])))
except Exception as e:
    print(f"Dataset loading error: {str(e)}")
    raise
    
def respond(
    message,
    history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for msg in history:
        messages.append({"role": "user", "content": msg[0]})
        if msg[1]:
            messages.append({"role": "assistant", "content": msg[1]})

    messages.append({"role": "user", "content": message})

    try:
        response = ""
        for chunk in client.chat.completions.create(
            model="CohereForAI/c4ai-command-r-plus-08-2024",
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
            messages=messages,
        ):
            if hasattr(chunk.choices[0].delta, 'content'):
                token = chunk.choices[0].delta.content
                if token is not None:
                    response += token
                    yield [{"role": "user", "content": message},
                          {"role": "assistant", "content": response}]
        
        return [{"role": "user", "content": message},
                {"role": "assistant", "content": response}]
    except Exception as e:
        print(f"Error in respond: {str(e)}")
        return [{"role": "user", "content": message},
                {"role": "assistant", "content": f"오류가 발생했습니다: {str(e)}"}]

def analyze_prompt(message):
    """LLM을 사용하여 프롬프트 분석"""
    try:
        analysis_prompt = f"""
        다음 요청을 분석하여 단백질 설계에 필요한 주요 특성을 추출하세요:
        요청: {message}
        
        다음 항목들을 분석해주세요:
        1. 주요 기능 (예: 치료, 결합, 촉매 등)
        2. 목표 환경 (예: 세포막, 수용성, 등)
        3. 필요한 구조적 특징
        4. 크기 및 복잡도 요구사항
        """
        
        response = client.chat.completions.create(
            model="CohereForAI/c4ai-command-r-plus-08-2024",
            messages=[{"role": "user", "content": analysis_prompt}],
            temperature=0.7
        )
        
        return response.choices[0].message.content
    except Exception as e:
        print(f"프롬프트 분석 중 오류: {str(e)}")
        return None

def search_protein_data(analysis, dataset):
    """분석 결과를 바탕으로 데이터셋에서 유사한 구조 검색"""
    try:
        # 키워드 추출
        keywords = extract_keywords(analysis)
        print("Extracted keywords:", keywords)
        
        # 데이터셋 구조 확인
        if not dataset or 'train' not in dataset:
            print("Invalid dataset structure")
            return []
            
        # 유사도 점수 계산
        scored_entries = []
        for entry in dataset['train']:
            try:
                score = calculate_similarity(keywords, entry)
                scored_entries.append((score, entry))
            except Exception as e:
                print(f"Error processing entry: {str(e)}")
                continue
        
        # 결과 정렬 및 반환
        scored_entries.sort(reverse=True)
        return scored_entries[:3]
        
    except Exception as e:
        print(f"데이터 검색 중 오류: {str(e)}")
        return []
        
def extract_parameters(analysis, similar_structures):
    """분석 결과와 유사 구조를 바탕으로 생성 파라미터 결정"""
    try:
        # 기본 파라미터 템플릿
        params = {
            'sequence_length': 100,
            'helix_bias': 0.02,
            'strand_bias': 0.02,
            'loop_bias': 0.1,
            'hydrophobic_target_score': 0
        }
        
        # 분석 결과에서 구조적 요구사항 파악
        if "막 투과" in analysis or "소수성" in analysis:
            params['hydrophobic_target_score'] = -2
            params['helix_bias'] = 0.03
        elif "수용성" in analysis or "가용성" in analysis:
            params['hydrophobic_target_score'] = 2
            params['loop_bias'] = 0.15
            
        # 유사 구조들의 특성 반영
        if similar_structures:
            avg_length = sum(len(s[1]['sequence']) for s in similar_structures) / len(similar_structures)
            params['sequence_length'] = int(avg_length)
            
            # 구조적 특성 분석 및 반영
            for _, structure in similar_structures:
                if 'secondary_structure' in structure:
                    helix_ratio = structure['secondary_structure'].count('H') / len(structure['secondary_structure'])
                    sheet_ratio = structure['secondary_structure'].count('E') / len(structure['secondary_structure'])
                    params['helix_bias'] = max(0.01, min(0.05, helix_ratio))
                    params['strand_bias'] = max(0.01, min(0.05, sheet_ratio))
        
        return params
    except Exception as e:
        print(f"파라미터 추출 중 오류: {str(e)}")
        return None

def process_chat(message, history):
    try:
        if any(keyword in message.lower() for keyword in ['protein', 'generate', '단백질', '생성', '치료']):
            # 1. LLM을 사용한 프롬프트 분석
            analysis = analyze_prompt(message)
            if not analysis:
                return history + [
                    {"role": "user", "content": message},
                    {"role": "assistant", "content": "요청 분석에 실패했습니다."}
                ]
            
            # 2. 유사 구조 검색
            similar_structures = search_protein_data(analysis, ds)
            if not similar_structures:
                return history + [
                    {"role": "user", "content": message},
                    {"role": "assistant", "content": "적합한 참조 구조를 찾지 못했습니다."}
                ]
            
            # 3. 생성 파라미터 결정
            params = extract_parameters(analysis, similar_structures)
            if not params:
                return history + [
                    {"role": "user", "content": message},
                    {"role": "assistant", "content": "파라미터 설정에 실패했습니다."}
                ]
            
            # 4. 단백질 생성
            try:
                protein_result = protein_diffusion_model(
                    sequence=None,
                    seq_len=params['sequence_length'],
                    helix_bias=params['helix_bias'],
                    strand_bias=params['strand_bias'],
                    loop_bias=params['loop_bias'],
                    secondary_structure=None,
                    aa_bias=None,
                    aa_bias_potential=None,
                    num_steps="25",
                    noise="normal",
                    hydrophobic_target_score=str(params['hydrophobic_target_score']),
                    hydrophobic_potential="2",
                    contigs=None,
                    pssm=None,
                    seq_mask=None,
                    str_mask=None,
                    rewrite_pdb=None
                )
                
                output_seq, output_pdb, structure_view, plddt_plot = next(protein_result)
                
                # 5. 결과 설명 생성
                explanation = f"""
                요청하신 기능에 맞는 단백질을 생성했습니다:
                
                분석된 요구사항:
                {analysis}
                
                설계된 구조적 특징:
                - 길이: {params['sequence_length']} 아미노산
                - 알파 헬릭스 비율: {params['helix_bias']*100:.1f}%
                - 베타 시트 비율: {params['strand_bias']*100:.1f}%
                - 루프 구조 비율: {params['loop_bias']*100:.1f}%
                - 소수성 점수: {params['hydrophobic_target_score']}
                
                참조된 유사 구조: {len(similar_structures)}개
                
                생성된 단백질의 3D 구조와 시퀀스를 확인하실 수 있습니다.
                """
                
                # 6. 결과 저장
                global current_protein_result
                current_protein_result = {
                    'sequence': output_seq,
                    'pdb': output_pdb,
                    'structure_view': structure_view,
                    'plddt_plot': plddt_plot,
                    'params': params
                }
                
                return history + [
                    {"role": "user", "content": message},
                    {"role": "assistant", "content": explanation}
                ]
            
            except Exception as e:
                return history + [
                    {"role": "user", "content": message},
                    {"role": "assistant", "content": f"단백질 생성 중 오류가 발생했습니다: {str(e)}"}
                ]
        else:
            return history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": "단백질 생성 관련 키워드를 포함해주세요."}
            ]
    except Exception as e:
        return history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": f"처리 중 오류가 발생했습니다: {str(e)}"}
        ]


def generate_protein(params):
    # 기존 protein_diffusion_model 함수 호출
    result = protein_diffusion_model(
        sequence=None,
        seq_len=params['sequence_length'],
        helix_bias=params['helix_bias'],
        strand_bias=params['strand_bias'],
        loop_bias=params['loop_bias'],
        secondary_structure=None,
        aa_bias=None,
        aa_bias_potential=None,
        num_steps="25",
        noise="normal",
        hydrophobic_target_score=str(params['hydrophobic_target_score']),
        hydrophobic_potential="2",
        contigs=None,
        pssm=None,
        seq_mask=None,
        str_mask=None,
        rewrite_pdb=None
    )
    return result

def generate_explanation(result, params):
    explanation = f"""
    생성된 단백질 분석:
    - 길이: {params['sequence_length']} 아미노산
    - 구조적 특징:
        * 알파 나선 비율: {params['helix_bias']*100}%
        * 베타 시트 비율: {params['strand_bias']*100}%
        * 루프 구조 비율: {params['loop_bias']*100}%
    - 특수 기능: {result.get('special_features', '없음')}
    """
    return explanation

# 체크포인트 파일 경로를 절대 경로로 수정
def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bias, 
                    secondary_structure, aa_bias, aa_bias_potential, 
                    num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
                    contigs, pssm, seq_mask, str_mask, rewrite_pdb):

    
    dssp_checkpoint = './SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'
    og_checkpoint = './SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'

    
    # 체크포인트 파일 존재 확인
    if not os.path.exists(dssp_checkpoint):
        raise FileNotFoundError(f"DSSP checkpoint file not found at: {dssp_checkpoint}")
    if not os.path.exists(og_checkpoint):
        raise FileNotFoundError(f"OG checkpoint file not found at: {og_checkpoint}")
    
    model_args = copy.deepcopy(args)


    # make sampler
    S = HuggingFace_sampler(args=model_args)

    # get random prefix 
    S.out_prefix = './tmp/'+secrets.token_hex(nbytes=10).upper()

    # set args
    S.args['checkpoint'] = None
    S.args['dump_trb'] = False
    S.args['dump_args'] = True
    S.args['save_best_plddt'] = True
    S.args['T'] = 20
    S.args['strand_bias'] = 0.0
    S.args['loop_bias'] = 0.0
    S.args['helix_bias'] = 0.0
    S.args['potentials'] = None
    S.args['potential_scale'] = None
    S.args['aa_composition'] = None


    # get sequence if entered and make sure all chars are valid
    alt_aa_dict = {'B':['D','N'],'J':['I','L'],'U':['C'],'Z':['E','Q'],'O':['K']}
    if sequence not in ['',None]:
        L = len(sequence)
        aa_seq = []
        for aa in sequence.upper():
            if aa in alt_aa_dict.keys():
                aa_seq.append(np.random.choice(alt_aa_dict[aa]))
            else:
                aa_seq.append(aa)

        S.args['sequence'] = aa_seq
    elif contigs not in ['',None]:
        S.args['contigs'] = [contigs]
    else:
        S.args['contigs'] = [f'{seq_len}']
        L = int(seq_len)
    
    print('DEBUG: ',rewrite_pdb)
    if rewrite_pdb not in ['',None]:
        S.args['pdb'] = rewrite_pdb.name

    if seq_mask not in ['',None]:
        S.args['inpaint_seq'] = [seq_mask]
    if str_mask not in ['',None]:
        S.args['inpaint_str'] = [str_mask]

    if secondary_structure in ['',None]:
        secondary_structure = None
    else:
        secondary_structure = ''.join(['E' if x == 'S' else x for x in secondary_structure])
        if L < len(secondary_structure):
            secondary_structure = secondary_structure[:len(sequence)]
        elif L == len(secondary_structure):
            pass
        else:
            dseq = L - len(secondary_structure)
            secondary_structure += secondary_structure[-1]*dseq
    

    # potentials
    potential_list = []
    potential_bias_list = []

    if aa_bias not in ['',None]:
        potential_list.append('aa_bias')
        S.args['aa_composition'] = aa_bias
        if aa_bias_potential in ['',None]:
            aa_bias_potential = 3
        potential_bias_list.append(str(aa_bias_potential))
    '''
    if target_charge not in ['',None]:
        potential_list.append('charge')
        if charge_potential in ['',None]:
            charge_potential = 1
        potential_bias_list.append(str(charge_potential))
        S.args['target_charge'] = float(target_charge)
        if target_ph in ['',None]:
            target_ph = 7.4
        S.args['target_pH'] = float(target_ph)
    '''
    
    if hydrophobic_target_score not in ['',None]:
        potential_list.append('hydrophobic')
        S.args['hydrophobic_score'] = float(hydrophobic_target_score)
        if hydrophobic_potential in ['',None]:
            hydrophobic_potential = 3
        potential_bias_list.append(str(hydrophobic_potential))
    
    if pssm not in ['',None]:
        potential_list.append('PSSM')
        potential_bias_list.append('5')
        S.args['PSSM'] = pssm.name
        

    if len(potential_list) > 0:
        S.args['potentials'] = ','.join(potential_list)
        S.args['potential_scale'] = ','.join(potential_bias_list)


    # normalise secondary_structure bias from range 0-0.3
    S.args['secondary_structure'] = secondary_structure
    S.args['helix_bias'] = helix_bias
    S.args['strand_bias'] = strand_bias
    S.args['loop_bias'] = loop_bias
    
    # set T
    if num_steps in ['',None]:
        S.args['T'] = 20
    else:
        S.args['T'] = int(num_steps)

    # noise
    if 'normal' in noise:
        S.args['sample_distribution'] = noise
        S.args['sample_distribution_gmm_means'] = [0]
        S.args['sample_distribution_gmm_variances'] = [1]
    elif 'gmm2' in noise:
        S.args['sample_distribution'] = noise
        S.args['sample_distribution_gmm_means'] = [-1,1]
        S.args['sample_distribution_gmm_variances'] = [1,1]
    elif 'gmm3' in noise:
        S.args['sample_distribution'] = noise
        S.args['sample_distribution_gmm_means'] = [-1,0,1]
        S.args['sample_distribution_gmm_variances'] = [1,1,1]



    if secondary_structure not in ['',None] or helix_bias+strand_bias+loop_bias > 0:
        S.args['checkpoint'] = dssp_checkpoint
        S.args['d_t1d'] = 29
        print('using dssp checkpoint')
    else:
        S.args['checkpoint'] = og_checkpoint
        S.args['d_t1d'] = 24
        print('using og checkpoint')
    

    for k,v in S.args.items():
        print(f"{k} --> {v}")
    
    # init S
    S.model_init()
    S.diffuser_init()
    S.setup()

    # sampling loop
    plddt_data = []
    for j in range(S.max_t):
        print(f'on step {j}')
        output_seq, output_pdb, plddt = S.take_step_get_outputs(j)
        plddt_data.append(plddt)
        yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
    
    output_seq, output_pdb, plddt = S.get_outputs()
 
    return output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)


def get_plddt_plot(plddt_data, max_t):
    x = [i+1 for i in range(len(plddt_data))]
    fig, ax = plt.subplots(figsize=(15,6))
    ax.plot(x,plddt_data,color='#661dbf', linewidth=3,marker='o')
    ax.set_xticks([i+1 for i in range(max_t)])
    ax.set_yticks([(i+1)/10 for i in range(10)])
    ax.set_ylim([0,1])
    ax.set_ylabel('model confidence (plddt)')
    ax.set_xlabel('diffusion steps (t)')
    return fig

def display_pdb(path_to_pdb):
    '''
        #function to display pdb in py3dmol
    '''
    pdb = open(path_to_pdb, "r").read()
    
    view = py3Dmol.view(width=500, height=500)
    view.addModel(pdb, "pdb")
    view.setStyle({'model': -1}, {"cartoon": {'colorscheme':{'prop':'b','gradient':'roygb','min':0,'max':1}}})#'linear', 'min': 0, 'max': 1, 'colors': ["#ff9ef0","#a903fc",]}}}) 
    view.zoomTo()
    output = view._make_html().replace("'", '"')
    print(view._make_html())
    x = f"""<!DOCTYPE html><html></center> {output} </center></html>"""  # do not use ' in this input
    
    return f"""<iframe height="500px" width="100%"  name="result" allow="midi; geolocation; microphone; camera;
                            display-capture; encrypted-media;" sandbox="allow-modals allow-forms
                            allow-scripts allow-same-origin allow-popups
                            allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
                            allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""

'''
    return f"""<iframe  style="width: 100%; height:700px" name="result" allow="midi; geolocation; microphone; camera; 
                            display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
                            allow-scripts allow-same-origin allow-popups 
                            allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
                            allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
'''

def get_motif_preview(pdb_id, contigs):
    try:
        input_pdb = fetch_pdb(pdb_id=pdb_id.lower() if pdb_id else None)
        if input_pdb is None:
            return gr.HTML("PDB ID를 입력해주세요"), None
        
        parse = parse_pdb(input_pdb)
        output_name = input_pdb

        pdb = open(output_name, "r").read()
        view = py3Dmol.view(width=500, height=500)
        view.addModel(pdb, "pdb")

        if contigs in ['',0]:
            contigs = ['0']
        else:
            contigs = [contigs]

        print('DEBUG: ',contigs)
        
        pdb_map = get_mappings(ContigMap(parse,contigs))
        print('DEBUG: ',pdb_map)
        print('DEBUG: ',pdb_map['con_ref_idx0'])
        roi = [x[1]-1 for x in pdb_map['con_ref_pdb_idx']]

        colormap = {0:'#D3D3D3', 1:'#F74CFF'}
        colors = {i+1: colormap[1] if i in roi else colormap[0] for i in range(parse['xyz'].shape[0])}
        view.setStyle({"cartoon": {"colorscheme": {"prop": "resi", "map": colors}}})
        view.zoomTo()
        output = view._make_html().replace("'", '"')
        print(view._make_html())
        x = f"""<!DOCTYPE html><html></center> {output} </center></html>"""  # do not use ' in this input
        
        return f"""<iframe height="500px" width="100%"  name="result" allow="midi; geolocation; microphone; camera;
                                display-capture; encrypted-media;" sandbox="allow-modals allow-forms
                                allow-scripts allow-same-origin allow-popups
                                allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
                                allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>""", output_name

    except Exception as e:
        return gr.HTML(f"오류가 발생했습니다: {str(e)}"), None

def fetch_pdb(pdb_id=None):
    if pdb_id is None or pdb_id == "":
        return None
    else:
        os.system(f"wget -qnc https://files.rcsb.org/view/{pdb_id}.pdb")
        return f"{pdb_id}.pdb"

# MSA AND PSSM GUIDANCE
def save_pssm(file_upload):
    filename = file_upload.name
    orig_name = file_upload.orig_name
    if filename.split('.')[-1] in ['fasta', 'a3m']:
        return msa_to_pssm(file_upload)
    return filename

def msa_to_pssm(msa_file):
    # Define the lookup table for converting amino acids to indices
    aa_to_index = {'A': 0, 'R': 1, 'N': 2, 'D': 3, 'C': 4, 'Q': 5, 'E': 6, 'G': 7, 'H': 8, 'I': 9, 'L': 10,
                'K': 11, 'M': 12, 'F': 13, 'P': 14, 'S': 15, 'T': 16, 'W': 17, 'Y': 18, 'V': 19, 'X': 20, '-': 21}
    # Open the FASTA file and read the sequences
    records = list(SeqIO.parse(msa_file.name, "fasta"))

    assert len(records) >= 1, "MSA must contain more than one protein sequecne."

    first_seq = str(records[0].seq)
    aligned_seqs = [first_seq]
    # print(aligned_seqs)
    # Perform sequence alignment using the Needleman-Wunsch algorithm
    aligner = Align.PairwiseAligner()
    aligner.open_gap_score = -0.7
    aligner.extend_gap_score = -0.3
    for record in records[1:]:
        alignment = aligner.align(first_seq, str(record.seq))[0]
        alignment = alignment.format().split("\n")
        al1 = alignment[0]
        al2 = alignment[2]
        al1_fin = ""
        al2_fin = ""
        percent_gap = al2.count('-')/ len(al2)
        if percent_gap > 0.4:
            continue
        for i in range(len(al1)):
            if al1[i] != '-':
                al1_fin += al1[i]
                al2_fin += al2[i]
        aligned_seqs.append(str(al2_fin))
    # Get the length of the aligned sequences
    aligned_seq_length = len(first_seq)
    # Initialize the position scoring matrix
    matrix = np.zeros((22, aligned_seq_length))
    # Iterate through the aligned sequences and count the amino acids at each position
    for seq in aligned_seqs:
        #print(seq)
        for i in range(aligned_seq_length):
            if i == len(seq):
                break
            amino_acid = seq[i]
            if amino_acid.upper() not in aa_to_index.keys():
                continue
            else:
                aa_index = aa_to_index[amino_acid.upper()]
            matrix[aa_index, i] += 1
    # Normalize the counts to get the frequency of each amino acid at each position
    matrix /= len(aligned_seqs)
    print(len(aligned_seqs))
    matrix[20:,]=0

    outdir = ".".join(msa_file.name.split('.')[:-1]) + ".csv"
    np.savetxt(outdir, matrix[:21,:].T, delimiter=",")
    return outdir

def get_pssm(fasta_msa, input_pssm):
    try:
        if input_pssm is not None:
            outdir = input_pssm.name
        elif fasta_msa is not None:
            outdir = save_pssm(fasta_msa)
        else:
            return gr.Plot(label="파일을 업로드해주세요"), None

        pssm = np.loadtxt(outdir, delimiter=",", dtype=float)
        fig, ax = plt.subplots(figsize=(15,6))
        plt.imshow(torch.permute(torch.tensor(pssm),(1,0)))
        return fig, outdir
    except Exception as e:
        return gr.Plot(label=f"오류가 발생했습니다: {str(e)}"), None

# 히어로 능력치 계산 함수 추가
def calculate_hero_stats(helix_bias, strand_bias, loop_bias, hydrophobic_score):
    stats = {
        'strength': strand_bias * 20,  # 베타시트 구조 기반
        'flexibility': helix_bias * 20, # 알파헬릭스 구조 기반
        'speed': loop_bias * 5,        # 루프 구조 기반
        'defense': abs(hydrophobic_score) if hydrophobic_score else 0
    }
    return stats

def toggle_seq_input(choice):
    if choice == "자동 설계":
        return gr.update(visible=True), gr.update(visible=False)
    else:  # "직접 입력"
        return gr.update(visible=False), gr.update(visible=True)

def toggle_secondary_structure(choice):
    if choice == "슬라이더로 설정":
        return (
            gr.update(visible=True),  # helix_bias
            gr.update(visible=True),  # strand_bias
            gr.update(visible=True),  # loop_bias
            gr.update(visible=False)  # secondary_structure
        )
    else:  # "직접 입력"
        return (
            gr.update(visible=False),  # helix_bias
            gr.update(visible=False),  # strand_bias
            gr.update(visible=False),  # loop_bias
            gr.update(visible=True)   # secondary_structure
        )


def create_radar_chart(stats):
    # 레이더 차트 생성 로직
    categories = list(stats.keys())
    values = list(stats.values())
    
    fig = go.Figure(data=go.Scatterpolar(
        r=values,
        theta=categories,
        fill='toself'
    ))
    
    fig.update_layout(
        polar=dict(
            radialaxis=dict(
                visible=True,
                range=[0, 1]
            )),
        showlegend=False
    )
    
    return fig

def generate_hero_description(name, stats, abilities):
    # 히어로 설명 생성 로직
    description = f"""
    히어로 이름: {name}
    
    주요 능력:
    - 근력: {'★' * int(stats['strength'] * 5)}
    - 유연성: {'★' * int(stats['flexibility'] * 5)}
    - 스피드: {'★' * int(stats['speed'] * 5)}
    - 방어력: {'★' * int(stats['defense'] * 5)}
    
    특수 능력: {', '.join(abilities)}
    """
    return description

def combined_generation(name, strength, flexibility, speed, defense, size, abilities,
                       sequence, seq_len, helix_bias, strand_bias, loop_bias, 
                       secondary_structure, aa_bias, aa_bias_potential, 
                       num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
                       contigs, pssm, seq_mask, str_mask, rewrite_pdb):
    try:
        # protein_diffusion_model 실행
        generator = protein_diffusion_model(
            sequence=None,
            seq_len=size,  # 히어로 크기를 seq_len으로 사용
            helix_bias=flexibility,  # 히어로 유연성을 helix_bias로 사용
            strand_bias=strength,    # 히어로 강도를 strand_bias로 사용
            loop_bias=speed,         # 히어로 스피드를 loop_bias로 사용
            secondary_structure=None,
            aa_bias=None,
            aa_bias_potential=None,
            num_steps="25",
            noise="normal",
            hydrophobic_target_score=str(-defense),  # 히어로 방어력을 hydrophobic score로 사용
            hydrophobic_potential="2",
            contigs=None,
            pssm=None,
            seq_mask=None,
            str_mask=None,
            rewrite_pdb=None
        )
        
        # 마지막 결과 가져오기
        final_result = None
        for result in generator:
            final_result = result
            
        if final_result is None:
            raise Exception("생성 결과가 없습니다")
            
        output_seq, output_pdb, structure_view, plddt_plot = final_result
        
        # 히어로 능력치 계산
        stats = calculate_hero_stats(flexibility, strength, speed, defense)
        
        # 모든 결과 반환
        return (
            create_radar_chart(stats),  # 능력치 차트
            generate_hero_description(name, stats, abilities),  # 히어로 설명
            output_seq,  # 단백질 서열
            output_pdb,  # PDB 파일
            structure_view,  # 3D 구조
            plddt_plot  # 신뢰도 차트
        )
    except Exception as e:
        print(f"Error in combined_generation: {str(e)}")
        return (
            None,
            f"에러: {str(e)}",
            None,
            None,
            gr.HTML("에러가 발생했습니다"),
            None
        )


def extract_parameters_from_chat(chat_response):
    """챗봇 응답에서 파라미터 추출"""
    try:
        params = {
            'sequence_length': 100,
            'helix_bias': 0.02,
            'strand_bias': 0.02,
            'loop_bias': 0.1,
            'hydrophobic_target_score': 0
        }
        
        # 응답 텍스트에서 값 추출
        if "길이:" in chat_response:
            length_match = re.search(r'길이: (\d+)', chat_response)
            if length_match:
                params['sequence_length'] = int(length_match.group(1))
                
        if "알파 헬릭스 비율:" in chat_response:
            helix_match = re.search(r'알파 헬릭스 비율: ([\d.]+)', chat_response)
            if helix_match:
                params['helix_bias'] = float(helix_match.group(1)) / 100
                
        if "베타 시트 비율:" in chat_response:
            strand_match = re.search(r'베타 시트 비율: ([\d.]+)', chat_response)
            if strand_match:
                params['strand_bias'] = float(strand_match.group(1)) / 100
                
        if "루프 구조 비율:" in chat_response:
            loop_match = re.search(r'루프 구조 비율: ([\d.]+)', chat_response)
            if loop_match:
                params['loop_bias'] = float(loop_match.group(1)) / 100
                
        if "소수성 점수:" in chat_response:
            hydro_match = re.search(r'소수성 점수: ([-\d.]+)', chat_response)
            if hydro_match:
                params['hydrophobic_target_score'] = float(hydro_match.group(1))
        
        return params
    except Exception as e:
        print(f"파라미터 추출 중 오류: {str(e)}")
        return None

def update_protein_display(chat_response):
    if "생성된 단백질 분석" in chat_response:
        params = extract_parameters_from_chat(chat_response)
        if params:
            result = generate_protein(params)
            stats = calculate_hero_stats(
                helix_bias=params['helix_bias'],
                strand_bias=params['strand_bias'],
                loop_bias=params['loop_bias'],
                hydrophobic_score=params['hydrophobic_target_score']
            )
            return {
                hero_stats: create_radar_chart(stats),
                hero_description: chat_response,
                output_seq: result[0],
                output_pdb: result[1],
                output_viewer: display_pdb(result[1]),
                plddt_plot: result[3]
            }
    return None

def process_chat_and_generate(message, history):
    try:
        # 1. 프롬프트 분석 및 파라미터 설정
        analysis = analyze_prompt(message)
        similar_structures = search_protein_data(analysis, ds)
        params = extract_parameters(analysis, similar_structures)
        
        # 2. 단백질 생성
        generator = protein_diffusion_model(
            sequence=None,
            seq_len=params['sequence_length'],
            helix_bias=params['helix_bias'],
            strand_bias=params['strand_bias'],
            loop_bias=params['loop_bias'],
            secondary_structure=None,
            aa_bias=None,
            aa_bias_potential=None,
            num_steps="25",  # 25단계 설정
            noise="normal",
            hydrophobic_target_score=str(params['hydrophobic_target_score']),
            hydrophobic_potential="2",
            contigs=None,
            pssm=None,
            seq_mask=None,
            str_mask=None,
            rewrite_pdb=None
        )
        
        # 3. 모든 단계의 결과 수집
        final_result = None
        for result in generator:
            final_result = result
            # 중간 결과 업데이트 (선택적)
            yield (
                history + [
                    {"role": "user", "content": message},
                    {"role": "assistant", "content": f"단백질 생성 중... {len(result[3].get_lines()) if result[3] else 0}단계 완료"}
                ],
                create_radar_chart(calculate_hero_stats(
                    params['helix_bias'],
                    params['strand_bias'],
                    params['loop_bias'],
                    params['hydrophobic_target_score']
                )),
                f"단백질 생성 진행 중... {len(result[3].get_lines()) if result[3] else 0}단계 완료",
                result[0],  # output_seq
                result[1],  # output_pdb
                result[2],  # structure_view
                result[3]   # plddt_plot
            )
        
        if final_result is None:
            raise Exception("생성 결과가 없습니다")
            
        output_seq, output_pdb, structure_view, plddt_plot = final_result
        
        # 4. 히어로 능력치 계산
        stats = calculate_hero_stats(
            helix_bias=params['helix_bias'],
            strand_bias=params['strand_bias'],
            loop_bias=params['loop_bias'],
            hydrophobic_score=params['hydrophobic_target_score']
        )
        
        # 5. 최종 설명 생성
        explanation = f"""
        요청하신 기능에 맞는 단백질을 생성했습니다:
        
        분석된 요구사항:
        {analysis}
        
        설계된 구조적 특징:
        - 길이: {params['sequence_length']} 아미노산
        - 알파 헬릭스 비율: {params['helix_bias']*100:.1f}%
        - 베타 시트 비율: {params['strand_bias']*100:.1f}%
        - 루프 구조 비율: {params['loop_bias']*100:.1f}%
        - 소수성 점수: {params['hydrophobic_target_score']}
        
        생성이 완료되었습니다. 총 {len(plddt_plot.get_lines()) if plddt_plot else 0}단계의 최적화를 수행했습니다.
        """
        
        # 6. 최종 결과 반환
        return (
            history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": explanation}
            ],
            create_radar_chart(stats),
            explanation,
            output_seq,
            output_pdb,
            structure_view,
            plddt_plot
        )
        
    except Exception as e:
        print(f"Error in process_chat_and_generate: {str(e)}")
        return (
            history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": f"오류가 발생했습니다: {str(e)}"}
            ],
            None, None, None, None, None, None
        )

def extract_keywords(analysis):
    """분석 텍스트에서 키워드 추출"""
    try:
        # 기본 키워드 추출
        keywords = []
        # 주요 기능 키워드
        if "치료" in analysis: keywords.extend(["therapeutic", "binding"])
        if "결합" in analysis: keywords.extend(["binding", "interaction"])
        if "촉매" in analysis: keywords.extend(["enzyme", "catalytic"])
        
        # 환경 키워드
        if "막" in analysis: keywords.extend(["membrane", "transmembrane"])
        if "수용성" in analysis: keywords.extend(["soluble", "hydrophilic"])
        if "소수성" in analysis: keywords.extend(["hydrophobic"])
        
        # 구조 키워드
        if "알파" in analysis or "나선" in analysis: keywords.append("helix")
        if "베타" in analysis or "시트" in analysis: keywords.append("sheet")
        if "루프" in analysis: keywords.append("loop")
        
        return list(set(keywords))  # 중복 제거
    except Exception as e:
        print(f"키워드 추출 중 오류: {str(e)}")
        return []

def calculate_similarity(keywords, entry):
    """키워드와 데이터셋 항목 간의 유사도 계산"""
    try:
        score = 0
        # 데이터셋 구조 확인 및 안전한 접근
        sequence = entry.get('sequence', '').lower() if isinstance(entry, dict) else str(entry).lower()
        
        # 데이터셋 구조 디버깅
        print("Entry structure:", type(entry))
        print("Entry content:", entry)
        
        for keyword in keywords:
            # 안전한 접근을 위한 수정
            description = entry.get('description', '') if isinstance(entry, dict) else ''
            if keyword in description.lower():
                score += 2
            if keyword in sequence:
                score += 1
            if isinstance(entry, dict) and 'secondary_structure' in entry:
                sec_structure = entry['secondary_structure']
                if keyword in ['helix'] and 'H' in sec_structure:
                    score += 1
                if keyword in ['sheet'] and 'E' in sec_structure:
                    score += 1
                if keyword in ['loop'] and 'L' in sec_structure:
                    score += 1
        return score
    except Exception as e:
        print(f"유사도 계산 중 상세 오류: {str(e)}")
        print("Entry:", entry)
        return 0



def download_checkpoint_files():
    """필요한 체크포인트 파일 다운로드"""
    try:
        import requests
        
        # 체크포인트 파일 URL (실제 URL로 교체 필요)
        dssp_url = "YOUR_DSSP_CHECKPOINT_URL"
        og_url = "YOUR_OG_CHECKPOINT_URL"
        
        # DSSP 체크포인트 다운로드
        if not os.path.exists(dssp_checkpoint):
            print("Downloading DSSP checkpoint...")
            response = requests.get(dssp_url)
            with open(dssp_checkpoint, 'wb') as f:
                f.write(response.content)
        
        # OG 체크포인트 다운로드
        if not os.path.exists(og_checkpoint):
            print("Downloading OG checkpoint...")
            response = requests.get(og_url)
            with open(og_checkpoint, 'wb') as f:
                f.write(response.content)
                
        print("Checkpoint files downloaded successfully")
    except Exception as e:
        print(f"Error downloading checkpoint files: {str(e)}")
        raise

# 시작 시 체크포인트 파일 확인 및 다운로드
try:
    download_checkpoint_files()
except Exception as e:
    print(f"Warning: Could not download checkpoint files: {str(e)}")

with gr.Blocks(theme='ParityError/Interstellar') as demo:
    with gr.Row():
        with gr.Column(scale=1):
            # 챗봇 인터페이스
            gr.Markdown("# 🤖 AI 단백질 설계 도우미")
            # 여기를 수정
            chatbot = gr.Chatbot(
                height=600,
                type='messages'  # 메시지 형식 지정
          
            )
            with gr.Row():
                msg = gr.Textbox(
                    label="메시지를 입력하세요",
                    placeholder="예: COVID-19를 치료할 수 있는 단백질을 생성해주세요",
                    lines=2,
                    scale=4
                )
                submit_btn = gr.Button("전송", variant="primary", scale=1)
            clear = gr.Button("대화 내용 지우기")



            
            with gr.Accordion("채팅 설정", open=False):
                system_message = gr.Textbox(
                    value="당신은 단백질 설계를 도와주는 전문가입니다.",
                    label="시스템 메시지"
                )
                max_tokens = gr.Slider(
                    minimum=1,
                    maximum=2048,
                    value=512,
                    step=1,
                    label="최대 토큰 수"
                )
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=4.0,
                    value=0.7,
                    step=0.1,
                    label="Temperature"
                )
                top_p = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top-P"
                )


            # 탭 인터페이스
            with gr.Tabs():
                with gr.TabItem("🦸‍♂️ 히어로 디자인"):
                    gr.Markdown("""
                    ### ✨ 당신만의 특별한 히어로를 만들어보세요!
                    각 능력치를 조절하면 히어로의 DNA가 자동으로 설계됩니다.
                    """)
                    
                    # 히어로 기본 정보
                    hero_name = gr.Textbox(
                        label="히어로 이름", 
                        placeholder="당신의 히어로 이름을 지어주세요!",
                        info="히어로의 정체성을 나타내는 이름을 입력하세요"
                    )
                    
                    # 능력치 설정
                    gr.Markdown("### 💪 히어로 능력치 설정")
                    with gr.Row():
                        strength = gr.Slider(
                            minimum=0.0, maximum=0.05, 
                            label="💪 초강력(근력)", 
                            value=0.02,
                            info="단단한 베타시트 구조로 강력한 힘을 생성합니다"
                        )
                        flexibility = gr.Slider(
                            minimum=0.0, maximum=0.05, 
                            label="🤸‍♂️ 유연성", 
                            value=0.02,
                            info="나선형 알파헬릭스 구조로 유연한 움직임을 가능하게 합니다"
                        )
                    
                    with gr.Row():
                        speed = gr.Slider(
                            minimum=0.0, maximum=0.20, 
                            label="⚡ 스피드", 
                            value=0.1,
                            info="루프 구조로 빠른 움직임을 구현합니다"
                        )
                        defense = gr.Slider(
                            minimum=-10, maximum=10, 
                            label="🛡️ 방어력", 
                            value=0,
                            info="음수: 수중 활동에 특화, 양수: 지상 활동에 특화"
                        )
                    
                    # 히어로 크기 설정
                    hero_size = gr.Slider(
                        minimum=50, maximum=200, 
                        label="📏 히어로 크기", 
                        value=100,
                        info="히어로의 전체적인 크기를 결정합니다"
                    )
                    
                    # 특수 능력 설정
                    with gr.Accordion("🌟 특수 능력", open=False):
                        gr.Markdown("""
                        특수 능력을 선택하면 히어로의 DNA에 특별한 구조가 추가됩니다.
                        - 자가 회복: 단백질 구조 복구 능력 강화
                        - 원거리 공격: 특수한 구조적 돌출부 형성
                        - 방어막 생성: 안정적인 보호층 구조 생성
                        """)
                        special_ability = gr.CheckboxGroup(
                            choices=["자가 회복", "원거리 공격", "방어막 생성"],
                            label="특수 능력 선택"
                        )
                    
                    # 생성 버튼
                    create_btn = gr.Button("🧬 히어로 생성!", variant="primary", scale=2)

                with gr.TabItem("🧬 히어로 DNA 설계"):
                    gr.Markdown("""
                    ### 🧪 히어로 DNA 고급 설정
                    히어로의 유전자 구조를 더 세밀하게 조정할 수 있습니다.
                    """)
                    
                    seq_opt = gr.Radio(
                        ["자동 설계", "직접 입력"],
                        label="DNA 설계 방식",
                        value="자동 설계"
                    )

                    sequence = gr.Textbox(
                        label="DNA 시퀀스", 
                        lines=1, 
                        placeholder='사용 가능한 아미노산: A,C,D,E,F,G,H,I,K,L,M,N,P,Q,R,S,T,V,W,Y (X는 무작위)',
                        visible=False
                    )
                    seq_len = gr.Slider(
                        minimum=5.0, maximum=250.0, 
                        label="DNA 길이", 
                        value=100, 
                        visible=True
                    )
                    
                    with gr.Accordion(label='🦴 골격 구조 설정', open=True):
                        gr.Markdown("""
                        히어로의 기본 골격 구조를 설정합니다.
                        - 나선형 구조: 유연하고 탄력있는 움직임
                        - 병풍형 구조: 단단하고 강력한 힘
                        - 고리형 구조: 빠르고 민첩한 움직임
                        """)
                        sec_str_opt = gr.Radio(
                            ["슬라이더로 설정", "직접 입력"],
                            label="골격 구조 설정 방식",
                            value="슬라이더로 설정"
                        )

                        secondary_structure = gr.Textbox(
                            label="골격 구조",
                            lines=1,
                            placeholder='H:나선형, S:병풍형, L:고리형, X:자동설정',
                            visible=False
                        )
                        
                        with gr.Column():
                            helix_bias = gr.Slider(
                                minimum=0.0, maximum=0.05,
                                label="나선형 구조 비율",
                                visible=True
                            )
                            strand_bias = gr.Slider(
                                minimum=0.0, maximum=0.05,
                                label="병풍형 구조 비율",
                                visible=True
                            )
                            loop_bias = gr.Slider(
                                minimum=0.0, maximum=0.20,
                                label="고리형 구조 비율",
                                visible=True
                            )
                    
                    with gr.Accordion(label='🧬 DNA 구성 설정', open=False):
                        gr.Markdown("""
                        특정 아미노산의 비율을 조절하여 히어로의 특성을 강화할 수 있습니다.
                        예시: W0.2,E0.1 (트립토판 20%, 글루탐산 10%)
                        """)
                        with gr.Row():
                            aa_bias = gr.Textbox(
                                label="아미노산 비율", 
                                lines=1, 
                                placeholder='예시: W0.2,E0.1'
                            )
                            aa_bias_potential = gr.Textbox(
                                label="강화 정도", 
                                lines=1, 
                                placeholder='1.0-5.0 사이 값 입력'
                            )

                    with gr.Accordion(label='🌍 환경 적응력 설정', open=False):
                        gr.Markdown("""
                        히어로의 환경 적응력을 조절합니다.
                        음수: 수중 활동에 특화, 양수: 지상 활동에 특화
                        """)
                        with gr.Row():
                            hydrophobic_target_score = gr.Textbox(
                                label="환경 적응 점수",
                                lines=1,
                                placeholder='예시: -5 (수중 활동에 특화)'
                            )
                            hydrophobic_potential = gr.Textbox(
                                label="적응력 강화 정도",
                                lines=1,
                                placeholder='1.0-2.0 사이 값 입력'
                            )

                    with gr.Accordion(label='⚙️ 고급 설정', open=False):
                        gr.Markdown("""
                        DNA 생성 과정의 세부 매개변수를 조정합니다.
                        """)
                        with gr.Row():
                            num_steps = gr.Textbox(
                                label="생성 단계",
                                lines=1,
                                placeholder='25 이하 권장'
                            )
                            noise = gr.Dropdown(
                                ['normal','gmm2 [-1,1]','gmm3 [-1,0,1]'],
                                label='노이즈 타입',
                                value='normal'
                            )
                    
                    design_btn = gr.Button("🧬 DNA 설계 생성!", variant="primary", scale=2)

                with gr.TabItem("🧪 히어로 유전자 강화"):
                    gr.Markdown("""
                    ### ⚡ 기존 히어로의 DNA 활용
                    강력한 히어로의 DNA 일부를 새로운 히어로에게 이식합니다.
                    """)
                    
                    gr.Markdown("공개된 히어로 DNA 데이터베이스에서 코드를 찾을 수 있습니다")
                    pdb_id_code = gr.Textbox(
                        label="히어로 DNA 코드",
                        lines=1,
                        placeholder='기존 히어로의 DNA 코드를 입력하세요 (예: 1DPX)'
                    )
                    
                    gr.Markdown("이식하고 싶은 DNA 영역을 선택하고 새로운 DNA를 추가할 수 있습니다")
                    contigs = gr.Textbox(
                        label="이식할 DNA 영역",
                        lines=1,
                        placeholder='예시: 15,A3-10,20-30'
                    )
                    
                    with gr.Row():
                        seq_mask = gr.Textbox(
                            label='능력 재설계',
                            lines=1,
                            placeholder='선택한 영역의 능력을 새롭게 디자인'
                        )
                        str_mask = gr.Textbox(
                            label='구조 재설계',
                            lines=1,
                            placeholder='선택한 영역의 구조를 새롭게 디자인'
                        )
                    
                    preview_viewer = gr.HTML()
                    rewrite_pdb = gr.File(label='히어로 DNA 파일')
                    preview_btn = gr.Button("🔍 미리보기", variant="secondary")
                    enhance_btn = gr.Button("⚡ 강화된 히어로 생성!", variant="primary", scale=2)

                with gr.TabItem("👑 히어로 가문"):
                    gr.Markdown("""
                    ### 🏰 위대한 히어로 가문의 유산
                    강력한 히어로 가문의 특성을 계승하여 새로운 히어로를 만듭니다.
                    """)
                    
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("히어로 가문의 DNA 정보가 담긴 파일을 업로드하세요")
                            fasta_msa = gr.File(label='가문 DNA 데이터')
                        with gr.Column():
                            gr.Markdown("이미 분석된 가문 특성 데이터가 있다면 업로드하세요")
                            input_pssm = gr.File(label='가문 특성 데이터')
                    
                    pssm = gr.File(label='분석된 가문 특성')
                    pssm_view = gr.Plot(label='가문 특성 분석 결과')
                    pssm_gen_btn = gr.Button("✨ 가문 특성 분석", variant="secondary")
                    inherit_btn = gr.Button("👑 가문의 힘 계승!", variant="primary", scale=2)

        # 오른쪽 열: 결과 표시
        with gr.Column(scale=1):
            gr.Markdown("## 🦸‍♂️ 히어로 프로필")
            hero_stats = gr.Plot(label="능력치 분석")
            hero_description = gr.Textbox(label="히어로 특성", lines=3)
            
            gr.Markdown("## 🧬 히어로 DNA 분석 결과")
            gr.Markdown("#### ⚡ DNA 안정성 점수")
            plddt_plot = gr.Plot(label='안정성 분석')
            gr.Markdown("#### 📝 DNA 시퀀스")
            output_seq = gr.Textbox(label="DNA 서열")
            gr.Markdown("#### 💾 DNA 데이터")
            output_pdb = gr.File(label="DNA 파일")
            gr.Markdown("#### 🔬 DNA 구조")
            output_viewer = gr.HTML()

    # 이벤트 연결
    # 챗봇 이벤트
    msg.submit(process_chat, [msg, chatbot], [chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

    seq_opt.change(
        fn=toggle_seq_input,
        inputs=[seq_opt],
        outputs=[seq_len, sequence],
        queue=False
    )



    sec_str_opt.change(
        fn=toggle_secondary_structure,
        inputs=[sec_str_opt],
        outputs=[helix_bias, strand_bias, loop_bias, secondary_structure],
        queue=False
    )

    preview_btn.click(
        get_motif_preview,
        inputs=[pdb_id_code, contigs],
        outputs=[preview_viewer, rewrite_pdb]
    )

    pssm_gen_btn.click(
        get_pssm,
        inputs=[fasta_msa, input_pssm],
        outputs=[pssm_view, pssm]
    )

    # 챗봇 기반 단백질 생성 결과 업데이트
    def update_protein_display(chat_response):
        if "생성된 단백질 분석" in chat_response:
            params = extract_parameters_from_chat(chat_response)
            result = generate_protein(params)
            return {
                hero_stats: create_radar_chart(calculate_hero_stats(params)),
                hero_description: chat_response,
                output_seq: result[0],
                output_pdb: result[1],
                output_viewer: display_pdb(result[1]),
                plddt_plot: result[3]
            }
        return None

    # 각 생성 버튼 이벤트 연결
    for btn in [create_btn, design_btn, enhance_btn, inherit_btn]:
        btn.click(
            combined_generation,
            inputs=[
                hero_name, strength, flexibility, speed, defense, hero_size, special_ability,
                sequence, seq_len, helix_bias, strand_bias, loop_bias, 
                secondary_structure, aa_bias, aa_bias_potential, 
                num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
                contigs, pssm, seq_mask, str_mask, rewrite_pdb
            ],
            outputs=[
                hero_stats, 
                hero_description, 
                output_seq,
                output_pdb,
                output_viewer,
                plddt_plot
            ]
        )

    # 이벤트 핸들러 연결
    msg.submit(
        fn=process_chat_and_generate,
        inputs=[msg, chatbot],
        outputs=[
            chatbot,
            hero_stats,
            hero_description,
            output_seq,
            output_pdb,
            output_viewer,
            plddt_plot
        ]
    )

    submit_btn.click(
        fn=process_chat_and_generate,
        inputs=[msg, chatbot],
        outputs=[
            chatbot,
            hero_stats,
            hero_description,
            output_seq,
            output_pdb,
            output_viewer,
            plddt_plot
        ]
    )

    # 채팅 내용 지우기
    clear.click(
        lambda: (None, None, None, None, None, None, None),
        None,
        [chatbot, hero_stats, hero_description, output_seq, output_pdb, output_viewer, plddt_plot],
        queue=False
    )

    # 챗봇 응답에 따른 결과 업데이트
    msg.submit(
        update_protein_display,
        inputs=[chatbot],
        outputs=[hero_stats, hero_description, output_seq, output_pdb, output_viewer, plddt_plot]
    )


    submit_btn.click(respond, 
                    [msg, chatbot, system_message, max_tokens, temperature, top_p], 
                    [chatbot])
    msg.submit(respond, 
              [msg, chatbot, system_message, max_tokens, temperature, top_p], 
              [chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

    # 이벤트 핸들러 수정
    msg.submit(
        fn=process_chat_and_generate,
        inputs=[msg, chatbot],
        outputs=[
            chatbot,
            hero_stats,
            hero_description,
            output_seq,
            output_pdb,
            output_viewer,
            plddt_plot
        ],
        show_progress=True
    ).then(
        lambda: None,  # 진행 상태 초기화
        None,
        None
    )

# 실행
demo.queue()
demo.launch(debug=True)