from dataclasses import dataclass
from typing import Any

import torch
from torch_geometric.data import Data
from rdkit import Chem, RDLogger

RDLogger.DisableLog("rdApp.*")


ATOM_FEATURES = {
    "atomic_num": list(range(1, 119)),
    "degree": [0, 1, 2, 3, 4, 5, 6],
    "formal_charge": [-3, -2, -1, 0, 1, 2, 3],
    "hybridization": [
        Chem.rdchem.HybridizationType.SP,
        Chem.rdchem.HybridizationType.SP2,
        Chem.rdchem.HybridizationType.SP3,
        Chem.rdchem.HybridizationType.SP3D,
        Chem.rdchem.HybridizationType.SP3D2,
    ],
    "num_h": [0, 1, 2, 3, 4],
    "is_aromatic": [False, True],
    "is_in_ring": [False, True],
}


def one_hot(value: Any, choices: list) -> list[int]:
    """미매칭이면 마지막 slot에 1 (OOV 처리)."""
    if value in choices:
        idx = choices.index(value)
    else:
        idx = len(choices) - 1
    result = [0] * len(choices)
    result[idx] = 1
    return result


def atom_features(atom: Chem.Atom) -> list[float]:
    features = []
    features += one_hot(atom.GetAtomicNum(), ATOM_FEATURES["atomic_num"])
    features += one_hot(atom.GetDegree(), ATOM_FEATURES["degree"])
    features += one_hot(atom.GetFormalCharge(), ATOM_FEATURES["formal_charge"])
    features += one_hot(atom.GetHybridization(), ATOM_FEATURES["hybridization"])
    features += one_hot(atom.GetTotalNumHs(), ATOM_FEATURES["num_h"])
    features += one_hot(atom.GetIsAromatic(), ATOM_FEATURES["is_aromatic"])
    features += one_hot(atom.IsInRing(), ATOM_FEATURES["is_in_ring"])
    return features


def bond_features(bond: Chem.Bond) -> list[float]:
    bt = bond.GetBondType()
    return [
        int(bt == Chem.rdchem.BondType.SINGLE),
        int(bt == Chem.rdchem.BondType.DOUBLE),
        int(bt == Chem.rdchem.BondType.TRIPLE),
        int(bt == Chem.rdchem.BondType.AROMATIC),
        int(bond.GetIsConjugated()),
        int(bond.IsInRing()),
    ]


def smiles_to_graph(smiles: str, drug_id: str = "") -> Data | None:
    """SMILES → PyG Data. 파싱 실패 시 None."""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None or mol.GetNumAtoms() == 0:
        return None
    canonical = Chem.MolToSmiles(mol, canonical=True)
    
    node_features = torch.tensor(
        [atom_features(atom) for atom in mol.GetAtoms()],
        dtype=torch.float,
    )
    
    edge_indices = []
    edge_attrs = []
    for bond in mol.GetBonds():
        i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
        feat = bond_features(bond)
        edge_indices += [[i, j], [j, i]]  # 양방향
        edge_attrs += [feat, feat]
    
    if not edge_indices:
        edge_index = torch.empty((2, 0), dtype=torch.long)
        edge_attr = torch.empty((0, 6), dtype=torch.float)
    else:
        edge_index = torch.tensor(edge_indices, dtype=torch.long).t().contiguous()
        edge_attr = torch.tensor(edge_attrs, dtype=torch.float)
    
    data = Data(x=node_features, edge_index=edge_index, edge_attr=edge_attr)
    data.smiles = canonical
    data.drug_id = drug_id
    return data

def protein_to_onehot(sequence: str, max_length: int = 1200) -> torch.Tensor:
    """단백질 서열 → (max_length, 21) one-hot."""
    aa_to_idx = {aa: i for i, aa in enumerate("ACDEFGHIKLMNPQRSTVWY")}
    seq = sequence[:max_length].upper()
    encoded = torch.zeros(max_length, 21)
    for i, aa in enumerate(seq):
        encoded[i, aa_to_idx.get(aa, 20)] = 1.0
    return encoded


