05 eGPS-oneBuilder_进化树 — 容器化系统发育树全套工作流

#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$ROOT_DIR/phylotree_builder_v0.0.1"

javac -cp "lib/*:java_tanglegram" -d java_tanglegram java_tanglegram/onebuilder/*.java java_tanglegram/tanglegram/*.java java_tanglegram/tree/alignment/indices/*.java java_tanglegram/tests/OneBuilderStandaloneTest.java java_tanglegram/tests/TanglegramStandaloneTest.java

echo "Java compilation completed."

#!/usr/bin/env bash
set -euo pipefail

script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)"
if [[ -n "${ONEBUILDER_REPO_ROOT:-}" ]]; then
  repo_root="$(cd "$ONEBUILDER_REPO_ROOT" && pwd -P)"
elif [[ -d "$PWD/phylotree_builder_v0.0.1" ]]; then
  repo_root="$(pwd -P)"
elif [[ -d "$script_dir/../phylotree_builder_v0.0.1" ]]; then
  repo_root="$(cd "$script_dir/.." && pwd -P)"
else
  echo "Error: could not locate repository root. Run this script from the repo root or set ONEBUILDER_REPO_ROOT." >&2
  exit 1
fi

source_dir="$repo_root/phylotree_builder_v0.0.1"
build_workspace="${ONEBUILDER_APPTAINER_DIR:-$repo_root/apptainer}"
stage_dir="$build_workspace/stage"
stage_app="$stage_dir/onebuilder"
stage_jre="$stage_dir/jre25"
image_name="${IMAGE_NAME:-onebuilder-0.0.1-linux64.sif}"
image_path="$build_workspace/$image_name"
java_home="${JAVA_HOME:-}"
full_rebuild="${FULL_REBUILD:-0}"
clean_stage="${CLEAN_STAGE:-$full_rebuild}"
refresh_runtime="${REFRESH_RUNTIME:-$clean_stage}"
refresh_jre="${REFRESH_JRE:-$clean_stage}"
force_prefix_rewrite="${FORCE_PREFIX_REWRITE:-$refresh_runtime}"
prefix_rewrite_scope="${PREFIX_REWRITE_SCOPE:-runtime}"

if [[ ! -d "$source_dir" ]]; then
  echo "Error: source directory not found: $source_dir" >&2
  exit 1
fi

source_runtime="$source_dir/.pixi/envs/default"

if [[ ! -d "$source_runtime" ]]; then
  echo "Error: local runtime software snapshot not found: $source_runtime" >&2
  exit 1
fi

if [[ ! -d "$source_dir/lib" ]]; then
  echo "Error: local Java lib directory not found: $source_dir/lib" >&2
  exit 1
fi

if [[ ! -x "$source_dir/third_party/mad/mad" ]]; then
  echo "Error: MAD executable not found: $source_dir/third_party/mad/mad" >&2
  exit 1
fi

if [[ ! -x "/opt/bioinfor/trimAI/bin/trimal" ]]; then
  echo "Error: trimAl executable not found: /opt/bioinfor/trimAI/bin/trimal" >&2
  exit 1
fi

if [[ -z "$java_home" ]]; then
  java_bin="$(command -v java || true)"
  if [[ -z "$java_bin" ]]; then
    echo "Error: java was not found on PATH." >&2
    exit 1
  fi
  java_home="$(cd "$(dirname "$java_bin")/.." && pwd)"
fi

javac_exe="${JAVAC_EXE:-$java_home/bin/javac}"
jlink_exe="${JLINK_EXE:-$java_home/bin/jlink}"

if [[ ! -x "$javac_exe" ]]; then
  echo "Error: javac not found: $javac_exe" >&2
  exit 1
fi

if [[ ! -x "$jlink_exe" ]]; then
  echo "Error: jlink not found: $jlink_exe" >&2
  exit 1
fi

echo "Compiling Java classes with $javac_exe"
(
  cd "$source_dir"
  "$javac_exe" -cp "lib/*:java_tanglegram" -d java_tanglegram \
    java_tanglegram/onebuilder/*.java \
    java_tanglegram/tanglegram/*.java \
    java_tanglegram/tests/OneBuilderStandaloneTest.java \
    java_tanglegram/tests/TanglegramStandaloneTest.java
)

echo "Preparing stage directory: $stage_dir"
if [[ "$clean_stage" == "1" ]]; then
  echo "FULL_REBUILD/CLEAN_STAGE requested; removing staged workspace."
  rm -rf "$stage_dir"
fi
mkdir -p "$stage_app" "$stage_jre"

echo "Syncing application files into stage"
rsync -a \
  --delete \
  --exclude='.pixi/' \
  --exclude='runtime/' \
  --exclude='__pycache__/' \
  --exclude='*.pyc' \
  --exclude='*.log' \
  --exclude='test*/' \
  --exclude='tests/tmp*/' \
  --exclude='third_party/mad2.2.zip' \
  "$source_dir/" "$stage_app/"

mkdir -p "$stage_app/runtime"
if [[ "$refresh_runtime" != "1" && -d "$stage_app/runtime/bin" ]]; then
  echo "Reusing staged runtime snapshot: $stage_app/runtime"
else
  echo "Refreshing staged runtime snapshot from $source_runtime"
  rm -rf "$stage_app/runtime"
  mkdir -p "$stage_app/runtime"
  rsync -a "$source_runtime/" "$stage_app/runtime/"
fi

rm -f "$stage_app/runtime/bin/pixi" "$stage_app/runtime/bin/pixi-global" 2>/dev/null || true

if [[ "$force_prefix_rewrite" == "1" ]]; then
  new_prefix="/opt/onebuilder/runtime"
  old_prefixes=(
    "$source_runtime"
    "/opt/BioInfo/phylotree_builder_v0.0.1/.pixi/envs/default"
    "/opt/onebuilder/.pixi/envs/default"
  )
  if [[ "$prefix_rewrite_scope" == "all" ]]; then
    prefix_rewrite_root="$stage_app"
  else
    prefix_rewrite_root="$stage_app/runtime"
  fi
  echo "Rewriting text prefixes under $prefix_rewrite_root"
  while IFS= read -r -d '' file_path; do
    if grep -Iq . "$file_path"; then
      for old_prefix in "${old_prefixes[@]}"; do
        OLD_PREFIX="$old_prefix" NEW_PREFIX="$new_prefix" perl -0pi -e 's/\Q$ENV{OLD_PREFIX}\E/$ENV{NEW_PREFIX}/g' "$file_path"
      done
    fi
  done < <(find "$prefix_rewrite_root" -type f -size -10M -print0)
else
  echo "Reusing existing staged prefix rewrite; set FORCE_PREFIX_REWRITE=1 to rescan text files."
fi

echo "Removing broken symlinks from staged runtime"
find "$stage_app" -xtype l -print -delete

if [[ "$refresh_jre" != "1" && -x "$stage_jre/bin/java" ]]; then
  echo "Reusing staged Java runtime: $stage_jre"
else
  echo "Creating Java 25 runtime with jlink"
  rm -rf "$stage_jre"
  "$jlink_exe" \
    --add-modules java.se,jdk.crypto.ec,jdk.localedata \
    --strip-debug \
    --no-man-pages \
    --no-header-files \
    --output "$stage_jre"
fi

build_runner="${APPTAINER_EXE:-apptainer}"
if ! command -v "$build_runner" >/dev/null 2>&1; then
  echo "Error: apptainer executable not found. Set APPTAINER_EXE if needed." >&2
  exit 1
fi

build_flags=()
if [[ -n "${APPTAINER_BUILD_FLAGS:-}" ]]; then
  read -r -a build_flags <<<"$APPTAINER_BUILD_FLAGS"
else
  build_flags=(--fakeroot)
fi

mkdir -p "$build_workspace"
echo "Building image: $image_path"
rm -f "$image_path"
(
  cd "$build_workspace"
  "$build_runner" build "${build_flags[@]}" "$image_path" "$script_dir/onebuilder.def"
)

echo "Built image: $image_path"

