from nupack import *
import subprocess
import os

# ---------- 独立的辅助函数 ----------
def get_sequence(Complex, result, skip_errors=True):
    """
    获取复合物的序列。
    - 若 Complex 为 None，返回字典 {复合物名: 序列}，跳过解析失败的复合物。
    - 若 Complex 为字符串或 Complex 对象，返回该复合物的序列字符串；失败时抛出异常（除非 skip_errors=False 时由调用者处理）。
    """
    # 统一转换为字符串名称
    if Complex is not None and not isinstance(Complex, str):
        Complex = Complex.name if hasattr(Complex, 'name') else str(Complex)

    def _get_seq_from_strands(strand_names, comp_name=None):
        """从链名列表获取拼接后的序列，如果失败则抛出异常"""
        seq_parts = []
        for s in strand_names:
            if not s.isidentifier():
                raise NameError(
                    f"复合物 '{comp_name}' 中的链名 '{s}' 不是有效的 Python 标识符，"
                    "请检查链的 name 属性是否包含空格或特殊字符。"
                )
            try:
                strand_obj = eval(s)
                if not isinstance(strand_obj, Strand):
                    raise TypeError(f"'{s}' 不是 Strand 对象")
                seq_parts.append(str(strand_obj))
            except NameError:
                raise NameError(f"链名 '{s}' 未定义，请确保该变量已创建为 Strand 对象。")
        return ''.join(seq_parts)

    if Complex is None:
        # 批量模式
        sequences = {}
        for comp in result:
            name = comp.name if hasattr(comp, 'name') else str(comp)
            try:
                # 去除括号，按 '+' 分割
                inp = ''.join([ch for ch in name if ch not in '()'])
                strands = inp.split('+')
                sequences[name] = _get_seq_from_strands(strands, comp_name=name)
            except Exception as e:
                if skip_errors:
                    print(f"⚠️ 跳过复合物 '{name}'，原因：{e}")
                else:
                    raise  # 如果不跳过，则抛出异常
        return sequences
    else:
        # 单复合物模式
        inp = ''.join([ch for ch in Complex if ch not in '()'])
        strands = inp.split('+')
        return _get_seq_from_strands(strands, comp_name=Complex)

def get_structure(Complex, result, skip_errors=True):
    """
    获取复合物的 MFE 结构（已去除 '+'）。
    - 若 Complex 为 None，返回字典 {复合物名: 结构}，跳过 result 中不存在的复合物。
    - 若 Complex 为字符串或 Complex 对象，返回该复合物的结构字符串；失败时抛出异常。
    """
    if Complex is not None and not isinstance(Complex, str):
        Complex = Complex.name if hasattr(Complex, 'name') else str(Complex)

    if Complex is None:
        structures = {}
        for comp in result:
            name = comp.name if hasattr(comp, 'name') else str(comp)
            try:
                structure = result[comp].mfe[0].structure
                structures[name] = str(structure).replace('+', '')
            except Exception as e:
                if skip_errors:
                    print(f"⚠️ 跳过复合物 '{name}'，获取结构失败：{e}")
                else:
                    raise
        return structures
    else:
        structure = result[Complex].mfe[0].structure
        return str(structure).replace('+', '')

# ---------- 绘图函数（带错误跳过） ----------
def danlian(Complex, result, output_folder="/home/zsy"):
    """
    绘制单个复合物的二级结构图。如果处理过程中出现错误，打印错误信息并返回，不中断程序。
    Complex 可为复合物名称字符串或 Complex 对象。
    """
    try:
        sequence = get_sequence(Complex, result, skip_errors=False)  # 这里不跳过，让异常抛出以便捕获
        best_structure = get_structure(Complex, result, skip_errors=False)
    except Exception as e:
        name = Complex if isinstance(Complex, str) else (Complex.name if hasattr(Complex, 'name') else str(Complex))
        print(f"❌ 处理复合物 '{name}' 时出错：{e}，已跳过绘图")
        return

    # 获取用于文件名的字符串名称
    if isinstance(Complex, str):
        seq_name = Complex
    else:
        seq_name = Complex.name if hasattr(Complex, 'name') else str(Complex)

    temp_file_path = os.path.join(output_folder, f"{seq_name}.seq")
    try:
        with open(temp_file_path, "w", newline='\n') as f:
            f.write(f">{seq_name}\n")
            f.write(f"{sequence}\n")
            f.write(f"{best_structure}\n")
    except Exception as e:
        print(f"❌ 写入临时文件失败：{e}")
        return

    try:
        input_filename = f"{seq_name}.seq"
        process = subprocess.run(["RNAplot", "--o", "svg", input_filename], cwd=output_folder, capture_output=True, text=True)
        if process.returncode == 0:
            print(f"🎉 绘图成功！请前往 {output_folder} 文件夹查看 {seq_name}_ss.svg")
            os.remove(temp_file_path)
        else:
            print(f"❌ RNAplot 运行失败，报错信息：\n{process.stderr}")
    except FileNotFoundError:
        print("❌ 找不到 RNAplot 命令，请确认环境变量。")
    except Exception as e:
        print(f"❌ 绘图过程中出现未知错误：{e}")


使用方法
算全部结果
for comp in my_result:
    danlian(comp, my_result, output_folder="/home/zsy")

算单个结果
# 假设您的计算结果保存在 my_result 中
target = '(Y2+Y3)'   # 替换为您的目标复合物名称
# 输出结构
print(get_structure(target, my_result))
# 输出序列
print(get_sequence(target, my_result))
# 绘制二级结构图
danlian(target, my_result, output_folder="/home/zsy")