class ProteinCNN(torch.nn.Module):
    """단백질 서열 → 임베딩 벡터 (baseline)."""

    def __init__(self, output_dim: int = 128):
        super().__init__()
        self.conv1 = torch.nn.Conv1d(21, 32, kernel_size=8, padding=3)
        self.conv2 = torch.nn.Conv1d(32, 64, kernel_size=8, padding=3)
        self.conv3 = torch.nn.Conv1d(64, output_dim, kernel_size=8, padding=3)
        self.pool = torch.nn.AdaptiveAvgPool1d(1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (B, max_length, 21)
        x = x.transpose(1, 2)  # (B, 21, L)
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = torch.relu(self.conv3(x))
        return self.pool(x).squeeze(-1)


class ProteinESM2(torch.nn.Module):
    """ESM2 임베딩 wrapper (편 02와 유사)."""

    MODEL_ID = "facebook/esm2_t6_8M_UR50D"  # 8M small · 매우 가벼움

    def __init__(self, device: str = "cuda"):
        super().__init__()
        from transformers import AutoTokenizer, AutoModel
        self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_ID)
        self.model = AutoModel.from_pretrained(self.MODEL_ID).to(device).eval()
        self.device = device
        self.output_dim = self.model.config.hidden_size  # 320 for 8M

    @torch.no_grad()
    def forward(self, sequences: list[str]) -> torch.Tensor:
        embs = []
        for seq in sequences:
            inputs = self.tokenizer(seq, return_tensors="pt", truncation=True, max_length=1024).to(self.device)
            out = self.model(**inputs)
            embs.append(out.last_hidden_state.mean(dim=1).squeeze(0))
        return torch.stack(embs).to(self.device)

from torch_geometric.nn import GATConv, global_mean_pool