# This script generates a "gold-standard" *simulated* benchmark dataset:
# - A high-quality multiple sequence alignment (MSA) of 12 taxa for a single protein (300 aa)
# - The codon-aligned DNA sequences (900 nt), back-translated with varied synonymous codons
# - A README/说明文件 with taxa, guide tree, simulation parameters, and quick-start tips
#
# Notes:
# - Alignment is gapless by construction (no indels), which is common for single-copy conserved coding genes.
# - Diversity comes from amino-acid substitutions along a guide tree; DNA adds synonymous variation.
# - Because we control the process, this serves as a "gold standard" where the true alignment is known.

import random, math, textwrap, json, os
from collections import defaultdict, namedtuple
from pathlib import Path

random.seed(42)

# ------------------ Setup ------------------
out_dir = Path("./simu")
out_dir.mkdir(parents=True, exist_ok=True)

N_AA = 300  # length of protein
aa_letters = list("ACDEFGHIKLMNPQRSTVWY")  # standard 20 aas

# Vertebrate-like taxa + an outgroup (Drosophila) for diversity
# Using shorter names (≤10 characters) for better compatibility
taxa = [
    "Human",      # Homo sapiens
    "Chimp",      # Pan troglodytes
    "Mouse",      # Mus musculus
    "Rat",        # Rattus norvegicus
    "Cow",        # Bos taurus
    "Dog",        # Canis lupus familiaris
    "Chicken",    # Gallus gallus
    "Finch",      # Taeniopygia guttata
    "Frog",       # Xenopus tropicalis
    "Zebrafish",  # Danio rerio
    "Coelacanth", # Latimeria chalumnae
    "Fly"         # Drosophila melanogaster
]

# Species tree based on phylogenetic data (branch lengths in million years, scaled for simulation)
# Original tree: (((((Homo_sapiens:6.5,Pan_troglodytes:6.5):73.5,(Mus_musculus:14.0,Rattus_norvegicus:14.0):66.0):240.0,(Bos_taurus:82.0,Canis_lupus_familiaris:82.0):238.0,(Gallus_gallus:100.0,Taeniopygia_guttata:100.0):220.0):30.0,Xenopus_tropicalis:350.0):65.0,Latimeria_chalumnae:415.0):15.0,Danio_rerio:430.0):180.0,Drosophila_melanogaster:610.0);
# Scaled by dividing by 1000 to get reasonable mutation rates for simulation
Tree = namedtuple("Tree", "name bl children")
def leaf(name, bl): return Tree(name, bl, [])
def node(name, bl, children): return Tree(name, bl, children)

guide_tree = node("Root", 0.0, [
    node("N1", 0.180, [  # 180 Myr
        node("N2", 0.015, [  # 15 Myr
            node("N3", 0.065, [  # 65 Myr
                node("N4", 0.030, [  # 30 Myr
                    node("N5", 0.240, [  # 240 Myr
                        node("Primates", 0.0735, [  # 73.5 Myr
                            leaf("Human", 0.0065),  # 6.5 Myr
                            leaf("Chimp", 0.0065),  # 6.5 Myr
                        ]),
                        node("Rodents", 0.066, [  # 66 Myr
                            leaf("Mouse", 0.014),  # 14 Myr
                            leaf("Rat", 0.014),  # 14 Myr
                        ]),
                    ]),
                    node("Carnivores", 0.238, [  # 238 Myr
                        leaf("Cow", 0.082),  # 82 Myr
                        leaf("Dog", 0.082),  # 82 Myr
                    ]),
                    node("Birds", 0.220, [  # 220 Myr
                        leaf("Chicken", 0.100),  # 100 Myr
                        leaf("Finch", 0.100),  # 100 Myr
                    ]),
                ]),
                leaf("Frog", 0.350),  # 350 Myr
            ]),
            leaf("Coelacanth", 0.415),  # 415 Myr
        ]),
        leaf("Zebrafish", 0.430),  # 430 Myr
    ]),
    leaf("Fly", 0.610),  # 610 Myr
])

# ------------------ Amino-acid mutation model ------------------
# Very simple model: along a branch with length L, each site mutates Poisson(mu * L) times.
# Each mutation replaces the current aa with a random *different* aa (uniform over 19 options).
MU = 0.7  # overall rate scaling to get ~15-40% divergence among vertebrates, higher vs fly

def mutate_sequence(seq, branch_length, mu=MU):
    n = len(seq)
    # Expected mutations per site: mu * branch_length
    mutated = list(seq)
    for i in range(n):
        k = random.poisson(mu * branch_length) if hasattr(random, "poisson") else \
            sum(1 for _ in range(random.randrange(0, max(1, int(5*mu*branch_length)+1))) if False)  # fallback

    # Python's random doesn't have poisson; implement simple Poisson via Knuth algorithm
    def poisson(lmbda):
        L = math.exp(-lmbda)
        k = 0
        p = 1.0
        while True:
            k += 1
            p *= random.random()
            if p <= L:
                return k - 1

    mutated = list(seq)
    for i in range(n):
        k = poisson(MU * branch_length)
        for _ in range(k):
            # change to a different amino acid
            current = mutated[i]
            choices = [a for a in aa_letters if a != current]
            mutated[i] = random.choice(choices)
    return "".join(mutated)

# ------------------ Generate ancestral protein ------------------
# Bias composition toward typical globular proteins (rough, not exact).
aa_background = {
    'A': 0.08, 'C': 0.02, 'D': 0.05, 'E': 0.06, 'F': 0.04,
    'G': 0.07, 'H': 0.02, 'I': 0.06, 'K': 0.06, 'L': 0.09,
    'M': 0.02, 'N': 0.04, 'P': 0.05, 'Q': 0.04, 'R': 0.05,
    'S': 0.07, 'T': 0.05, 'V': 0.07, 'W': 0.01, 'Y': 0.03
}

def random_protein(n):
    letters, weights = zip(*aa_background.items())
    # cumulative selection
    import bisect
    cum = []
    c = 0.0
    for w in weights:
        c += w
        cum.append(c)
    cum = [x/c for x in cum]
    s = []
    for _ in range(n):
        r = random.random()
        idx = next(i for i,x in enumerate(cum) if r <= x)
        s.append(letters[idx])
    return "".join(s)

ancestor = random_protein(N_AA)

# ------------------ Evolve along tree ------------------
def traverse(tree, parent_seq, accum_bl, results):
    current_seq = mutate_sequence(parent_seq, tree.bl)
    if tree.children:
        for ch in tree.children:
            traverse(ch, current_seq, accum_bl + tree.bl, results)
    else:
        # leaf
        results[tree.name] = current_seq

results_protein = {}
traverse(guide_tree, ancestor, 0.0, results_protein)

# Ensure all taxa present (in case of structure mismatch)
assert set(results_protein.keys()) == set(taxa), "Taxa mismatch in simulation"

# ------------------ Back-translate to DNA with synonymous variation ------------------
# Genetic code (standard) mapping aa -> codons (excluding stop)
codon_table = {
    'A': ["GCT","GCC","GCA","GCG"],
    'C': ["TGT","TGC"],
    'D': ["GAT","GAC"],
    'E': ["GAA","GAG"],
    'F': ["TTT","TTC"],
    'G': ["GGT","GGC","GGA","GGG"],
    'H': ["CAT","CAC"],
    'I': ["ATT","ATC","ATA"],
    'K': ["AAA","AAG"],
    'L': ["TTA","TTG","CTT","CTC","CTA","CTG"],
    'M': ["ATG"],
    'N': ["AAT","AAC"],
    'P': ["CCT","CCC","CCA","CCG"],
    'Q': ["CAA","CAG"],
    'R': ["CGT","CGC","CGA","CGG","AGA","AGG"],
    'S': ["TCT","TCC","TCA","TCG","AGT","AGC"],
    'T': ["ACT","ACC","ACA","ACG"],
    'V': ["GTT","GTC","GTA","GTG"],
    'W': ["TGG"],
    'Y': ["TAT","TAC"],
}

def back_translate(protein, bias=None):
    dna = []
    for aa in protein:
        codons = codon_table[aa]
        # Optional bias per taxon (e.g., GC-rich vs AT-rich). Here we randomize lightly.
        c = random.choice(codons)
        dna.append(c)
    return "".join(dna)

