#!/usr/bin/env python3
"""
KaKs Analysis Utilities
Ka/Ks分析工具函数
"""

import hashlib
import logging
import re
import shutil
import subprocess
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from Bio import SeqIO

logger = logging.getLogger(__name__)


def compute_checksum(filepath: Path) -> str:
    """计算文件SHA256校验值"""
    sha256 = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            sha256.update(chunk)
    return sha256.hexdigest()


def check_tool(tool_name: str) -> Optional[str]:
    """检查工具是否可用"""
    path = shutil.which(tool_name)
    if path:
        return path
    return None


def get_tool_version(tool_name: str) -> str:
    """获取工具版本信息"""
    try:
        if tool_name == "muscle":
            result = subprocess.run(
                ["muscle", "-version"],
                capture_output=True, text=True, timeout=5
            )
            return result.stdout.strip().split('\n')[0] if result.stdout else "unknown"
        elif tool_name == "yn00" or tool_name == "codeml":
            result = subprocess.run(
                [tool_name],
                capture_output=True, text=True, timeout=5
            )
            output = result.stderr.strip() or result.stdout.strip()
            lines = output.split('\n')
            for line in lines:
                if "version" in line.lower() or "paml" in line.lower():
                    return line.strip()
            return "PAML (version unknown)"
        elif tool_name == "pal2nal":
            return "pal2nal v14.1"
    except Exception as e:
        logger.warning(f"Failed to get {tool_name} version: {e}")
    return "not_found"


def parse_cds_fasta(fasta_path: Path) -> Dict[str, str]:
    """解析CDS FASTA文件，返回基因ID到序列的字典"""
    sequences = {}
    for record in SeqIO.parse(fasta_path, "fasta"):
        gene_id = record.id
        sequences[gene_id] = str(record.seq)
    return sequences