class GAT_DTI(torch.nn.Module):
    """분자 GAT + 단백질 CNN/ESM + concat MLP → DTI 예측."""

    def __init__(
        self,
        atom_feat_dim: int = 133,     # atom_features 출력 차원
        gat_hidden: int = 128,
        gat_heads: int = 4,
        protein_dim: int = 128,
        fusion_dim: int = 256,
        task: str = "classification",  # "regression" or "classification"
    ):
        super().__init__()
        self.gat1 = GATConv(atom_feat_dim, gat_hidden, heads=gat_heads, dropout=0.1)
        self.gat2 = GATConv(gat_hidden * gat_heads, gat_hidden, heads=1, dropout=0.1)
        self.protein_encoder = ProteinCNN(output_dim=protein_dim)
        self.fusion = torch.nn.Sequential(
            torch.nn.Linear(gat_hidden + protein_dim, fusion_dim),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.2),
            torch.nn.Linear(fusion_dim, fusion_dim // 2),
            torch.nn.ReLU(),
        )
        self.head = torch.nn.Linear(fusion_dim // 2, 1)
        self.task = task

    def forward(self, mol_data: Data, protein_onehot: torch.Tensor, return_attention: bool = False) -> torch.Tensor:
        x, edge_index = mol_data.x, mol_data.edge_index
        if return_attention:
            x, (_, alpha1) = self.gat1(x, edge_index, return_attention_weights=True)
        else:
            x = self.gat1(x, edge_index)
        x = torch.relu(x)
        x = self.gat2(x, edge_index)
        mol_emb = global_mean_pool(x, mol_data.batch)  # (B, gat_hidden)
        
        prot_emb = self.protein_encoder(protein_onehot)  # (B, protein_dim)
        combined = torch.cat([mol_emb, prot_emb], dim=-1)
        hidden = self.fusion(combined)
        out = self.head(hidden).squeeze(-1)
        
        if self.task == "classification":
            out = torch.sigmoid(out)
        
        if return_attention:
            return out, alpha1
        return out

from torch_geometric.loader import DataLoader as PyGLoader


class DTIDataset(torch.utils.data.Dataset):
    """DTI 학습 데이터셋."""

    def __init__(self, records: list[dict]):
        self.records = []
        for r in records:
            g = smiles_to_graph(r["smiles"], drug_id=r.get("drug_id", ""))
            if g is None:
                continue
            g.protein = protein_to_onehot(r["protein_sequence"])
            g.y = torch.tensor(r["label"], dtype=torch.float)
            g.target_id = r.get("target_id", "")
            self.records.append(g)

    def __len__(self):
        return len(self.records)

    def __getitem__(self, idx):
        return self.records[idx]


def train_dti(
    train_records: list[dict],
    val_records: list[dict],
    epochs: int = 50,
    batch_size: int = 64,
    lr: float = 1e-3,
    device: str = "cuda",
    task: str = "classification",
) -> torch.nn.Module:
    train_ds = DTIDataset(train_records)
    val_ds = DTIDataset(val_records)
    train_loader = PyGLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_loader = PyGLoader(val_ds, batch_size=batch_size)
    
    model = GAT_DTI(task=task).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)
    criterion = torch.nn.BCELoss() if task == "classification" else torch.nn.MSELoss()
    
    best_val_metric = 0.0 if task == "classification" else float("inf")
    for epoch in range(epochs):
        # Train
        model.train()
        total_loss = 0.0
        for batch in train_loader:
            batch = batch.to(device)
            optimizer.zero_grad()
            pred = model(batch, batch.protein.view(batch.num_graphs, -1, 21))
            loss = criterion(pred, batch.y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total_loss += loss.item()
        
        # Validate
        val_metric = evaluate_model(model, val_loader, device, task)
        print(f"Epoch {epoch+1}: train_loss={total_loss/len(train_loader):.4f}, val_metric={val_metric:.4f}")
        
        # Best 저장
        if (task == "classification" and val_metric > best_val_metric) or \
           (task == "regression" and val_metric < best_val_metric):
            best_val_metric = val_metric
            torch.save(model.state_dict(), "best_dti.pt")
    
    return model


def evaluate_model(model, loader, device: str, task: str) -> float:
    from sklearn.metrics import roc_auc_score, mean_squared_error
    model.eval()
    all_preds, all_labels = [], []
    with torch.no_grad():
        for batch in loader:
            batch = batch.to(device)
            pred = model(batch, batch.protein.view(batch.num_graphs, -1, 21))
            all_preds.extend(pred.cpu().tolist())
            all_labels.extend(batch.y.cpu().tolist())
    if task == "classification":
        return roc_auc_score(all_labels, all_preds)
    return mean_squared_error(all_labels, all_preds, squared=False)

import random


def cold_target_split(
    records: list[dict], val_target_ratio: float = 0.2, seed: int = 42,
) -> tuple[list, list]:
    """Cold-target split: validation 표적은 학습에 등장 안 함."""
    random.seed(seed)
    all_targets = list({r["target_id"] for r in records})
    random.shuffle(all_targets)
    n_val = int(len(all_targets) * val_target_ratio)
    val_targets = set(all_targets[:n_val])
    train = [r for r in records if r["target_id"] not in val_targets]
    val = [r for r in records if r["target_id"] in val_targets]
    return train, val


def cold_drug_split(
    records: list[dict], val_drug_ratio: float = 0.2, seed: int = 42,
) -> tuple[list, list]:
    """Cold-drug split: validation 약물은 학습에 등장 안 함."""
    random.seed(seed)
    all_drugs = list({r["drug_id"] for r in records})
    random.shuffle(all_drugs)
    n_val = int(len(all_drugs) * val_drug_ratio)
    val_drugs = set(all_drugs[:n_val])
    train = [r for r in records if r["drug_id"] not in val_drugs]
    val = [r for r in records if r["drug_id"] in val_drugs]
    return train, val

import pandas as pd


MPRO_SEQUENCE = (  # SARS-CoV-2 Mpro (UniProt P0DTD1, 306 residues)
    "SGFRKMAFPSGKVEGCMVQVTCGTTTLNGLWLDDVVYCPRHVICTSEDMLNPNYEDLLIRKSNHNFLVQAGNVQLRVIGHSMQNCVLKLKVDTANPKTPKYKFVR"
    "IQPGQTFSVLACYNGSPSGVYQCAMRPNFTIKGSFLNGSCGSVGFNIDYDCVSFCYMHHMELPTGVHAGTDLEGNFYGPFVDRQTAQAAGTDTTITVNVLAWLYA"
    "AVINGDRWFLNRFTTTLNDFNLVAMKYNYEPLTQDHVDILGPLSAQTGIAVLDMCASLKELLQNGMNGRTILGSALLEDEFTPFDVVRQCSGVTFQ"
)


KNOWN_MPRO_INHIBITORS = [
    # nirmatrelvir (Paxlovid 성분)
    "CC1(C)C2CC1C(NC(=O)C1CCCN1C(=O)C(NC(=O)OC(C)(C)C)C(C)(C)C)C(=O)NC(C#N)CC2=O",
    # ensitrelvir (Xocova) 등 예시
    # 실전에서는 ChEMBL · DrugBank 등에서 fetch
]


def screen_mpro_repurposing(
    model: torch.nn.Module,
    drug_smiles_list: list[tuple[str, str]],  # (drug_id, smiles)
    protein_sequence: str = MPRO_SEQUENCE,
    top_k: int = 20,
    device: str = "cuda",
) -> pd.DataFrame:
    """승인 약물 라이브러리에서 Mpro 결합 후보 상위 K 리턴."""
    model.eval()
    prot_encoded = protein_to_onehot(protein_sequence).unsqueeze(0).to(device)
    
    results = []
    for drug_id, smiles in drug_smiles_list:
        g = smiles_to_graph(smiles, drug_id)
        if g is None:
            continue
        g.batch = torch.zeros(g.x.size(0), dtype=torch.long)
        g = g.to(device)
        with torch.no_grad():
            pred = model(g, prot_encoded).item()
        results.append({
            "drug_id": drug_id,
            "smiles": smiles,
            "mpro_score": pred,
        })
    
    df = pd.DataFrame(results).sort_values("mpro_score", ascending=False)
    return df.head(top_k)


def validate_against_known_inhibitors(
    model: torch.nn.Module,
    all_screened: pd.DataFrame,
    known_inhibitors_smiles: list[str],
    top_k: int = 20,
) -> dict:
    """상위 K에 알려진 inhibitor가 얼마나 포함되나?"""
    known_set = set(known_inhibitors_smiles)
    top_smiles = set(all_screened.head(top_k)["smiles"].tolist())
    recall = len(top_smiles & known_set) / max(len(known_set), 1)
    return {
        "top_k": top_k,
        "recall_of_known": recall,
        "known_in_top_k": list(top_smiles & known_set),
    }

def visualize_attention_on_molecule(
    model: GAT_DTI,
    smiles: str,
    protein_sequence: str,
    output_path: str = "attention.png",
    device: str = "cuda",
) -> None:
    """GAT attention을 분자 위 색상 오버레이."""
    from rdkit.Chem import Draw
    from rdkit.Chem.Draw import rdMolDraw2D
    from matplotlib.cm import get_cmap
    
    g = smiles_to_graph(smiles)
    g.batch = torch.zeros(g.x.size(0), dtype=torch.long)
    g = g.to(device)
    prot = protein_to_onehot(protein_sequence).unsqueeze(0).to(device)
    
    model.eval()
    with torch.no_grad():
        pred, alpha = model(g, prot, return_attention=True)
    
    # alpha: (num_edges, num_heads). 원자별 attention aggregate
    edge_index = g.edge_index.cpu().numpy()
    alpha_np = alpha.mean(dim=1).cpu().numpy()  # (num_edges,)
    n_atoms = g.x.size(0)
    atom_attn = np.zeros(n_atoms)
    for e_idx, dest in enumerate(edge_index[1]):
        atom_attn[dest] += float(alpha_np[e_idx])
    atom_attn = atom_attn / max(atom_attn.max(), 1e-8)
    
    # RDKit로 색상 highlight
    mol = Chem.MolFromSmiles(smiles)
    cmap = get_cmap("Reds")
    highlight_colors = {i: cmap(float(atom_attn[i]))[:3] for i in range(n_atoms)}
    
    drawer = rdMolDraw2D.MolDraw2DCairo(500, 500)
    drawer.drawOptions().addAtomIndices = False
    rdMolDraw2D.PrepareAndDrawMolecule(
        drawer, mol,
        highlightAtoms=list(range(n_atoms)),
        highlightAtomColors=highlight_colors,
    )
    drawer.FinishDrawing()
    with open(output_path, "wb") as f:
        f.write(drawer.GetDrawingText())
    print(f"Attention 시각화 저장: {output_path} (예측 score={pred.item():.3f})")

def full_pipeline(
    training_records: list[dict],
    mpro_repurposing_library: list[tuple[str, str]],
    output_dir: Path,
    device: str = "cuda",
) -> None:
    """DTI 모델 학습 → Mpro 스크리닝 → attention 시각화."""
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print("[1/4] Cold-target split")
    train, val = cold_target_split(training_records, val_target_ratio=0.2)
    print(f"      train {len(train)} / val {len(val)}")
    
    print("[2/4] GAT_DTI 학습")
    model = train_dti(train, val, epochs=50, task="classification", device=device)
    
    print("[3/4] Mpro 재활용 스크리닝")
    top20 = screen_mpro_repurposing(model, mpro_repurposing_library, top_k=20, device=device)
    top20.to_csv(output_dir / "mpro_top20.csv", index=False)
    
    print("[4/4] 상위 3개 attention 시각화")
    for _, row in top20.head(3).iterrows():
        visualize_attention_on_molecule(
            model, row["smiles"], MPRO_SEQUENCE,
            output_path=str(output_dir / f"attn_{row['drug_id']}.png"),
            device=device,
        )
    
    # 알려진 inhibitor 재현 검증
    validation = validate_against_known_inhibitors(
        model, top20, KNOWN_MPRO_INHIBITORS,
    )
    print(f"알려진 inhibitor recall@20: {validation['recall_of_known']:.2f}")