results_dna = {tax: back_translate(seq) for tax, seq in results_protein.items()}

# Sanity checks
for tax, dna in results_dna.items():
    assert len(dna) == 3 * N_AA
    assert all(dna[i:i+3] in sum(codon_table.values(), []) for i in range(0, len(dna), 3))

# ------------------ Write FASTA (aligned) ------------------
def wrap_fasta(seq, width=60):
    return "\n".join(seq[i:i+width] for i in range(0, len(seq), width))

prot_fa_path = out_dir / "gold_standard_protein_aligned.fasta"
dna_fa_path  = out_dir / "gold_standard_cds_aligned.fasta"

with prot_fa_path.open("w") as f:
    for tax in taxa:
        f.write(f">{tax}\n{wrap_fasta(results_protein[tax])}\n")

with dna_fa_path.open("w") as f:
    for tax in taxa:
        f.write(f">{tax}\n{wrap_fasta(results_dna[tax])}\n")

# ------------------ Compute pairwise identity to Human for README ------------------
def pid(a, b):
    assert len(a) == len(b)
    same = sum(1 for x,y in zip(a,b) if x==y)
    return 100.0 * same / len(a)

human_prot = results_protein["Human"]
human_dna  = results_dna["Human"]

summary_rows = []
for tax in taxa:
    p1 = pid(human_prot, results_protein[tax])
    p2 = pid(human_dna, results_dna[tax])
    summary_rows.append((tax, len(results_protein[tax]), f"{p1:.1f}%", len(results_dna[tax]), f"{p2:.1f}%"))

# ------------------ README ------------------
readme_path = out_dir / "README_说明.txt"
guide_newick = "(((((Homo_sapiens:0.0065,Pan_troglodytes:0.0065):0.0735,(Mus_musculus:0.014,Rattus_norvegicus:0.014):0.066):0.240,(Bos_taurus:0.082,Canis_lupus_familiaris:0.082):0.238,(Gallus_gallus:0.100,Taeniopygia_guttata:0.100):0.220):0.030,Xenopus_tropicalis:0.350):0.065,Latimeria_chalumnae:0.415):0.015,Danio_rerio:0.430):0.180,Drosophila_melanogaster:0.610);"

readme = f"""
金标准（模拟）基准数据集：单拷贝保守蛋白基因（300 aa）
====================================================

文件列表
--------
1. 蛋白质比对（FASTA，已对齐）：{prot_fa_path.name}
2. 编码序列（FASTA，已按密码子对齐）：{dna_fa_path.name}
3. 本说明：{readme_path.name}

核心特性
--------
- 12 个物种（人、黑猩猩、小鼠、大鼠、牛、犬、鸡、斑胸草雀、爪蟾、斑马鱼、腔棘鱼、果蝇）
- 蛋白长度 300 aa；DNA 长度 900 nt（严格 3 的倍数，无终止密码子）
- 无插入缺失（gapless）但含有充分替换差异（蛋白与人类序列的身份度约 60–95%；DNA 更低）
- DNA 通过同义密码子回译，增加碱基层面的多样性；蛋白/DNA 天然等长，可直接建树
- 因为由已知树形与模型模拟得到，**对齐即为真值**（gold standard）

指南树（Newick）
----------------
{guide_newick}

仿真模型（简要）
----------------
- 祖先蛋白序列按经验氨基酸背景频率随机生成（长度 300）
- 沿上述指南树进行替换：每条分支长度 L 下，位点突变次数 ~ Poisson(μ·L)，每次随机换成不同氨基酸（均匀 19 选 1）
- 全局替换率 μ = {MU}
- 回译采用标准遗传密码表，随机选择同义密码子（引入同义差异）；未模拟插入缺失

快速上手（示例）
----------------
- IQ-TREE（蛋白）： `iqtree2 -s {prot_fa_path.name} -m MFP -bb 1000 -nt AUTO`
- IQ-TREE（DNA，密码子分区）：
  `iqtree2 -s {dna_fa_path.name} -st CODON -m MFP+MERGE -bb 1000 -nt AUTO`
- RAxML-NG（蛋白）： `raxml-ng --msa {prot_fa_path.name} --model JTT+G --bs-trees 200`
- PhyloBayes/MrBayes 等亦可直接使用该比对

与“真实基因”的差异与用途
------------------------
- 本数据集为**可控仿真**，因此适合用于：方法对比、参数调优、脚本联调、管线回归测试
- 若您需要某个**具体真实基因**（例如 FZD8）跨物种的高质量真实比对，我也可以基于公共数据库拉取正交单拷贝并做严格筛选与 MAFFT-L-INS-i / PRANK-codon 对齐

序列摘要（相对于人类序列）
--------------------------
物种\t蛋白长度(aa)\t蛋白身份度\tDNA长度(nt)\tDNA身份度
"""[1:]

for row in summary_rows:
    readme += f"{row[0]}\t{row[1]}\t{row[2]}\t{row[3]}\t{row[4]}\n"

with readme_path.open("w", encoding="utf-8") as f:
    f.write(readme)

def to_newick(nt: Tree) -> str:
    if not nt.children:
        return f"{nt.name}:{nt.bl}"
    children = ",".join(to_newick(child) for child in nt.children)
    if nt.bl:
        return f"({children}){nt.name}:{nt.bl}"
    return f"({children}){nt.name}"

print(to_newick(guide_tree) + ";")

# Return paths for the user
ret = {"protein_fasta": str(prot_fa_path), "dna_fasta": str(dna_fa_path), "readme": str(readme_path), "dir": str(out_dir)}

print(ret)

library(ape)
library(TreeDist)

normalize_language <- function(value) {
  text <- tolower(trimws(if (length(value) >= 1) value[1] else ""))
  if (text %in% c("zh", "zh-cn", "zh_cn", "chinese", "中文", "cn")) {
    return("chinese")
  }
  return("english")
}

runtime_text <- function(language, english, chinese) {
  if (normalize_language(language) == "chinese") chinese else english
}

# 读取TSV文件并提取树文件路径
read_tree_metadata <- function(tsv_file) {
  # 读取TSV文件
  metadata <- read.table(tsv_file, sep = "\t", header = FALSE, 
                        col.names = c("Method", "TreeFile"), 
                        stringsAsFactors = FALSE)
  metadata <- metadata[
    !is.na(metadata$TreeFile) &
      nzchar(trimws(metadata$TreeFile)) &
      toupper(trimws(metadata$TreeFile)) != "NULL",
    ,
    drop = FALSE
  ]
  return(metadata)
}

# 读取所有进化树
read_all_trees <- function(metadata, language = "english") {
  trees <- list()
  tree_names <- c()
  
  for (i in 1:nrow(metadata)) {
    method <- metadata$Method[i]
    tree_file <- metadata$TreeFile[i]
    
    # 检查文件是否存在
    if (file.exists(tree_file)) {
      # 读取树文件
      tree <- read.tree(tree_file)
      if (inherits(tree, "multiPhylo")) {
        cat(runtime_text(language, "Warning: multiple trees found; using first tree for:", "警告: 检测到多棵树，仅使用第一棵树:"), method, "\n")
        tree <- tree[[1]]
      }
      trees[[length(trees) + 1]] <- tree
      tree_names <- c(tree_names, method)
      cat(runtime_text(language, "Successfully read tree:", "成功读取树:"), method, "\n")
    } else {
      cat(runtime_text(language, "Warning: file does not exist:", "警告: 文件不存在:"), tree_file, "\n")
    }
  }
  
  return(list(trees = trees, names = tree_names))
}

as_scalar_distance <- function(value) {
  numeric_value <- suppressWarnings(as.numeric(value))
  numeric_value <- numeric_value[is.finite(numeric_value)]
  if (length(numeric_value) == 0) {
    return(NA_real_)
  }
  return(numeric_value[1])
}

