from dataclasses import dataclass
from pathlib import Path

import torch
import numpy as np
from Bio.PDB import PDBParser


@dataclass
class ActiveSiteConstraint:
    """활성 부위 조건 정의."""
    site_name: str                          # e.g. "serine_protease_triad"
    residue_positions: list[int]             # 생성 서열에서 활성 부위 residue 위치 (0-based)
    residue_types: list[str]                 # 각 위치의 아미노산 (예: ["S", "H", "D"])
    coordinates_backbone: np.ndarray         # (N_residues, 4, 3): N, CA, C, O 좌표
    ideal_distances: dict[tuple[int, int], float]  # 활성 부위 residue 간 이상적 거리 (Å)


AA_3TO1 = {
    "ALA": "A", "ARG": "R", "ASN": "N", "ASP": "D", "CYS": "C",
    "GLN": "Q", "GLU": "E", "GLY": "G", "HIS": "H", "ILE": "I",
    "LEU": "L", "LYS": "K", "MET": "M", "PHE": "F", "PRO": "P",
    "SER": "S", "THR": "T", "TRP": "W", "TYR": "Y", "VAL": "V",
}


def load_catalytic_triad_from_pdb(
    pdb_path: Path,
    triad_selection: list[tuple[str, int, str]],  # [(chain, resid, expected_aa)]
    site_name: str = "serine_protease_triad",
) -> ActiveSiteConstraint:
    """PDB에서 catalytic triad 좌표 추출.
    
    Chymotrypsin (4CHA) 예: [("A", 195, "S"), ("A", 57, "H"), ("A", 102, "D")]
    """
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure("target", str(pdb_path))
    
    residue_types = []
    coordinates = []
    positions = []
    for i, (chain_id, resid, expected_aa) in enumerate(triad_selection):
        try:
            chain = structure[0][chain_id]
            residue = chain[resid]
        except KeyError:
            raise ValueError(f"PDB 결측: {chain_id}:{resid}")
        aa = AA_3TO1.get(residue.get_resname().upper(), "X")
        if aa != expected_aa:
            raise ValueError(f"활성 부위 {chain_id}:{resid} 아미노산 mismatch: 기대 {expected_aa}, 실제 {aa}")
        residue_types.append(aa)
        coords = np.array([
            residue["N"].get_coord(),
            residue["CA"].get_coord(),
            residue["C"].get_coord(),
            residue["O"].get_coord(),
        ])
        coordinates.append(coords)
        positions.append(i)
    
    coordinates_arr = np.stack(coordinates, axis=0)
    
    # 활성 부위 residue 간 CA-CA 이상적 거리 (chymotrypsin 실측 기반)
    ca_positions = coordinates_arr[:, 1, :]  # CA만
    ideal_dist = {}
    for i in range(len(triad_selection)):
        for j in range(i + 1, len(triad_selection)):
            d = float(np.linalg.norm(ca_positions[i] - ca_positions[j]))
            ideal_dist[(i, j)] = d
    
    return ActiveSiteConstraint(
        site_name=site_name,
        residue_positions=positions,
        residue_types=residue_types,
        coordinates_backbone=coordinates_arr,
        ideal_distances=ideal_dist,
    )


# 예시: chymotrypsin catalytic triad
CHYMOTRYPSIN_TRIAD_4CHA = [
    ("A", 57, "H"),   # His57 (histidine base)
    ("A", 102, "D"),  # Asp102 (aspartate stabilizer)
    ("A", 195, "S"),  # Ser195 (serine nucleophile)
]

from esm.models.esm3 import ESM3
from esm.sdk.api import ESMProtein, GenerationConfig


