from pathlib import Path
from dataclasses import dataclass

import requests


@dataclass
class MSAResult:
    """複数の配列アライメントの結果。"""
    query_id: str
    a3m_content: str  # a3m形式のMSA
    depth: int  # MSAの深さ（配列の数）


def build_msa_remote(
    query_sequence: str,
    query_id: str = "target",
    api_base: str = "https://api.colabfold.com",
    max_depth: int = 5000,
) -> MSAResult:
    """ColabFold MSAサーバー（MMseqs2バックエンド）を使用してリモートMSAを構築します。

    Boltzの公式ドキュメントでは、ColabFoldまたはBioLM APIが推奨されています[4]。
    """
    # 実際のAPI呼び出し（例：ColabFold MMseqs2 API）
    payload = {"query": f">{query_id}\n{query_sequence}", "mode": "env"}
    resp = requests.post(f"{api_base}/ticket/msa", data=payload, timeout=600)
    resp.raise_for_status()
    ticket = resp.json()["id"]

    # 完了をポーリングし、待機します
    import time
    while True:
        status = requests.get(f"{api_base}/ticket/{ticket}", timeout=30).json()
        if status["status"] == "COMPLETE":
            break
        elif status["status"] == "ERROR":
            raise RuntimeError(f"MSAの構築に失敗しました：{status}")
        time.sleep(5)

    # a3m結果をダウンロードします
    a3m_resp = requests.get(f"{api_base}/result/download/{ticket}", timeout=60)
    a3m_resp.raise_for_status()
    a3m_content = a3m_resp.text
    depth = a3m_content.count("\n>")

    return MSAResult(query_id=query_id, a3m_content=a3m_content, depth=depth)


def save_a3m(msa: MSAResult, output_dir: Path) -> Path:
    """a3mファイルを保存します。"""
    output_dir.mkdir(parents=True, exist_ok=True)
    path = output_dir / f"{msa.query_id}.a3m"
    path.write_text(msa.a3m_content)
    return path

from rdkit import Chem
from rdkit.Chem import AllChem, Descriptors, Lipinski
import pandas as pd


def canonicalize_smiles(smiles: str) -> str | None:
    """SMILESの標準形を返します。解析に失敗した場合はNoneを返します。"""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    return Chem.MolToSmiles(mol, canonical=True)


def lipinski_filter(smiles: str) -> dict:
    """リピンスキーの5則が満たされているか確認し、物理化学的特性を計算します。"""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return {"valid": False}
    props = {
        "MW": Descriptors.MolWt(mol),
        "LogP": Descriptors.MolLogP(mol),
        "HBA": Lipinski.NumHAcceptors(mol),
        "HBD": Lipinski.NumHDonors(mol),
        "RotB": Lipinski.NumRotatableBonds(mol),
    }
    violations = sum([
        props["MW"] > 500,
        props["LogP"] > 5,
        props["HBA"] > 10,
        props["HBD"] > 5,
    ])
    props["valid"] = True
    props["lipinski_violations"] = violations
    props["lipinski_pass"] = violations <= 1
    return props


def prepare_ligand_library(smiles_list: list[str]) -> pd.DataFrame:
    """リガンドライブラリを標準化し、物理化学的フィルタを適用します。"""
    records = []
    for i, raw in enumerate(smiles_list):
        canon = canonicalize_smiles(raw)
        if canon is None:
            continue
        props = lipinski_filter(canon)
        records.append({
            "ligand_id": f"L{i:04d}",
            "smiles_raw": raw,
            "smiles_canonical": canon,
            **props,
        })
    df = pd.DataFrame(records)
    return df[df["lipinski_pass"]].reset_index(drop=True)

import yaml