def parse_ortholog_pairs(pairs_file: Path) -> List[Tuple[str, str]]:
    """解析同源基因对文件，格式: gene1\tgene2"""
    pairs = []
    with open(pairs_file, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            parts = line.split('\t')
            if len(parts) >= 2:
                gene1 = parts[0].strip()
                gene2 = parts[1].strip()
                pairs.append((gene1, gene2))
    return pairs


def extract_sequences(cds_dict: Dict[str, str], gene_ids: List[str], output_path: Path) -> int:
    """从CDS字典中提取指定基因的序列并写入文件"""
    count = 0
    with open(output_path, 'w') as f:
        for gene_id in gene_ids:
            if gene_id in cds_dict:
                seq = cds_dict[gene_id]
                f.write(f">{gene_id}\n")
                for i in range(0, len(seq), 60):
                    f.write(seq[i:i+60] + '\n')
                count += 1
            else:
                logger.warning(f"Gene {gene_id} not found in CDS file")
    return count


def run_muscle_alignment(input_fasta: Path, output_aln: Path, maxiters: int = 2, cpu: int = 8) -> bool:
    """运行MUSCLE进行蛋白序列比对"""
    muscle_path = check_tool("muscle")
    if not muscle_path:
        logger.error("muscle not found in PATH")
        return False
    
    try:
        cmd = [
            "muscle", "-align", str(input_fasta), "-output", str(output_aln),
            "-maxiters", str(maxiters), "-threads", str(cpu)
        ]
        logger.info(f"Running MUSCLE: {' '.join(cmd)}")
        subprocess.run(cmd, check=True, capture_output=True, text=True)
        return output_aln.exists() and output_aln.stat().st_size > 0
    except Exception as e:
        logger.error(f"MUSCLE error: {e}")
        return False


def run_pal2nal(protein_aln: Path, cds_fasta: Path, output_codon: Path, output_format: str = "fasta") -> bool:
    """运行pal2nal将蛋白比对转换为密码子比对"""
    pal2nal_path = check_tool("pal2nal.pl") or check_tool("pal2nal")
    if not pal2nal_path:
        logger.error("pal2nal.pl not found in PATH")
        return False
    
    try:
        cmd = [
            "perl", str(pal2nal_path), str(protein_aln), str(cds_fasta),
            "-output", output_format, "-nogap", "-nomismatch"
        ]
        logger.info(f"Running pal2nal: {' '.join(cmd)}")
        with open(output_codon, 'w') as out_f:
            subprocess.run(cmd, check=True, stdout=out_f, stderr=subprocess.PIPE, text=True)
        return output_codon.exists() and output_codon.stat().st_size > 0
    except Exception as e:
        logger.error(f"pal2nal error: {e}")
        return False


def create_yn00_ctl(seqfile: Path, outfile: Path, genetic_code: int = 0) -> Path:
    """创建yn00的控制文件"""
    ctl_path = seqfile.parent / "yn00.ctl"
    ctl_content = f"""seqfile = {seqfile}
outfile = {outfile}
noisy = 9
verbose = 1
icode = {genetic_code}
weighting = 0
commonf3x4 = 0
"""
    with open(ctl_path, 'w') as f:
        f.write(ctl_content)
    return ctl_path


def run_yn00(codon_aln: Path, output_dir: Path, genetic_code: int = 0) -> Optional[Path]:
    """运行yn00计算Ka/Ks"""
    yn00_path = check_tool("yn00")
    if not yn00_path:
        logger.error("yn00 not found in PATH")
        return None
    
    try:
        output_file = output_dir / "yn00.out"
        ctl_file = create_yn00_ctl(codon_aln, output_file, genetic_code)
        cmd = [yn00_path, str(ctl_file)]
        logger.info(f"Running yn00: {' '.join(cmd)}")
        subprocess.run(cmd, cwd=output_dir, check=True, capture_output=True, text=True)
        return output_file if output_file.exists() and output_file.stat().st_size > 0 else None
    except Exception as e:
        logger.error(f"yn00 error: {e}")
        return None


def parse_yn00_output(yn00_file: Path) -> List[Dict]:
    """解析yn00输出文件，提取Ka/Ks结果"""
    results = []
    current_pair = {}
    
    with open(yn00_file, 'r') as f:
        content = f.read()
    
    lines = content.split('\n')
    i = 0
    while i < len(lines):
        line = lines[i].strip()
        if line.startswith('(') and 'vs.' in line:
            match = re.search(r'\(([^)]+)\s+vs\.\s+([^)]+)\)', line)
            if match:
                gene1 = match.group(1).strip()
                gene2 = match.group(2).strip()
                current_pair = {'gene1': gene1, 'gene2': gene2}
                
                j = i + 1
                while j < len(lines) and j < i + 20:
                    next_line = lines[j].strip()
                    if 'Ka' in next_line and 'Ks' in next_line:
                        if j + 1 < len(lines):
                            values_line = lines[j + 1].strip()
                            parts = values_line.split()
                            if len(parts) >= 4:
                                try:
                                    current_pair['Ka'] = float(parts[0])
                                    current_pair['Ks'] = float(parts[1])
                                    current_pair['Ka_Ks'] = float(parts[2]) if parts[2] != 'nan' else float('nan')
                                    current_pair['method'] = parts[3] if len(parts) > 3 else 'YN00'
                                    results.append(current_pair.copy())
                                except ValueError:
                                    pass
                        break
                    j += 1
        i += 1
    
    return results


def calculate_divergence_time(ks_values: List[float], neutral_rate: float = 7.08e-9) -> Dict:
    """计算分歧时间"""
    valid_ks = [ks for ks in ks_values if ks > 0 and not (ks != ks)]
    
    if not valid_ks:
        return {'mean_ks': None, 'median_ks': None, 'divergence_time_my': None}
    
    mean_ks = sum(valid_ks) / len(valid_ks)
    median_ks = sorted(valid_ks)[len(valid_ks) // 2]
    
    # T = Ks / (2 * neutral_rate)
    divergence_time = mean_ks / (2 * neutral_rate) / 1e6  # 转换为百万年
    
    return {
        'mean_ks': mean_ks,
        'median_ks': median_ks,
        'divergence_time_my': divergence_time,
        'pair_count': len(valid_ks)
    }