#!/usr/bin/env python3
"""Quick diagnostic: can the model distinguish drugs?

Tests:
1. Within-drug ranking: for same drug, different doses, do they rank near each other?
2. Cross-drug ranking: do different drugs get different scores?
3. Drug class separation: do drugs with same target category get similar scores?
4. Drug gate value: is the gate allowing drug information to flow?

Usage:
    python scripts/drug_diagnostic.py \
        --checkpoint outputs/causal_flow_drug/best_checkpoint.pt \
        --preprocessed data/processed/sciplex3_k562_24h.pt \
        --gene-map data/processed/sciplex3_k562_24h_gene_map.json \
        --smiles data/chembl_smiles.csv \
        --target-num-genes 2085
"""

import argparse
import json
import os
import sys
import warnings
from collections import defaultdict

warnings.filterwarnings("ignore")
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

import numpy as np
import torch

from gidflow.data.sciplex_dataset import Sciplex3Dataset
from gidflow.models import CausalFlowGIDModel


def load_model(checkpoint_path: str, num_genes: int, device: torch.device) -> CausalFlowGIDModel:
    """Load model with architecture inference from state dict."""
    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
    sd = ckpt["model_state_dict"]
    has_drug = "drug_gate" in sd

    encoder_output = sd["source_encoder.mlp.6.weight"].shape[0]
    encoder_hidden = sd["source_encoder.mlp.0.weight"].shape[0]
    gap_output = sd["gap_encoder.mlp.6.weight"].shape[0]
    gap_hidden = sd["gap_encoder.mlp.0.weight"].shape[0]
    planner_hidden = sd["causal_planner.per_gene_mlp.0.weight"].shape[0]
    src_layers = len([k for k in sd if "source_encoder" in k and ".mlp." in k and "weight" in k])
    n_layers = (src_layers - 1) // 3 + 1

    pert_enc_in = sd["flow_response.pert_encoder.0.weight"].shape[1]
    drug_emb_dim = pert_enc_in - num_genes if has_drug else 0
    flow_pert_emb_dim = sd["flow_response.pert_encoder.4.weight"].shape[0]
    flow_hidden_dim = sd["flow_response.pert_encoder.0.weight"].shape[0]

    gene_emb_shape = sd.get("causal_planner.causal_estimator.gene_embedding.weight", torch.zeros(1)).shape
    causal_gene_emb_dim = gene_emb_shape[1] if len(gene_emb_shape) > 1 else 128

    model = CausalFlowGIDModel(
        num_genes=num_genes,
        encoder_hidden=encoder_hidden,
        encoder_output=encoder_output,
        gap_hidden=gap_hidden,
        gap_output=gap_output,
        causal_gene_emb_dim=causal_gene_emb_dim,
        causal_n_heads=4,
        causal_n_layers=2,
        planner_hidden=planner_hidden,
        planner_n_layers=n_layers,
        flow_latent_dim=256,
        flow_hidden_dim=flow_hidden_dim,
        flow_n_layers=3,
        flow_time_embed_dim=128,
        flow_pert_emb_dim=flow_pert_emb_dim,
        n_layers=n_layers,
        use_cooccurrence=True,
        use_latent=False,
        drug_emb_dim=drug_emb_dim,
        use_drug_encoder=has_drug,
        use_drug_gene_bridge=has_drug,
    ).to(device)

    model.load_state_dict(sd, strict=False)
    model.eval()
    print(f"  Loaded: num_genes={num_genes}, drug_emb_dim={drug_emb_dim}, use_drug={has_drug}")
    gate_val = torch.sigmoid(model.drug_gate).item()
    print(f"  Drug gate: {gate_val:.4f}")
    return model


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--preprocessed", required=True)
    parser.add_argument("--gene-map", default="")
    parser.add_argument("--smiles", default="")
    parser.add_argument("--target-num-genes", type=int, default=None)
    parser.add_argument("--device", default="cuda")
    args = parser.parse_args()

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")

    print("=== Loading dataset ===")
    ds = Sciplex3Dataset(
        h5ad_path="",
        n_hvg=2000,
        preprocessed_path=args.preprocessed,
        target_num_genes=args.target_num_genes,
        drug_smiles_csv=args.smiles,
    )
    print(f"  {len(ds._conditions)} conditions, {ds.num_genes} genes, {len(ds.unique_drugs)} drugs")

    print("\n=== Loading model ===")
    model = load_model(args.checkpoint, ds.num_genes, device)

    # Build condition lookup
    conditions = ds._conditions if not hasattr(ds, 'dataset') else ds.dataset._conditions
    get_emb = lambda name: ds._get_drug_embedding(name) if not hasattr(ds, 'dataset') else ds.dataset._get_drug_embedding(name)
    get_X = ds._X if not hasattr(ds, 'dataset') else ds.dataset._X

    # Sample cells for each condition
    rng = np.random.default_rng(42)
    cond_data = []
    for cond in conditions:
        ns = min(32, len(cond["vehicle_cell_idx"]))
        nt = min(32, len(cond["drug_cell_idx"]))
        src_idx = rng.choice(cond["vehicle_cell_idx"], size=ns, replace=False)
        tgt_idx = rng.choice(cond["drug_cell_idx"], size=nt, replace=False)
        src = torch.from_numpy(get_X[src_idx]).float()
        tgt = torch.from_numpy(get_X[tgt_idx]).float()
        emb = get_emb(cond["drug_name"])
        pert = torch.from_numpy(cond["pert_vec"]).float()
        cond_data.append({
            "drug": cond["drug_name"],
            "dose": cond["dose"],
            "target": cond.get("target", ""),
            "src": src,
            "tgt": tgt,
            "emb": emb,
            "pert": pert,
        })

    print(f"\n=== Diagnostic Tests ===")

    # Test 1: Flow matching loss for same condition (should be low)
    print("\n1. Self-score (should be lowest possible):")
    self_scores = []
    for cd in cond_data[:10]:
        src = cd["src"].unsqueeze(0).to(device)
        tgt = cd["tgt"].unsqueeze(0).to(device)
        pert = cd["pert"].unsqueeze(0).to(device)
        emb = cd["emb"].unsqueeze(0).to(device)
        src_a, tgt_a, _, _ = model._align_populations(src, tgt)
        with torch.no_grad():
            loss = model.flow_response(src_a, pert, torch.tensor([0.5], device=device), tgt_a, drug_emb=emb)
        self_scores.append(-loss.item())
    print(f"  Mean self-score (n=10): {np.mean(self_scores):.4f} ± {np.std(self_scores):.4f}")

    # Test 2: Score distribution for same drug vs different drugs
    print("\n2. Score distribution:")
    # Pick a drug with multiple conditions
    drug_conds = {}
    for i, cd in enumerate(cond_data):
        drug_conds.setdefault(cd["drug"], []).append(i)

    # For first 5 drugs with ≥3 conditions, compute within-drug vs cross-drug scores
    within_drug_scores = []
    cross_drug_scores = []

    for drug, indices in list(drug_conds.items())[:5]:
        if len(indices) < 3:
            continue
        # Within-drug: score condition 0 against other conditions of same drug
        cd0 = cond_data[indices[0]]
        for idx in indices[1:]:
            cd = cond_data[idx]
            src = cd0["src"].unsqueeze(0).to(device)
            tgt = cd["tgt"].unsqueeze(0).to(device)
            pert = cd["pert"].unsqueeze(0).to(device)
            emb = cd["emb"].unsqueeze(0).to(device)
            src_a, tgt_a, _, _ = model._align_populations(src, tgt)
            with torch.no_grad():
                loss = model.flow_response(src_a, pert, torch.tensor([0.5], device=device), tgt_a, drug_emb=emb)
            within_drug_scores.append(-loss.item())

        # Cross-drug: score against 10 random conditions of different drugs
        other_indices = [i for i in range(len(cond_data)) if i not in indices][:10]
        for idx in other_indices:
            cd = cond_data[idx]
            src = cd0["src"].unsqueeze(0).to(device)
            tgt = cd["tgt"].unsqueeze(0).to(device)
            pert = cd["pert"].unsqueeze(0).to(device)
            emb = cd["emb"].unsqueeze(0).to(device)
            src_a, tgt_a, _, _ = model._align_populations(src, tgt)
            with torch.no_grad():
                loss = model.flow_response(src_a, pert, torch.tensor([0.5], device=device), tgt_a, drug_emb=emb)
            cross_drug_scores.append(-loss.item())

    print(f"  Within-drug (same drug, diff dose):  {np.mean(within_drug_scores):.4f} ± {np.std(within_drug_scores):.4f} (n={len(within_drug_scores)})")
    print(f"  Cross-drug (diff drug):              {np.mean(cross_drug_scores):.4f} ± {np.std(cross_drug_scores):.4f} (n={len(cross_drug_scores)})")
    if within_drug_scores and cross_drug_scores:
        print(f"  Difference (within - cross):         {np.mean(within_drug_scores) - np.mean(cross_drug_scores):.4f}")
        print(f"  → Positive = model gives higher score to same-drug conditions (GOOD)")
        print(f"  → Negative = model gives higher score to different drugs (BAD)")

    # Test 3: Drug class separation
    print("\n3. Drug class separation:")
    target_scores = defaultdict(list)
    for cd in cond_data:
        src = cd["src"].unsqueeze(0).to(device)
        tgt = cd["tgt"].unsqueeze(0).to(device)
        pert = cd["pert"].unsqueeze(0).to(device)
        emb = cd["emb"].unsqueeze(0).to(device)
        src_a, tgt_a, _, _ = model._align_populations(src, tgt)
        with torch.no_grad():
            loss = model.flow_response(src_a, pert, torch.tensor([0.5], device=device), tgt_a, drug_emb=emb)
        target_scores[cd["target"]].append(-loss.item())

    # Show variance within vs between target classes
    within_class_vars = []
    for target, scores in target_scores.items():
        if len(scores) >= 3:
            within_class_vars.append(np.var(scores))

    all_scores = [s for scores in target_scores.values() for s in scores]
    print(f"  Total conditions scored: {len(all_scores)}")
    print(f"  Target classes with ≥3 conditions: {len(within_class_vars)}")
    print(f"  Mean variance within drug class: {np.mean(within_class_vars):.4f}")
    print(f"  Overall score variance: {np.var(all_scores):.4f}")
    print(f"  → Low within-class variance = model groups same-target drugs (GOOD)")

    # Test 4: Drug embedding statistics
    print("\n4. Drug embedding statistics:")
    all_embs = torch.stack([get_emb(c["drug_name"]) for c in conditions])
    norms = all_embs.norm(dim=1)
    print(f"  Embedding norms: min={norms.min():.4f}, max={norms.max():.4f}, mean={norms.mean():.4f}")
    nz = (norms > 0.01).sum().item()
    print(f"  Non-zero embeddings: {nz}/{len(all_embs)} ({100*nz/len(all_embs):.1f}%)")

    # Test 5: Random pair score distribution
    print("\n5. Random pair score distribution:")
    random_scores = []
    for _ in range(100):
        i, j = rng.choice(len(cond_data), size=2, replace=False)
        cd_i, cd_j = cond_data[i], cond_data[j]
        src = cd_i["src"].unsqueeze(0).to(device)
        tgt = cd_j["tgt"].unsqueeze(0).to(device)
        pert = cd_j["pert"].unsqueeze(0).to(device)
        emb = cd_j["emb"].unsqueeze(0).to(device)
        src_a, tgt_a, _, _ = model._align_populations(src, tgt)
        with torch.no_grad():
            loss = model.flow_response(src_a, pert, torch.tensor([0.5], device=device), tgt_a, drug_emb=emb)
        random_scores.append(-loss.item())
    print(f"  Random pair score: {np.mean(random_scores):.4f} ± {np.std(random_scores):.4f}")
    print(f"  Self-score (same cond): {np.mean(self_scores):.4f} ± {np.std(self_scores):.4f}")
    print(f"  → Self-score should be HIGHER than random pair score (GOOD)")


if __name__ == "__main__":
    main()