# 计算距离矩阵
calculate_distance_matrices <- function(trees, tree_names, language = "english") {
  n_trees <- length(tree_names)
  
  # 初始化距离矩阵
  tree_distance_matrix <- matrix(0, nrow = n_trees, ncol = n_trees)
  rf_distance_matrix <- matrix(0, nrow = n_trees, ncol = n_trees)
  
  # 设置行名和列名
  rownames(tree_distance_matrix) <- tree_names
  colnames(tree_distance_matrix) <- tree_names
  rownames(rf_distance_matrix) <- tree_names
  colnames(rf_distance_matrix) <- tree_names
  
  # 计算所有唯一树对之间的距离，并镜像到对称矩阵的另一侧
  for (i in 1:(n_trees - 1)) {
    for (j in (i + 1):n_trees) {
      tree1 <- trees[[i]]
      tree2 <- trees[[j]]
      
      # 计算TreeDistance
      tree_dist <- as_scalar_distance(TreeDistance(tree1, tree2))
      tree_distance_matrix[i, j] <- tree_dist
      tree_distance_matrix[j, i] <- tree_dist
      
      # 计算Robinson-Foulds距离
      rf_dist <- as_scalar_distance(RobinsonFoulds(tree1, tree2))
      rf_distance_matrix[i, j] <- rf_dist
      rf_distance_matrix[j, i] <- rf_dist
      
      cat(runtime_text(language, "Calculated:", "计算完成:"), tree_names[i], "vs", tree_names[j], 
          "TreeDist:", tree_dist, "RF:", rf_dist, "\n")
    }
  }
  
  return(list(tree_distance = tree_distance_matrix, 
              rf_distance = rf_distance_matrix))
}

# 主函数
main <- function(tsv_file = "phylo_results/tree_meta_data.tsv", language = "english") {
  # 读取元数据
  cat(runtime_text(language, "Reading tree metadata...", "读取树文件元数据..."), "\n")
  metadata <- read_tree_metadata(tsv_file)
  print(metadata)
  
  # 读取所有树
  cat("\n", runtime_text(language, "Reading tree files...", "读取进化树文件..."), "\n", sep = "")
  tree_data <- read_all_trees(metadata, language)
  trees <- tree_data$trees
  tree_names <- tree_data$names
  
  if (length(trees) < 2) {
    stop(runtime_text(language, "At least 2 trees are required to calculate distance matrices", "至少需要2棵树才能计算距离矩阵"))
  }
  
  # 计算距离矩阵
  cat("\n", runtime_text(language, "Calculating distance matrices...", "计算距离矩阵..."), "\n", sep = "")
  distance_matrices <- calculate_distance_matrices(trees, tree_names, language)
  
  # 输出结果
  cat("\n", runtime_text(language, "=== TreeDistance matrix ===", "=== TreeDistance 距离矩阵 ==="), "\n", sep = "")
  print(distance_matrices$tree_distance)
  
  cat("\n", runtime_text(language, "=== Robinson-Foulds matrix ===", "=== Robinson-Foulds 距离矩阵 ==="), "\n", sep = "")
  print(distance_matrices$rf_distance)
  
  # 保存结果到文件
  write.table(distance_matrices$tree_distance, 
              "tree_distance_matrix.tsv",
              sep = "\t", row.names = TRUE, col.names = TRUE, quote = FALSE)
  write.table(distance_matrices$rf_distance, 
              "rf_distance_matrix.tsv",
              sep = "\t", row.names = TRUE, col.names = TRUE, quote = FALSE)
  
  cat("\n", runtime_text(language, "Distance matrices saved to:", "距离矩阵已保存到:"), "\n", sep = "")
  cat("- tree_distance_matrix.tsv\n")
  cat("- rf_distance_matrix.tsv\n")
  
  return(distance_matrices)
}

# 命令行参数处理
args <- commandArgs(trailingOnly = TRUE)

# 检查命令行参数
if (!(length(args) %in% c(1, 2))) {
  cat("Usage: Rscript cal_pair_wise_tree_dist.R <tree_meta_data.tsv> [language]\n")
  cat("Example: Rscript cal_pair_wise_tree_dist.R phylo_results/tree_meta_data.tsv english\n")
  quit(status = 1)
}

# 获取输入文件路径
tsv_file <- args[1]
language <- if (length(args) >= 2) normalize_language(args[2]) else "english"

# 检查输入文件是否存在
if (!file.exists(tsv_file)) {
  cat(runtime_text(language, "Error: input file does not exist:", "错误: 输入文件不存在:"), tsv_file, "\n")
  quit(status = 1)
}

# 运行主函数
cat(runtime_text(language, "Processing file:", "开始处理文件:"), tsv_file, "\n")
result <- main(tsv_file, language)
cat(runtime_text(language, "Done!", "处理完成!"), "\n")

#!/usr/bin/env python3
"""Parse MrBayes NEXUS trees and export Newick output."""

import re


def _is_chinese(language):
    return str(language or "").strip().lower() in {"zh", "zh-cn", "zh_cn", "chinese", "中文", "cn"}


def _runtime_text(language, english, chinese):
    return chinese if _is_chinese(language) else english


