import numpy as np
import pandas as pd
from rdkit.Chem import PandasTools
from rdkit import Chem
from rdkit.Chem import DataStructs
from rdkit.Chem import rdFingerprintGenerator
from rdkit.Chem import AllChem

import mlx.core as mx
from mlxmolkit.fp_uint32 import fp_uint8_to_uint32
from mlxmolkit.fused_tanimoto_nlist import fused_neighbor_list_metal
from mlxmolkit.butina import butina_from_neighbor_list_csr

def rdkit_fps_and_bytes(mols, fp_radius: int, fp_nbits: int, n_cpu_threads: int):
    """Generate fps ONCE: return both RDKit fps and packed uint8 bytes for Metal."""
    from rdkit import DataStructs
    from rdkit.Chem import rdFingerprintGenerator

    generator = rdFingerprintGenerator.GetMorganGenerator(radius=fp_radius, fpSize=fp_nbits)
    fps = generator.GetFingerprints(mols, numThreads=n_cpu_threads)

    nbytes = (fp_nbits + 7) // 8
    out = np.zeros((len(mols), nbytes), dtype=np.uint8)
    for i, bv in enumerate(fps):
        bits = np.zeros((fp_nbits,), dtype=np.uint8)
        DataStructs.ConvertToNumpyArray(bv, bits)
        out[i] = np.packbits(bits, bitorder="little")[:nbytes]

    return fps, out

def run_rdkit_workflow(fps, n_mols: int, distance_threshold: float):
    """Blog workflow: BulkTanimotoSimilarity → ClusterData."""
    from rdkit.DataStructs import BulkTanimotoSimilarity
    from rdkit.ML.Cluster.Butina import ClusterData

    t0 = time.time()
    distances = []
    for i in range(n_mols):
        distances.extend(BulkTanimotoSimilarity(fps[i], fps[:i], returnDistance=True))
    t_sim = time.time() - t0

    t0 = time.time()
    clusters = ClusterData(
        np.array(distances), n_mols, distance_threshold,
        isDistData=True, distFunc=None, reordering=True,
    )
    t_clust = time.time() - t0

    return {
        "similarity": t_sim,
        "clustering": t_clust,
        "total": t_sim + t_clust,
        "n_clusters": len(clusters),
        "cluster_sizes": sorted([len(c) for c in clusters], reverse=True),
        "AllClusters": clusters,
    }

def run_mlx_workflow(fp_bytes_np: np.ndarray, similarity_cutoff: float):
    """
    Full Metal pipeline (like nvMolKit):
      1. Fused Tanimoto+threshold → CSR neighbor list (GPU, no N×N matrix)
      2. Butina greedy (CPU on CSR)
    """
    import mlx.core as mx
    from mlxmolkit.fp_uint32 import fp_uint8_to_uint32
    from mlxmolkit.fused_tanimoto_nlist import fused_neighbor_list_metal
    from mlxmolkit.butina import butina_from_neighbor_list_csr

    n_mols = fp_bytes_np.shape[0]
    fp_mx = mx.array(fp_bytes_np)
    fp_u32 = fp_uint8_to_uint32(fp_mx)

    # Fused: Tanimoto + threshold + CSR neighbor list (all on Metal, no N×N matrix)
    t0 = time.time()
    offsets, indices = fused_neighbor_list_metal(fp_u32, similarity_cutoff)
    t_fused = time.time() - t0

    # Butina greedy on CPU (CSR)
    t0 = time.time()
    result = butina_from_neighbor_list_csr(offsets, indices, n_mols, similarity_cutoff)
    t_butina = time.time() - t0

    n_edges = int(np.diff(offsets).sum())
    mem_saved_mb = (n_mols * n_mols * 4) / 1e6

    return {
        "fused_tanimoto_nlist": t_fused,
        "butina": t_butina,
        "total": t_fused + t_butina,
        "n_clusters": len(result.clusters),
        "cluster_sizes": sorted([len(c) for c in result.clusters], reverse=True),
        "n_edges": n_edges,
        "mem_saved_mb": mem_saved_mb,
        "NewClusters": result.clusters
    }

PandasTools.WriteSDF(dataframe, 'output.sdf', molColName='ROMol', idName=None, properties=list(dataframe.columns), allNumeric=False, forceV3000=False)