class ESM3Designer:
    """ESM3 조건부 단백질 디자인."""

    def __init__(self, device: str = "cuda", model_name: str = "esm3-sm-open-v1"):
        self.device = device
        self.model = ESM3.from_pretrained(model_name).to(device).eval()
        self.model_name = model_name

    @torch.no_grad()
    def design_with_active_site(
        self,
        constraint: ActiveSiteConstraint,
        target_length: int = 245,       # chymotrypsin 유사 길이
        num_samples: int = 100,
        temperature: float = 0.7,
        seed: int | None = None,
    ) -> list[dict]:
        """활성 부위 조건 유지하며 신규 서열 생성.
        
        Returns: [{sequence, seed, active_site_preserved: bool}]
        """
        if seed is not None:
            torch.manual_seed(seed)
            np.random.seed(seed)
        
        generated = []
        for i in range(num_samples):
            # 초기 단백질: 서열 masked, 활성 부위 residue 위치는 고정
            seq_list = ["_"] * target_length
            
            # 활성 부위 위치 (target_length 상 임의 배치, 예: 중반)
            triad_positions = self._distribute_triad_positions(
                target_length, len(constraint.residue_positions),
            )
            for pos, aa in zip(triad_positions, constraint.residue_types):
                seq_list[pos] = aa
            initial_sequence = "".join(seq_list)
            
            protein = ESMProtein(sequence=initial_sequence)
            
            # ESM3 iterative decoding
            config = GenerationConfig(
                track="sequence",
                num_steps=max(target_length // 4, 20),
                temperature=temperature,
            )
            try:
                result = self.model.generate(protein, config)
                # 활성 부위 residue 실제로 유지됐는지 사후 확인
                preserved = all(
                    result.sequence[pos] == aa
                    for pos, aa in zip(triad_positions, constraint.residue_types)
                )
                generated.append({
                    "sequence": result.sequence,
                    "seed": (seed or 0) + i,
                    "triad_positions": triad_positions,
                    "active_site_preserved": preserved,
                    "length": len(result.sequence),
                })
            except Exception as e:
                print(f"[{i}] 생성 실패: {e}")
        
        return generated

    def _distribute_triad_positions(
        self, target_length: int, n_triad: int, spread_ratio: float = 0.4,
    ) -> list[int]:
        """활성 부위 residue를 target_length 상 균등 분포.
        
        실제 chymotrypsin은 57, 102, 195 위치. 245 residue 상 21%, 42%, 80% 위치.
        신규 디자인에서는 유사 상대 위치 유지.
        """
        chymo_ref = [57, 102, 195]
        chymo_length = 245
        positions = [
            int(target_length * (p / chymo_length))
            for p in chymo_ref[:n_triad]
        ]
        return positions

def compute_pairwise_identity(seq_a: str, seq_b: str) -> float:
    """길이가 같으면 position-wise, 다르면 alignment."""
    if len(seq_a) == len(seq_b):
        matches = sum(1 for a, b in zip(seq_a, seq_b) if a == b)
        return matches / len(seq_a)
    # 길이 다르면 global alignment
    from Bio import pairwise2
    aln = pairwise2.align.globalxx(seq_a, seq_b, one_alignment_only=True)[0]
    return aln.score / max(len(seq_a), len(seq_b))


def diversity_filter(
    sequences: list[str],
    max_pairwise_identity: float = 0.7,
) -> list[str]:
    """탐욕적 필터: 이미 채택한 서열과 identity 낮은 것만 유지."""
    kept = []
    for seq in sequences:
        keep = True
        for existing in kept:
            if compute_pairwise_identity(seq, existing) > max_pairwise_identity:
                keep = False
                break
        if keep:
            kept.append(seq)
    return kept

def blast_against_swissprot(
    query_seq: str,
    swissprot_fasta_path: Path,
    e_threshold: float = 1e-5,
) -> list[dict]:
    """BLAST로 자연 서열 유사도 검사.
    
    실전은 로컬 BLAST DB (makeblastdb + blastp) 활용.
    """
    import subprocess
    import tempfile
    
    with tempfile.NamedTemporaryFile(mode="w", suffix=".fasta", delete=False) as query_f:
        query_f.write(f">query\n{query_seq}\n")
        query_path = query_f.name
    
    try:
        result = subprocess.run(
            ["blastp",
             "-query", query_path,
             "-db", str(swissprot_fasta_path),
             "-outfmt", "6 qseqid sseqid pident evalue bitscore",
             "-evalue", str(e_threshold),
             "-max_target_seqs", "5"],
            capture_output=True, text=True, check=True,
        )
        hits = []
        for line in result.stdout.strip().split("\n"):
            if not line:
                continue
            fields = line.split("\t")
            hits.append({
                "subject": fields[1],
                "identity": float(fields[2]),
                "e_value": float(fields[3]),
                "bit_score": float(fields[4]),
            })
        return hits
    except (subprocess.CalledProcessError, FileNotFoundError):
        return []


def natural_similarity_check(
    seq: str,
    max_identity: float = 30.0,  # %
) -> tuple[bool, float]:
    """seq가 자연 서열과 max_identity 이하로 novel한지.
    
    Returns: (is_novel, max_identity_found)
    """
    hits = blast_against_swissprot(seq, Path("/path/to/swissprot.fasta"))
    if not hits:
        return True, 0.0
    max_id = max(h["identity"] for h in hits)
    return max_id < max_identity, max_id

import subprocess
import yaml


def predict_structure_boltz(sequence: str, output_dir: Path, seq_id: str) -> Path:
    """편 11의 Boltz-2 파이프라인으로 구조 예측."""
    yaml_content = {
        "version": 1,
        "sequences": [{"protein": {"id": "A", "sequence": sequence}}],
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    yaml_path = output_dir / f"{seq_id}.yaml"
    with open(yaml_path, "w") as f:
        yaml.safe_dump(yaml_content, f)
    
    try:
        subprocess.run([
            "boltz", "predict", str(yaml_path),
            "--out_dir", str(output_dir),
            "--use_msa_server",
        ], check=True, capture_output=True)
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"Boltz-2 실패 {seq_id}: {e.stderr.decode()[:500]}")
    
    return output_dir / seq_id / f"{seq_id}_model_0.pdb"


def kabsch_rmsd(coords_a: np.ndarray, coords_b: np.ndarray) -> float:
    """Kabsch 정렬 후 RMSD 계산."""
    a = coords_a - coords_a.mean(axis=0)
    b = coords_b - coords_b.mean(axis=0)
    h = a.T @ b
    u, _, vt = np.linalg.svd(h)
    d = np.sign(np.linalg.det(vt.T @ u.T))
    correction = np.eye(3)
    correction[2, 2] = d
    r = vt.T @ correction @ u.T
    a_aligned = a @ r.T
    return float(np.sqrt(np.mean(np.sum((a_aligned - b) ** 2, axis=1))))


def check_active_site_recovery(
    predicted_pdb: Path,
    triad_positions: list[int],
    original_constraint: ActiveSiteConstraint,
    rmsd_threshold: float = 2.0,
) -> tuple[bool, float, float]:
    """예측 구조가 활성 부위를 재현하나?
    
    Returns: (recovered, rmsd, mean_plddt)
    """
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure("predicted", str(predicted_pdb))
    residues_list = list(structure.get_residues())
    
    predicted_coords = []
    plddt_values = []
    for pos in triad_positions:
        if pos >= len(residues_list):
            return False, float("inf"), 0.0
        residue = residues_list[pos]
        try:
            coords = np.array([
                residue["N"].get_coord(),
                residue["CA"].get_coord(),
                residue["C"].get_coord(),
                residue["O"].get_coord(),
            ])
        except KeyError:
            return False, float("inf"), 0.0
        predicted_coords.append(coords)
        # pLDDT는 CA atom의 B-factor에 저장 (AF2/Boltz 관례)
        plddt_values.append(float(residue["CA"].get_bfactor()))
    
    predicted_arr = np.stack(predicted_coords, axis=0).reshape(-1, 3)
    reference_arr = original_constraint.coordinates_backbone.reshape(-1, 3)
    rmsd = kabsch_rmsd(predicted_arr, reference_arr)
    mean_plddt = float(np.mean(plddt_values))
    return rmsd < rmsd_threshold, rmsd, mean_plddt

from dataclasses import asdict
import pandas as pd


@dataclass
class DesignCandidate:
    sequence: str
    length: int
    seed: int
    active_site_preserved_in_seq: bool
    triad_positions: list[int]
    pdb_path: str | None
    active_site_rmsd: float
    plddt_active_site: float
    natural_identity_max: float
    novelty_score: float
    composite_score: float


def full_design_pipeline(
    catalytic_pdb: Path,
    triad_selection: list[tuple[str, int, str]],
    target_length: int,
    output_dir: Path,
    num_samples: int = 100,
    device: str = "cuda",
) -> pd.DataFrame:
    """활성 부위 → 신규 서열 생성 → 다양성 · novelty · 검증 → 랭킹."""
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print("[1/6] 활성 부위 constraint 로딩")
    constraint = load_catalytic_triad_from_pdb(catalytic_pdb, triad_selection)
    print(f"      triad: {constraint.residue_types}, distances: {constraint.ideal_distances}")
    
    print("[2/6] ESM3 조건부 생성")
    designer = ESM3Designer(device=device)
    generated = designer.design_with_active_site(
        constraint, target_length, num_samples, temperature=0.7, seed=42,
    )
    preserved = [g for g in generated if g["active_site_preserved"]]
    print(f"      생성 {len(generated)}, 활성 부위 유지 {len(preserved)}")
    
    print("[3/6] 다양성 필터")
    preserved_seqs = [g["sequence"] for g in preserved]
    diverse_seqs = diversity_filter(preserved_seqs, max_pairwise_identity=0.7)
    diverse = [g for g in preserved if g["sequence"] in diverse_seqs]
    print(f"      다양성 후 {len(diverse)}")
    
    print("[4/6] 자연 서열 novelty 검사")
    novel_candidates = []
    for g in diverse:
        is_novel, max_id = natural_similarity_check(g["sequence"])
        g["natural_identity_max"] = max_id
        g["novelty_score"] = 1.0 - (max_id / 100.0)
        if is_novel:
            novel_candidates.append(g)
    print(f"      novel (identity ≤ 30%) {len(novel_candidates)}")
    
    print(f"[5/6] Boltz-2 구조 검증 ({len(novel_candidates)}개)")
    candidates_verified = []
    for i, g in enumerate(novel_candidates):
        seq_id = f"design_{g['seed']}"
        try:
            pdb = predict_structure_boltz(g["sequence"], output_dir / "structures", seq_id)
            recovered, rmsd, plddt = check_active_site_recovery(
                pdb, g["triad_positions"], constraint,
            )
            candidates_verified.append(DesignCandidate(
                sequence=g["sequence"],
                length=g["length"],
                seed=g["seed"],
                active_site_preserved_in_seq=g["active_site_preserved"],
                triad_positions=g["triad_positions"],
                pdb_path=str(pdb),
                active_site_rmsd=rmsd,
                plddt_active_site=plddt,
                natural_identity_max=g["natural_identity_max"],
                novelty_score=g["novelty_score"],
                composite_score=0.0,  # 하단에서 계산
            ))
        except Exception as e:
            print(f"      [{i}] {seq_id} 실패: {e}")
    
    print(f"[6/6] Composite 랭킹 · CSV 저장")
    if not candidates_verified:
        return pd.DataFrame()
    
    df = pd.DataFrame([asdict(c) for c in candidates_verified])
    # composite: plddt (40%) + 활성 부위 RMSD (40%) + novelty (20%)
    df["composite_score"] = (
        df["plddt_active_site"] / 100 * 0.4 +
        (1 - df["active_site_rmsd"].clip(0, 5) / 5.0) * 0.4 +
        df["novelty_score"] * 0.2
    )
    df = df.sort_values("composite_score", ascending=False)
    df.to_csv(output_dir / "design_candidates.csv", index=False)
    print(f"      상위 5 composite: {df.head(5)['composite_score'].tolist()}")
    return df


# 실행 예시 (chymotrypsin catalytic triad → 신규 serine protease 100개)
# result = full_design_pipeline(
#     catalytic_pdb=Path("./4CHA.pdb"),
#     triad_selection=CHYMOTRYPSIN_TRIAD_4CHA,
#     target_length=245,
#     output_dir=Path("./serine_protease_design"),
#     num_samples=100,
# )

def run_rfdiffusion_motif(
    active_site_pdb: Path,
    triad_residues: list[int],
    output_dir: Path,
    num_designs: int = 100,
) -> list[Path]:
    """RFdiffusion motif scaffolding.
    
    catalytic triad를 motif로 고정하고 나머지 backbone 생성.
    실전은 RosettaCommons RFdiffusion GitHub 예시 참조.
    """
    # 개념 stub: RFdiffusion CLI 호출
    # subprocess.run(["python", "run_inference.py", ...])
    return []  # 실전 구현 시 완성


def run_proteinmpnn(
    backbone_pdb: Path,
    output_dir: Path,
    num_sequences: int = 10,
) -> list[str]:
    """ProteinMPNN으로 backbone에 서열 배치."""
    # subprocess.run(["python", "protein_mpnn_run.py", ...])
    return []