def parse_nexus_tree(nexus_file_path, output_nwk_path=None, language="english"):
    """Parse a NEXUS tree file and optionally export a Newick copy."""

    # 读取 nexus 文件
    with open(nexus_file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    # 查找树的定义部分
    # MrBayes 输出通常在 begin trees; 和 end; 之间（不区分大小写）
    tree_block_pattern = r'begin\s+trees;(.*?)end;'
    tree_block_match = re.search(tree_block_pattern, content, re.DOTALL | re.IGNORECASE)

    if not tree_block_match:
        raise ValueError(_runtime_text(language, "TREES block not found", "未找到 TREES 块"))

    tree_block = tree_block_match.group(1)

    # 提取 TRANSLATE 信息
    translate_dict = {}
    translate_pattern = r'translate\s+(.*?);'
    translate_match = re.search(translate_pattern, tree_block, re.DOTALL | re.IGNORECASE)

    if translate_match:
        translate_content = translate_match.group(1)
        # 解析每一行的映射关系
        for line in translate_content.split(','):
            line = line.strip()
            if line:
                # 移除可能的制表符和多余空格
                parts = re.split(r'\s+', line.strip())
                if len(parts) >= 2:
                    translate_dict[parts[0]] = parts[1]

    # 查找树的定义行
    # 格式为: tree tree_name = [&U] (newick_string);
    tree_pattern = r'tree\s+\w+\s*=\s*(?:\[&[^\]]*\])?\s*(\([^;]+\));'
    tree_match = re.search(tree_pattern, tree_block, re.DOTALL | re.IGNORECASE)

    if not tree_match:
        raise ValueError(_runtime_text(language, "Tree definition not found", "未找到树的定义"))

    # 获取 newick 字符串
    newick_string = tree_match.group(1).strip()

    # 清理 newick 字符串，移除可能的注释
    # MrBayes 会在节点上添加支持值等信息，格式如 [&prob=0.95,length_mean=0.15]
    newick_string = re.sub(r'\[&[^\]]*\]', '', newick_string)

    # 替换数字标识符为实际物种名
    # 需要确保只替换作为节点标识符的数字，不替换分支长度中的数字
    for num, name in translate_dict.items():
        # 匹配模式：数字前面是 ( 或 , 或开头，后面是 [ 或 : 或 , 或 )
        pattern = r'(?<=[\(,])' + re.escape(num) + r'(?=[\[\:,\)])'
        newick_string = re.sub(pattern, name, newick_string)

    # 确保字符串以分号结尾
    if not newick_string.endswith(';'):
        newick_string += ';'

    # 如果没有指定输出路径，自动生成
    if output_nwk_path is None:
        output_nwk_path = nexus_file_path + '.nwk'

    # 直接导出清洗后的 Newick 文本
    with open(output_nwk_path, 'w', encoding='utf-8') as handle:
        handle.write(newick_string)
        handle.write('\n')

    leaf_count = len(translate_dict) if translate_dict else (newick_string.count(',') + 1 if newick_string else 0)

    print(_runtime_text(language, f"Successfully parsed nexus file: {nexus_file_path}", f"成功解析 nexus 文件: {nexus_file_path}"))
    print(_runtime_text(language, f"Tree exported in Newick format: {output_nwk_path}", f"树已导出为 newick 格式: {output_nwk_path}"))
    print(_runtime_text(language, f"Tree contains {leaf_count} leaf nodes", f"树包含 {leaf_count} 个叶节点"))

    return newick_string


if __name__ == "__main__":
    # experiments/trial/main.py
    import sys
    from pathlib import Path

    # 将 ref_fun 所在目录（即项目根目录下的 ref_fun）加入 Python 路径
    for i in range(0, 5):
        print("i:", i)
        yy = Path(__file__).resolve().parents[i]
        print(yy)

#!/usr/bin/env python3

from __future__ import annotations

from dataclasses import dataclass


ENGLISH = "english"
CHINESE = "chinese"


def normalize_language(value) -> str:
    text = str(value or "").strip().lower()
    if text in {"zh", "zh-cn", "zh_cn", "chinese", "中文", "cn"}:
        return CHINESE
    return ENGLISH


def runtime_language(runtime_config) -> str:
    if not isinstance(runtime_config, dict):
        return ENGLISH
    run_config = runtime_config.get("run")
    if not isinstance(run_config, dict):
        return ENGLISH
    return normalize_language(run_config.get("language"))


def translate_runtime_text(message: str, language: str) -> str:
    if not isinstance(message, str) or not message:
        return message

    replacements = ENGLISH_TO_CHINESE if normalize_language(language) == CHINESE else CHINESE_TO_ENGLISH
    translated = message
    for source, target in replacements:
        translated = translated.replace(source, target)
    return translated


@dataclass(frozen=True)
class RuntimeTranslator:
    language: str = ENGLISH

    @classmethod
    def from_runtime_config(cls, runtime_config):
        return cls(runtime_language(runtime_config))

    def is_chinese(self) -> bool:
        return normalize_language(self.language) == CHINESE

    def text(self, english: str, chinese: str) -> str:
        return chinese if self.is_chinese() else english

    def method_label(self, method_key: str) -> str:
        labels = {
            "distance": self.text("Distance method", "距离法"),
            "maximum_likelihood": self.text("Maximum likelihood method", "极大似然法"),
            "bayesian": self.text("Bayesian method", "贝叶斯法"),
            "parsimony": self.text("Parsimony method", "简约法"),
            "protein_structure": self.text("Protein structure", "蛋白质结构"),
        }
        return labels[method_key]

    def marker_complete(self, method_key: str) -> str:
        return self.text(
            f"===={self.method_label(method_key)} complete===========================================================",
            f"===={self.method_label(method_key)}完成===========================================================",
        )

    def marker_failed(self, method_key: str) -> str:
        return self.text(
            f"===={self.method_label(method_key)} failed===========================================================",
            f"===={self.method_label(method_key)}失败===========================================================",
        )

    def marker_skipped(self, method_key: str) -> str:
        return self.text(
            f"===={self.method_label(method_key)} skipped by runtime config====",
            f"===={self.method_label(method_key)}已按运行时配置跳过====",
        )


class LocalizedLogger:
    def __init__(self, logger, translator: RuntimeTranslator):
        self._logger = logger
        self._translator = translator

    def __getattr__(self, name):
        return getattr(self._logger, name)

    def debug(self, message, *args, **kwargs):
        self._logger.debug(self._translate(message), *args, **kwargs)

    def info(self, message, *args, **kwargs):
        self._logger.info(self._translate(message), *args, **kwargs)

    def warning(self, message, *args, **kwargs):
        self._logger.warning(self._translate(message), *args, **kwargs)

    def error(self, message, *args, **kwargs):
        self._logger.error(self._translate(message), *args, **kwargs)

    def _translate(self, message):
        if isinstance(message, str):
            return translate_runtime_text(message, self._translator.language)
        return message


ENGLISH_TO_CHINESE = [
    ("Changed to work_dir: ", "当前所在目录： "),
    ("Starting distance-based inference...", "开始距离法建树..."),
    ("Distance-based inference completed", "距离法建树完成"),
    ("Distance-based method failed: ", "距离法建树失败: "),
    ("Starting maximum likelihood inference (IQ-TREE)...", "开始极大似然法建树..."),
    ("Maximum likelihood inference completed", "极大似然法建树完成"),
    ("Maximum likelihood method failed: ", "极大似然法建树失败: "),
    ("Starting Bayesian inference (MrBayes)...", "开始贝叶斯法建树..."),
    ("Bayesian inference completed", "贝叶斯法建树完成"),
    ("MrBayes consensus tree output not found", "未找到贝叶斯树输出文件"),
    ("MrBayes path not explicitly set; will try to use system 'mb' executable", "MrBayes路径未显式指定，将尝试使用系统中的 `mb`"),
    ("Bayesian method failed: ", "贝叶斯法建树失败: "),
    ("Starting parsimony inference...", "开始简约法建树..."),
    ("Parsimony inference completed", "简约法建树完成"),
    ("Parsimony method failed: ", "简约法建树失败: "),
    ("Generating tree visualizations...", "生成进化树可视化图片..."),
    ("Visualization completed", "可视化完成"),
    ("Tree distance calculation completed", "完成计算树的距离"),
    ("Tree distance calculation failed with command: ", "树距离计算失败，命令: "),
    ("Tree distance script execution failed with command ", "树距离脚本执行失败，命令 "),
    ("Tree distance heatmaps saved: ", "树距离热图已保存: "),
    (" and ", " 和 "),
    ("Failed to generate tree distance heatmaps: ", "生成树距离热图失败: "),
    ("Summary saved: ", "结果总结已保存: "),
    ("MAD executable not found: ", "MAD 工具未找到: "),
    ("MAD rerooting failed: ", "MAD 定根失败: "),
    ("Rerooting completed: ", "定根完成，"),
    ("Restored names in tree: ", "已恢复树文件名称: "),
    ("Failed to restore names in ", "恢复树文件名称失败 "),
    ("Restoring original sequence names in trees...", "恢复树文件中的原始序列名称..."),
    ("Protein phylogenetic pipeline completed!", "进化树构建管道完成!"),
    ("Phylogenetic pipeline completed!", "进化树构建管道完成!"),
    ("Results saved to: ", "结果保存在: "),
    ("Starting protein phylogenetic pipeline...This is: ", "开始进化树构建管道...This is: "),
    ("Starting phylogenetic pipeline...This is: ", "开始进化树构建管道...This is: "),
    ("Checking required software...", "检查所需软件..."),
    ("Input file does not exist: ", "输入文件不存在: "),
    ("At least 3 sequences are required to build a phylogenetic tree", "至少需要3条序列才能构建进化树"),
    ("Input file validation succeeded: ", "输入文件验证成功: "),
    (" sequences", "条序列"),
    ("Input file format error: ", "输入文件格式错误: "),
    ("Creating output directory: ", "创建输出目录: "),
    ("Converting to PHYLIP format: ", "转换为PHYLIP格式: "),
    ("Setting temporary naming convention for input sequences...", "为输入序列设置临时命名规则..."),
    ("IQ-TREE not found", "IQ-TREE未找到"),
    ("MrBayes not found", "MrBayes未找到"),
    ("Required PHYLIP commands not found: ", "未找到必需的 PHYLIP 命令: "),
    ("Tree file not found, skipping: ", "树文件不存在，跳过: "),
    ("No name mapping available; skipping name restoration in trees", "没有名称映射，跳过树文件中的名称恢复"),
    ("Standard error:", "标准错误:"),
]


CHINESE_TO_ENGLISH = [(target, source) for source, target in ENGLISH_TO_CHINESE]

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import copy
import json
from pathlib import Path


PROTEIN_DEFAULTS = {
    "distance": {
        "enabled": True,
        "neighbor_method": "NJ",
        "neighbor_outgroup_index": None,
        "protdist_menu_overrides": [],
        "neighbor_menu_overrides": [],
    },
    "maximum_likelihood": {
        "enabled": True,
        "bootstrap_replicates": 1000,
        "model_strategy": "MFP",
        "model_set": "WAG,LG,JTT,Dayhoff,DCMut,rtREV,cpREV,VT,Blosum62,mtMam,mtArt,HIVb,HIVw",
        "threads": None,
        "threads_max": None,
        "seed": None,
        "safe": False,
        "keep_ident": False,
        "quiet": True,
        "verbose": False,
        "redo": True,
        "memory_limit": None,
        "outgroup": None,
        "sequence_type": None,
        "alrt": None,
        "abayes": False,
        "extra_args": [],
    },
    "bayesian": {
        "enabled": True,
        "protein_model_prior": "mixed",
        "rates": "invgamma",
        "ngen": 100000,
        "samplefreq": 100,
        "printfreq": 1000,
        "diagnfreq": 5000,
        "nruns": None,
        "nchains": None,
        "temp": None,
        "stoprule": None,
        "stopval": None,
        "burnin": None,
        "burninfrac": None,
        "relburnin": None,
        "command_block": [],
    },
    "parsimony": {
        "enabled": True,
        "protpars_outgroup_index": None,
        "protpars_print_steps": True,
        "protpars_print_sequences": True,
        "protpars_menu_overrides": [],
    },
    "protein_structure": {
        "enabled": False,
        "backend": "foldseek",
        "use_structure_manifest": False,
        "structure_manifest_file": None,
        "prostt5_model_path": None,
        "sequence_only_mode": "prostt5",
        "similarity_rule": "mean_qtmscore_ttmscore",
        "missing_distance": "13.81551056",
        "tree_builder_method": "NJ",
        "threads": 0,
        "sensitivity": 9.5,
        "evalue": 10.0,
        "max_seqs": 1000,
        "coverage_threshold": 0.0,
        "coverage_mode": 0,
        "alignment_type": 2,
        "tmscore_threshold": 0.0,
        "exhaustive_search": False,
        "exact_tmscore": False,
        "gpu": False,
        "verbosity": 3,
        "extra_args": [],
    },
    "reroot": {
        "method": "MAD",
        "ladderization": {
            "direction": "UP",
            "sort_by_clade_size": True,
            "sort_by_leaf_name_string": True,
            "sort_by_branch_length": True,
        },
    },
}


DNA_DEFAULTS = {
    "distance": {
        "enabled": True,
        "dnadist_model": "F84",
        "dnadist_transition_transversion_ratio": 2.0,
        "dnadist_empirical_base_frequencies": True,
        "neighbor_method": "NJ",
        "neighbor_outgroup_index": None,
        "dnadist_menu_overrides": [],
        "neighbor_menu_overrides": [],
    },
    "maximum_likelihood": {
        "enabled": True,
        "bootstrap_replicates": 1000,
        "model_strategy": "MFP",
        "model_set": "",
        "threads": None,
        "threads_max": None,
        "seed": None,
        "safe": False,
        "keep_ident": False,
        "quiet": True,
        "verbose": False,
        "redo": True,
        "memory_limit": None,
        "outgroup": None,
        "sequence_type": None,
        "alrt": None,
        "abayes": False,
        "extra_args": [],
    },
    "bayesian": {
        "enabled": True,
        "rates": "invgamma",
        "ngen": 100000,
        "samplefreq": 100,
        "printfreq": 1000,
        "diagnfreq": 5000,
        "nst": 6,
        "nruns": None,
        "nchains": None,
        "temp": None,
        "stoprule": None,
        "stopval": None,
        "burnin": None,
        "burninfrac": None,
        "relburnin": None,
        "command_block": [],
    },
    "parsimony": {
        "enabled": True,
        "dnapars_outgroup_index": None,
        "dnapars_transversion_parsimony": False,
        "dnapars_menu_overrides": [],
    },
    "protein_structure": {
        "enabled": False,
        "backend": "foldseek",
        "use_structure_manifest": False,
        "structure_manifest_file": None,
        "prostt5_model_path": None,
        "sequence_only_mode": "prostt5",
        "similarity_rule": "mean_qtmscore_ttmscore",
        "missing_distance": "13.81551056",
        "tree_builder_method": "NJ",
        "threads": 0,
        "sensitivity": 9.5,
        "evalue": 10.0,
        "max_seqs": 1000,
        "coverage_threshold": 0.0,
        "coverage_mode": 0,
        "alignment_type": 2,
        "tmscore_threshold": 0.0,
        "exhaustive_search": False,
        "exact_tmscore": False,
        "gpu": False,
        "verbosity": 3,
        "extra_args": [],
    },
    "reroot": {
        "method": "MAD",
        "ladderization": {
            "direction": "UP",
            "sort_by_clade_size": True,
            "sort_by_leaf_name_string": True,
            "sort_by_branch_length": True,
        },
    },
}


PROTEIN_FIXED_MODELS = {"LG", "WAG", "JTT"}
DNA_FIXED_MODELS = {"GTR", "HKY", "JC"}


def load_runtime_config(config_path):
    if config_path is None:
        return {}

    path = Path(config_path)
    if not path.exists():
        raise FileNotFoundError(f"Runtime config does not exist: {path}")

    content = path.read_text(encoding="utf-8").strip()
    if not content:
        return {}

    payload = json.loads(content)
    if not isinstance(payload, dict):
        raise ValueError("Runtime config must be a JSON object.")
    _resolve_runtime_config_paths(payload, path.parent)
    return payload


def _resolve_runtime_config_paths(payload, config_dir):
    methods = payload.get("methods")
    if not isinstance(methods, dict):
        return
    protein_structure = methods.get("protein_structure")
    if not isinstance(protein_structure, dict):
        return
    for path_key in ("structure_manifest_file", "prostt5_model_path"):
        raw_path = protein_structure.get(path_key)
        if not raw_path:
            continue
        resolved_path = Path(str(raw_path).strip()).expanduser()
        if not resolved_path.is_absolute():
            resolved_path = (config_dir / resolved_path).resolve()
        protein_structure[path_key] = str(resolved_path)


def protein_runtime_settings(runtime_config):
    return _merged_settings(PROTEIN_DEFAULTS, runtime_config, "PROTEIN")


def dna_runtime_settings(runtime_config):
    return _merged_settings(DNA_DEFAULTS, runtime_config, "DNA_CDS")


def _merged_settings(defaults, runtime_config, input_type):
    merged = copy.deepcopy(defaults)
    methods = (runtime_config or {}).get("methods", {})
    reroot = (runtime_config or {}).get("reroot", {})
    if isinstance(reroot, dict):
        _merge_reroot(merged["reroot"], reroot)

    if isinstance(methods, dict):
        for method_name, overrides in methods.items():
            if method_name not in merged or not isinstance(overrides, dict):
                continue
            if method_name == "maximum_likelihood":
                _merge_maximum_likelihood(merged[method_name], overrides)
            elif method_name == "bayesian":
                _merge_bayesian(merged[method_name], overrides)
            elif method_name == "distance":
                _merge_distance(merged[method_name], overrides)
            elif method_name == "parsimony":
                _merge_parsimony(merged[method_name], overrides)
            elif method_name == "protein_structure":
                _merge_protein_structure(merged[method_name], overrides, input_type)
            else:
                _merge_scalar_overrides(merged[method_name], overrides)

    _normalize_input_type_specific_settings(merged, input_type)

    return merged


def _merge_reroot(target, overrides):
    method = str(overrides.get("method") or target.get("method") or "MAD").strip()
    if method not in {"MAD", "root-at-middle-point"}:
        method = "MAD"
    target["method"] = method
    ladderization = target.setdefault(
        "ladderization",
        {
            "direction": "UP",
            "sort_by_clade_size": True,
            "sort_by_leaf_name_string": True,
            "sort_by_branch_length": True,
        },
    )
    override_ladderization = overrides.get("ladderization")
    if isinstance(override_ladderization, dict):
        direction = str(override_ladderization.get("direction") or ladderization.get("direction") or "UP").strip().upper()
        if direction not in {"UP", "DOWN"}:
            direction = "UP"
        ladderization["direction"] = direction
        if "sort_by_clade_size" in override_ladderization:
            ladderization["sort_by_clade_size"] = bool(override_ladderization.get("sort_by_clade_size"))
        if "sort_by_leaf_name_string" in override_ladderization:
            ladderization["sort_by_leaf_name_string"] = bool(override_ladderization.get("sort_by_leaf_name_string"))
        if "sort_by_branch_length" in override_ladderization:
            ladderization["sort_by_branch_length"] = bool(override_ladderization.get("sort_by_branch_length"))


def _normalize_input_type_specific_settings(merged, input_type):
    maximum_likelihood = merged.get("maximum_likelihood")
    if not isinstance(maximum_likelihood, dict):
        return

    model_strategy = _normalize_tree_build_model_strategy(
        maximum_likelihood.get("model_strategy")
    )
    maximum_likelihood["model_strategy"] = model_strategy
    raw_model_set = str(maximum_likelihood.get("model_set") or "").strip()

    if input_type == "DNA_CDS":
        if model_strategy in DNA_FIXED_MODELS:
            maximum_likelihood["model_set"] = ""
            return
        maximum_likelihood["model_set"] = _filter_model_set(raw_model_set, DNA_FIXED_MODELS)
        return

    if model_strategy in PROTEIN_FIXED_MODELS:
        maximum_likelihood["model_set"] = ""
        return
    maximum_likelihood["model_set"] = raw_model_set


def _merge_maximum_likelihood(target, overrides):
    _merge_scalar_overrides(target, overrides, skip_keys={"iqtree"})
    iqtree = overrides.get("iqtree")
    if not isinstance(iqtree, dict):
        target["model_strategy"] = _normalize_tree_build_model_strategy(
            target.get("model_strategy")
        )
        return

    _merge_scalar_overrides(target, iqtree, skip_keys={"common", "advanced", "extra_args"})
    _merge_section_into_target(target, iqtree.get("common"))
    _merge_section_into_target(target, iqtree.get("advanced"))
    if "extra_args" in iqtree:
        target["extra_args"] = _normalize_string_list(iqtree.get("extra_args"))
    target["model_strategy"] = _normalize_tree_build_model_strategy(
        target.get("model_strategy")
    )


def _merge_bayesian(target, overrides):
    _merge_scalar_overrides(target, overrides, skip_keys={"mrbayes"})
    mrbayes = overrides.get("mrbayes")
    if not isinstance(mrbayes, dict):
        return

    _merge_scalar_overrides(target, mrbayes, skip_keys={"common", "advanced", "command_block"})
    _merge_section_into_target(target, mrbayes.get("common"))
    _merge_section_into_target(target, mrbayes.get("advanced"))
    if "command_block" in mrbayes:
        target["command_block"] = _normalize_string_list(mrbayes.get("command_block"))


def _merge_distance(target, overrides):
    _merge_scalar_overrides(target, overrides, skip_keys={"protdist", "dnadist", "neighbor"})
    _merge_phylip_program(target, overrides, "protdist", "protdist_menu_overrides")
    _merge_phylip_program(target, overrides, "dnadist", "dnadist_menu_overrides")
    _merge_phylip_program(target, overrides, "neighbor", "neighbor_menu_overrides")


def _merge_parsimony(target, overrides):
    _merge_scalar_overrides(target, overrides, skip_keys={"protpars", "dnapars"})
    _merge_phylip_program(target, overrides, "protpars", "protpars_menu_overrides")
    _merge_phylip_program(target, overrides, "dnapars", "dnapars_menu_overrides")


def _merge_protein_structure(target, overrides, input_type):
    _merge_scalar_overrides(target, overrides, skip_keys={"foldseek"})
    foldseek = overrides.get("foldseek")
    if isinstance(foldseek, dict):
        _merge_section_into_target(target, foldseek.get("common"))
        _merge_section_into_target(target, foldseek.get("advanced"))
        if "extra_args" in foldseek:
            target["extra_args"] = _normalize_string_list(foldseek.get("extra_args"))
    target["backend"] = "foldseek"
    if input_type != "PROTEIN":
        target["enabled"] = False
    target["use_structure_manifest"] = bool(target.get("use_structure_manifest", False))
    manifest_file = target.get("structure_manifest_file")
    target["structure_manifest_file"] = str(manifest_file).strip() if manifest_file else None
    prostt5_model_path = target.get("prostt5_model_path")
    target["prostt5_model_path"] = str(prostt5_model_path).strip() if prostt5_model_path else None
    sequence_only_mode = str(target.get("sequence_only_mode") or "prostt5").strip().lower()
    target["sequence_only_mode"] = sequence_only_mode or "prostt5"
    similarity_rule = str(target.get("similarity_rule") or "mean_qtmscore_ttmscore").strip()
    target["similarity_rule"] = similarity_rule or "mean_qtmscore_ttmscore"
    missing_distance = target.get("missing_distance")
    target["missing_distance"] = str(missing_distance).strip() if missing_distance is not None else "13.81551056"
    target["tree_builder_method"] = _normalize_protein_structure_tree_builder_method(
        target.get("tree_builder_method")
    )
    target["threads"] = _normalize_int(target.get("threads"), 0, 0, 512)
    target["sensitivity"] = _normalize_float(target.get("sensitivity"), 9.5, 1.0, 20.0)
    target["evalue"] = _normalize_float(target.get("evalue"), 10.0, 0.0, None)
    target["max_seqs"] = _normalize_int(target.get("max_seqs"), 1000, 1, None)
    target["coverage_threshold"] = _normalize_float(target.get("coverage_threshold"), 0.0, 0.0, 1.0)
    target["coverage_mode"] = _normalize_int(target.get("coverage_mode"), 0, 0, 5)
    target["alignment_type"] = _normalize_int(target.get("alignment_type"), 2, 0, 2)
    target["tmscore_threshold"] = _normalize_float(target.get("tmscore_threshold"), 0.0, 0.0, 1.0)
    target["exhaustive_search"] = _normalize_bool(target.get("exhaustive_search"), False)
    target["exact_tmscore"] = _normalize_bool(target.get("exact_tmscore"), False)
    target["gpu"] = _normalize_bool(target.get("gpu"), False)
    target["verbosity"] = _normalize_int(target.get("verbosity"), 3, 0, 3)
    target["extra_args"] = _normalize_string_list(target.get("extra_args"))


def _normalize_protein_structure_tree_builder_method(value):
    method = str(value or "NJ").strip()
    if method.lower() in {"swiftnj", "swift nj"}:
        return "SwiftNJ"
    return "NJ"


def _normalize_int(value, default, minimum=None, maximum=None):
    try:
        normalized = int(value)
    except (TypeError, ValueError):
        normalized = default
    if minimum is not None:
        normalized = max(minimum, normalized)
    if maximum is not None:
        normalized = min(maximum, normalized)
    return normalized


def _normalize_float(value, default, minimum=None, maximum=None):
    try:
        normalized = float(value)
    except (TypeError, ValueError):
        normalized = default
    if minimum is not None:
        normalized = max(minimum, normalized)
    if maximum is not None:
        normalized = min(maximum, normalized)
    return normalized


def _normalize_bool(value, default=False):
    if isinstance(value, bool):
        return value
    if value is None:
        return default
    text = str(value).strip().lower()
    if text in {"1", "true", "yes", "y", "on"}:
        return True
    if text in {"0", "false", "no", "n", "off"}:
        return False
    return default


def _merge_phylip_program(target, overrides, program_key, target_key):
    program = overrides.get(program_key)
    if not isinstance(program, dict):
        return

    _merge_section_into_target(target, program.get("common"))
    _merge_section_into_target(target, program.get("advanced"))
    if "menu_overrides" in program:
        target[target_key] = _normalize_string_list(program.get("menu_overrides"))


def _merge_section_into_target(target, section):
    if not isinstance(section, dict):
        return
    for key, value in section.items():
        if value is not None:
            target[key] = value


def _merge_scalar_overrides(target, overrides, skip_keys=None):
    skipped = set(skip_keys or ())
    for key, value in overrides.items():
        if key in skipped or isinstance(value, dict):
            continue
        if key.endswith("_args") or key.endswith("_block") or key.endswith("_overrides"):
            target[key] = _normalize_string_list(value)
        elif value is not None:
            target[key] = value


def _normalize_string_list(value):
    if value is None:
        return []
    if isinstance(value, str):
        return [value]
    if isinstance(value, (list, tuple)):
        return [str(item) for item in value if item is not None]
    return [str(value)]


def _normalize_tree_build_model_strategy(value):
    strategy = str(value).strip() if value is not None else ""
    if not strategy:
        return "MFP"
    if strategy in {"MF", "TESTNEWONLY", "TESTNEW"}:
        return "MFP"
    if strategy == "TESTONLY":
        return "TEST"
    return strategy


def _filter_model_set(raw_model_set, allowed_models):
    if not raw_model_set:
        return ""
    tokens = []
    seen = set()
    for item in raw_model_set.split(","):
        normalized = str(item).strip().upper()
        if not normalized or normalized not in allowed_models or normalized in seen:
            continue
        seen.add(normalized)
        tokens.append(normalized)
    return ",".join(tokens)

#!/usr/bin/env python3
"""Render combined heatmaps for tree distance summaries."""

from __future__ import annotations

import csv
import math
from pathlib import Path
from typing import List, Sequence, Tuple

import matplotlib as mpl

mpl.set_loglevel("warning")
mpl.use("Agg", force=True)

import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap, Normalize
from matplotlib.patches import Rectangle


def _read_distance_matrix(matrix_path: Path) -> Tuple[List[str], List[List[float]]]:
    with open(matrix_path, "r", encoding="utf-8", newline="") as handle:
        rows = [row for row in csv.reader(handle, delimiter="\t") if row]

    if len(rows) < 2:
        raise ValueError(f"Matrix file is empty or incomplete: {matrix_path}")

    header = rows[0]
    labels = header[1:] if header and header[0].strip() == "" else header
    matrix: List[List[float]] = []
    row_labels: List[str] = []

    for row in rows[1:]:
        if len(row) == len(labels) + 1:
            row_labels.append(row[0])
            values = row[1:]
        elif len(row) == len(labels):
            values = row
        else:
            raise ValueError(f"Unexpected row width in matrix file: {matrix_path}")

        matrix.append([float(value) for value in values])

    if row_labels and row_labels != labels:
        raise ValueError(f"Row and column labels do not match in {matrix_path}")

    if any(len(row) != len(labels) for row in matrix):
        raise ValueError(f"Matrix shape is inconsistent in {matrix_path}")

    return labels, matrix


def _format_value(value: float) -> str:
    if abs(value - round(value)) < 1e-9:
        return str(int(round(value)))
    return f"{value:.3f}".rstrip("0").rstrip(".")


def _build_cmap(hex_colors: Sequence[str]) -> LinearSegmentedColormap:
    return LinearSegmentedColormap.from_list("tree_summary", list(hex_colors))


def _draw_single_heatmap(
    ax, labels: Sequence[str], matrix: Sequence[Sequence[float]], title: str, cmap
):
    size = len(labels)
    display_matrix: List[List[float]] = []
    off_diag_values: List[float] = []

    for row_index, row in enumerate(matrix):
        display_row: List[float] = []
        for col_index, value in enumerate(row):
            if row_index == col_index:
                display_row.append(math.nan)
            else:
                display_row.append(value)
                off_diag_values.append(value)
        display_matrix.append(display_row)

    if off_diag_values:
        vmin = min(off_diag_values)
        vmax = max(off_diag_values)
        if abs(vmax - vmin) < 1e-12:
            vmax = vmin + 1.0
    else:
        vmin, vmax = 0.0, 1.0

    cmap = cmap.copy()
    cmap.set_bad("#f3f4f6")

    image = ax.imshow(
        display_matrix,
        cmap=cmap,
        norm=Normalize(vmin=vmin, vmax=vmax),
        interpolation="nearest",
        aspect="equal",
    )

    ax.set_title(title, fontsize=15, fontweight="semibold", pad=12)
    ax.set_xticks(
        range(size), labels=labels, rotation=35, ha="right", rotation_mode="anchor"
    )
    ax.set_yticks(range(size), labels=labels)
    ax.tick_params(axis="both", which="major", labelsize=10)

    ax.set_xticks([index - 0.5 for index in range(1, size)], minor=True)
    ax.set_yticks([index - 0.5 for index in range(1, size)], minor=True)
    ax.grid(which="minor", color="#ffffff", linestyle="-", linewidth=1.4)
    ax.tick_params(which="minor", bottom=False, left=False)

    for spine in ax.spines.values():
        spine.set_visible(False)

    for row_index, row in enumerate(matrix):
        for col_index, value in enumerate(row):
            if row_index == col_index:
                ax.add_patch(
                    Rectangle(
                        (col_index - 0.5, row_index - 0.5),
                        1,
                        1,
                        facecolor="#f3f4f6",
                        edgecolor="#ffffff",
                        linewidth=1.4,
                    )
                )
                text_color = "#6b7280"
            else:
                normalized = (value - vmin) / (vmax - vmin) if vmax > vmin else 0.0
                text_color = "white" if normalized > 0.62 else "#111827"

            ax.text(
                col_index,
                row_index,
                _format_value(value),
                ha="center",
                va="center",
                fontsize=10,
                color=text_color,
                fontweight="semibold" if row_index != col_index else "normal",
            )

    return image


def create_tree_distance_heatmaps(
    tree_distance_matrix: Path | str,
    rf_distance_matrix: Path | str,
    output_png: Path | str,
    output_pdf: Path | str,
    figure_title: str = "Tree Distance Comparison Across Inference Methods",
    tree_title: str = "TreeDist Matrix",
    rf_title: str = "Robinson-Foulds Matrix",
    tree_colorbar_label: str = "TreeDist",
    rf_colorbar_label: str = "RF Distance",
) -> Tuple[Path, Path]:
    tree_distance_matrix = Path(tree_distance_matrix)
    rf_distance_matrix = Path(rf_distance_matrix)
    output_png = Path(output_png)
    output_pdf = Path(output_pdf)

    labels, tree_distance = _read_distance_matrix(tree_distance_matrix)
    rf_labels, rf_distance = _read_distance_matrix(rf_distance_matrix)

    if labels != rf_labels:
        raise ValueError("TreeDist and RF matrices must use the same method order")

    plt.rcParams["font.sans-serif"] = ["DejaVu Sans", "Arial"]
    plt.rcParams["axes.unicode_minus"] = False

    fig, axes = plt.subplots(1, 2, figsize=(15, 7.8), constrained_layout=True)
    fig.patch.set_facecolor("#fcfcfd")
    fig.suptitle(
        figure_title,
        fontsize=18,
        fontweight="bold",
        y=1.02,
    )

    tree_cmap = _build_cmap(["#e0f2fe", "#38bdf8", "#0f766e"])
    rf_cmap = _build_cmap(["#ede9fe", "#8b5cf6", "#4c1d95"])

    tree_image = _draw_single_heatmap(
        axes[0], labels, tree_distance, tree_title, tree_cmap
    )
    rf_image = _draw_single_heatmap(
        axes[1], labels, rf_distance, rf_title, rf_cmap
    )

    tree_colorbar = fig.colorbar(tree_image, ax=axes[0], shrink=0.82, pad=0.02)
    tree_colorbar.set_label(tree_colorbar_label, fontsize=10)
    rf_colorbar = fig.colorbar(rf_image, ax=axes[1], shrink=0.82, pad=0.02)
    rf_colorbar.set_label(rf_colorbar_label, fontsize=10)

    output_png.parent.mkdir(parents=True, exist_ok=True)
    output_pdf.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_png, dpi=300, bbox_inches="tight", facecolor=fig.get_facecolor())
    fig.savefig(output_pdf, bbox_inches="tight", facecolor=fig.get_facecolor())
    plt.close(fig)

    return output_png, output_pdf


def main(summary_dir: Path | str) -> Tuple[Path, Path]:
    summary_dir = Path(summary_dir)
    return create_tree_distance_heatmaps(
        summary_dir / "tree_distance_matrix.tsv",
        summary_dir / "rf_distance_matrix.tsv",
        summary_dir / "tree_distance_heatmaps.png",
        summary_dir / "tree_distance_heatmaps.pdf",
    )


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Create combined heatmaps for tree distance summary matrices."
    )
    parser.add_argument(
        "summary_dir",
        nargs="?",
        default=".",
        help="Directory containing tree_distance_matrix.tsv and rf_distance_matrix.tsv",
    )
    arguments = parser.parse_args()
    png_path, pdf_path = main(arguments.summary_dir)
    print(f"Saved heatmaps: {png_path} and {pdf_path}")