def build_boltz_config(
    protein_sequence: str,
    protein_msa_path: Path,
    ligand_smiles: str,
    ligand_id: str,
    output_dir: Path,
) -> Path:
    """Boltz YAML設定ファイルを生成します。タンパク質-リガンドペアごとに1つです。"""
    config = {
        "version": 1,
        "sequences": [
            {
                "protein": {
                    "id": "A",
                    "sequence": protein_sequence,
                    "msa": str(protein_msa_path),
                }
            },
            {
                "ligand": {
                    "id": "L",
                    "smiles": ligand_smiles,
                }
            },
        ],
        "constraints": [],
        "properties": [
            {"affinity": {"binder": "L"}}  # 結合親和性の予測を有効にします（Boltz-2の新機能）
        ],
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    yaml_path = output_dir / f"{ligand_id}.yaml"
    with open(yaml_path, "w") as f:
        yaml.safe_dump(config, f, sort_keys=False)
    return yaml_path

import subprocess
import json


@dataclass
class BoltzResult:
    """Boltz-2予測結果。"""
    ligand_id: str
    pdb_path: Path  # 予測された3D複合構造
    plddt_mean: float  # 構造の信頼度（0〜100）
    iptm: float  # インターフェースpTM（0〜1）
    affinity_log_ic50: float  # 予測されたlog10（IC50 mol/L）（値が低いほど良い）
    inference_time_sec: float


def run_boltz_prediction(
    yaml_path: Path,
    output_dir: Path,
    device: str = "cuda",
    use_msa_server: bool = False,
) -> BoltzResult:
    """Boltz-2 CLIを実行します（実際にはPython APIも使用できます）。"""
    import time
    start = time.time()

    cmd = [
        "boltz", "predict",
        str(yaml_path),
        "--out_dir", str(output_dir),
        "--devices", "1",
        "--accelerator", "gpu" if device == "cuda" else "cpu",
    ]
    if use_msa_server:
        cmd.append("--use_msa_server")

    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
    elapsed = time.time() - start

    # Boltzはoutput_dir/{name}/{name}_model_0.pdb + confidence.jsonを生成します
    name = yaml_path.stem
    pdb_path = output_dir / name / f"{name}_model_0.pdb"
    conf_path = output_dir / name / f"confidence_{name}_model_0.json"

    with open(conf_path) as f:
        conf = json.load(f)

    # 結合親和性は別のJSONに格納されます（Boltz-2の新機能）
    affinity_path = output_dir / name / f"affinity_{name}.json"
    affinity_data = json.loads(affinity_path.read_text()) if affinity_path.exists() else {}

    return BoltzResult(
        ligand_id=name,
        pdb_path=pdb_path,
        plddt_mean=float(conf.get("complex_plddt", 0.0)),
        iptm=float(conf.get("iptm", 0.0)),
        affinity_log_ic50=float(affinity_data.get("affinity_pred_value", 0.0)),
        inference_time_sec=elapsed,
    )

from tqdm import tqdm


def batch_screen(
    protein_sequence: str,
    protein_msa: MSAResult,
    ligand_library: pd.DataFrame,
    work_dir: Path,
    device: str = "cuda",
) -> pd.DataFrame:
    """すべてのリガンドライブラリをスクリーニングします。"""
    msa_path = save_a3m(protein_msa, work_dir / "msa")

    results = []
    failures = []
    for _, row in tqdm(ligand_library.iterrows(), total=len(ligand_library), desc="スクリーニング中"):
        try:
            yaml_path = build_boltz_config(
                protein_sequence=protein_sequence,
                protein_msa_path=msa_path,
                ligand_smiles=row["smiles_canonical"],
                ligand_id=row["ligand_id"],
                output_dir=work_dir / "configs",
            )
            boltz_result = run_boltz_prediction(
                yaml_path=yaml_path,
                output_dir=work_dir / "predictions",
                device=device,
                use_msa_server=False,
            )
            results.append({
                "ligand_id": boltz_result.ligand_id,
                "smiles": row["smiles_canonical"],
                "affinity_log_ic50": boltz_result.affinity_log_ic50,
                "plddt": boltz_result.plddt_mean,
                "iptm": boltz_result.iptm,
                "pdb_path": str(boltz_result.pdb_path),
                "inference_sec": boltz_result.inference_time_sec,
                "MW": row.get("MW"),
                "LogP": row.get("LogP"),
            })
        except Exception as e:
            failures.append({"ligand_id": row["ligand_id"], "error": str(e)})

    df_results = pd.DataFrame(results).sort_values("affinity_log_ic50")
    df_failures = pd.DataFrame(failures)
    df_failures.to_csv(work_dir / "failures.csv", index=False)
    return df_results

def render_pymol(
    pdb_path: Path,
    output_png: Path,
    ray_trace: bool = True,
) -> None:
    """ヘッドレスPyMOLレンダリング。"""
    script = f"""
load {pdb_path}, complex
hide everything
show cartoon, complex and polymer
show sticks, complex and organic
color grey70, complex and polymer
color yellow, complex and organic
bg_color white
zoom complex and organic, 5
{"ray 1200, 900" if ray_trace else ""}
png {output_png}, dpi=150
quit
"""
    script_path = output_png.parent / "render.pml"
    script_path.write_text(script)
    subprocess.run(["pymol", "-cq", str(script_path)], check=True)


def render_top_k(results_df: pd.DataFrame, output_dir: Path, k: int = 20) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    for _, row in results_df.head(k).iterrows():
        render_pymol(
            pdb_path=Path(row["pdb_path"]),
            output_png=output_dir / f"{row['ligand_id']}.png",
        )

def full_screening_pipeline(
    protein_sequence: str,
    protein_id: str,
    smiles_library: list[str],
    work_dir: Path,
    device: str = "cuda",
    top_k: int = 20,
) -> pd.DataFrame:
    """FASTA + SMILESライブラリ -> ランク付けされたDataFrame。"""
    print(f"[1/6] リモートMSA構築のリクエスト（クエリの長さ={len(protein_sequence)})")
    msa = build_msa_remote(protein_sequence, query_id=protein_id)
    print(f"      MSAの深さ={msa.depth}")

    print(f"[2/6] リガンドライブラリの標準化（生={len(smiles_library)})")
    library = prepare_ligand_library(smiles_library)
    print(f"      リピンスキーパス={len(library)}")

    print(f"[3/6] バッチスクリーニングの開始")
    results = batch_screen(protein_sequence, msa, library, work_dir, device)
    results.to_csv(work_dir / "screening_results.csv", index=False)

    print(f"[4/6] 上位{top_k} PyMOLレンダリング")
    render_top_k(results, work_dir / "top_k_renders", k=top_k)

    print(f"[5/6] 完了：work_dir={work_dir}")
    print(f"[6/6] 上位5つのリガンドの概要：")
    print(results.head(5)[["ligand_id", "affinity_log_ic50", "plddt", "iptm"]].to_string(index=False))
    return results

## Chai-1 アンサンブル検証 (オプションの拡張)

Chai-1 は Boltz と同様の最先端のオープンな基盤モデルであるため、上位候補を Chai-1 で再評価することで、誤検出を減らすことができます [2]。

```python
from chai_lab.chai1 import run_inference


def chai_reverification(
    protein_sequence: str,
    ligand_smiles: str,
    output_dir: Path,
    device: str = "cuda",
) -> dict:
    """Chai-1 を使用して、同じペアを再予測します。Boltz の結果とアンサンブルスコアリングを行います。"""
    fasta_str = f">protein|name=A\n{protein_sequence}\n>ligand|name=L\n{ligand_smiles}\n"
    fasta_path = output_dir / "chai_input.fasta"
    fasta_path.write_text(fasta_str)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    result = run_inference(
        fasta_file=fasta_path,
        output_dir=output_dir,
        num_trunk_recycles=3,
        num_diffn_timesteps=200,
        seed=42,
        device=device,
        use_esm_embeddings=True,
    )
    return {
        "pdb_path": result[0],
        "iptm": float(result[0].iptm),  # 例として、実際の API を参照してください。
    }


def ensemble_ranking(
    boltz_score: float,
    chai_score: float,
    weight: float = 0.5,
) -> float:
    """2 つのモデルからのスコアの重み付き平均 (値が小さいほど良い)。"""
    return weight * boltz_score + (1 - weight) * chai_score