from dataclasses import dataclass
from pathlib import Path

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


@dataclass
class ActiveSiteConstraint:
    \"\"\"Definir condiciones del sitio activo.\"\"\"
    site_name: str                          # e.g. "serine_protease_triad"
    residue_positions: list[int]             # Posición (0-based) del residuo en el sitio activo en la secuencia generada
    residue_types: list[str]                 # Aminoácidos en cada posición (ej. [\"S\", \"H\", \"D\"])
    coordinates_backbone: np.ndarray         # (N_residuos, 4, 3): Coordenadas N, CA, C, O
    ideal_distances: dict[tuple[int, int], float]  # Distancia ideal entre residuos del sitio activo (Å)


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:
    \"\"\"Extraer coordenadas del triplete catalítico desde PDB.

    Ejemplo de Quimotripsina (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 faltante: {chain_id}:{resid}\")
        aa = AA_3TO1.get(residue.get_resname().upper(), "X")
        if aa != expected_aa:
            raise ValueError(f\"Discrepancia de aminoácido en sitio activo {chain_id}:{resid}: esperado {expected_aa}, obtenido {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)

    # Distancia ideal CA-CA entre residuos del sitio activo (basado en datos experimentales de quimotripsina)
    ca_positions = coordinates_arr[:, 1, :]  # Solo 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,
    )


# Ejemplo: tríada catalítica de quimotripsina
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:
    """Diseño de proteínas condicional de 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,       # longitud similar a la quimotripsina
        num_samples: int = 100,
        temperature: float = 0.7,
        seed: int | None = None,
    ) -> list[dict]:
        """Genera nuevas secuencias manteniendo las condiciones del sitio activo.

        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):
            # Proteína inicial: secuencia enmascarada, posiciones de residuos del sitio activo fijas
            seq_list = ["_"] * target_length

            # Posiciones del sitio activo (distribución aleatoria dentro de target_length, ej.: mitad)
            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)
                # Verificar post-hoc si el residuo del sitio activo se mantuvo realmente
                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}] Fallo en la generación: {e}")

        return generated

    def _distribute_triad_positions(
        self, target_length: int, n_triad: int, spread_ratio: float = 0.4,
    ) -> list[int]:
        """Distribución uniforme de residuos del sitio activo a lo largo de target_length.

        La quimotripsina real está en las posiciones 57, 102 y 195. En una cadena de 245 residuos, esto corresponde al 21%, 42% y 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:
    """Si las longitudes son iguales, posición por posición; si no, alineamiento."""
    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)
    # Si las longitudes difieren, alineamiento global
    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]:
    """Filtro voraz: mantiene solo aquellos con baja identidad respecto a las secuencias ya aceptadas."""
    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]:
    """Verificación de similitud de secuencias naturales mediante BLAST.

    En la práctica, se utiliza una base de datos local de BLAST (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]:
    """Comprueba si seq es lo suficientemente novedosa, con identidad máxima respecto a secuencias naturales.

    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:
    """Predicción de estructura mediante el pipeline Boltz-2 del capítulo 11."""
    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"Fallo de 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:
    """Cálculo de RMSD tras el alineamiento de Kabsch."""
    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]:
    """¿Reproduce la estructura predicha el sitio activo?

    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 se almacena en el factor B del átomo CA (convención 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:
    """Sitio activo → generación de nueva secuencia → diversidad, novedad y validación → clasificación."""
    output_dir.mkdir(parents=True, exist_ok=True)

    print("[1/6] Carga de restricciones del sitio activo")
    constraint = load_catalytic_triad_from_pdb(catalytic_pdb, triad_selection)
    print(f"      triad: {constraint.residue_types}, distances: {constraint.ideal_distances}")

    print("[2/6] Generación condicional de 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"      Generados {len(generated)}, sitio activo conservado {len(preserved)}")

    print("[3/6] Filtro de diversidad")
    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"      Tras diversidad {len(diverse)}")

    print("[4/6] Prueba de novedad de secuencia natural")
    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] Verificación estructural con Boltz-2 ({len(novel_candidates)} casos)")
    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,  # calculado más abajo
            ))
        except Exception as e:
            print(f"      [{i}] {seq_id} fallido: {e}")

    print(f"[6/6] Clasificación compuesta · Guardar CSV")
    if not candidates_verified:
        return pd.DataFrame()

    df = pd.DataFrame([asdict(c) for c in candidates_verified])
    # compuesto: plddt (40%) + RMSD del sitio activo (40%) + novedad (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"      Top 5 compuestos: {df.head(5)['composite_score'].tolist()}")
    return df


# Ejemplo de ejecución (tríada catalítica de la quimotripsina → 100 nuevas serina proteasas)
# 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.

    Para casos prácticos, consulte el ejemplo de GitHub de RosettaCommons RFdiffusion.
    """
    """
    # Concepto stub: llamada a la CLI de RFdiffusion
    # subprocess.run(["python", "run_inference.py", ...])
    return []  # Implementación completa en producción


def run_proteinmpnn(
    backbone_pdb: Path,
    output_dir: Path,
    num_sequences: int = 10,
) -> list[str]:
    """Coloca la secuencia en el backbone con ProteinMPNN."""
    # subprocess.run(["python", "protein_mpnn_run.py", ...])
    return []