# 导入所需模块
import os
import numpy as np
import matplotlib.pyplot as plt
from pyplogo.extractors.structure_based import StructureExtractor
from pyplogo.visualizers.secondary_structure import SecondaryStructureVisualizer
from pyplogo.visualizers.themes import get_theme, get_all_themes
from pyplogo.utils.parsers import parse_fasta, validate_sequence
from pyplogo.utils.formatters import format_sequence, format_secondary_structure
from pyplogo.data.aa_properties import AA_PROPERTIES, AA_COLORS, SS_COLORS

# 设置 matplotlib 中文字体支持（如果需要）
plt.rcParams['font.family'] = ['DejaVu Sans', 'SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

# 初始化提取器
extractor = StructureExtractor()

# 示例 PDB 文件路径（需要替换为实际文件路径）
pdb_file = r"example/example.pdb"  # 请替换为实际文件路径
chain_id = "A"  # 指定要提取的链

try:
    # 从 PDB 文件提取结构数据
    structure_data = extractor.from_pdb(pdb_file, chain_id)
    
    print(f"提取的序列长度: {len(structure_data.sequence)}")
    print(f"二级结构长度: {len(structure_data.secondary_structure)}")
    print(f"二硫键数量: {len(structure_data.disulfide_bonds)}")
    
    # 显示前50个氨基酸和对应的二级结构
    print("\n序列 (前50):")
    print(structure_data.sequence[:50])
    
    print("\n二级结构 (前50):")
    print(structure_data.secondary_structure[:50])
    
    # 初始化可视化器
    visualizer = SecondaryStructureVisualizer(
        row_length=50,  # 每行显示50个氨基酸
        amino_acid_spacing=28,  # 氨基酸间距
        aa_font_size=12,  # 氨基酸字体大小
        aa_font_family=['DejaVu Sans', 'SimHei']  # 使用支持中文的字体
    )
    
    # 创建可视化图形
    fig = visualizer.create_figure(
        sequence=structure_data.sequence,
        secondary_structure=structure_data.secondary_structure,
        title="蛋白质二级结构可视化",
        show_residue_numbers=True,  # 显示残基编号
        show_disulfide_bonds=True,  # 显示二硫键
        disulfide_bonds=structure_data.disulfide_bonds,  # 二硫键数据
        figsize=(12, 8),  # 图像大小
    )
    
    # 显示图形
    plt.tight_layout()
    plt.show()
    
except Exception as e:
    print(f"错误: {e}")
    print("请确保提供了有效的 PDB 文件路径")