from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Tuple, Optional, List

import numpy as np
import pandas as pd
from rdkit import Chem, DataStructs, RDLogger
from rdkit.Chem import AllChem, Descriptors, Lipinski, Crippen


@dataclass(frozen=True)
class FeatConfig:
    fp_radius: int = 2
    fp_nbits: int = 2048


def _morgan_fp(mol, radius: int, nbits: int) -> np.ndarray:
    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=nbits)
    arr = np.zeros((nbits,), dtype=np.int8)
    DataStructs.ConvertToNumpyArray(fp, arr)
    return arr


def _physchem_from_mol(mol) -> Dict[str, float]:
    return {
        "rd_mw": float(Descriptors.MolWt(mol)),
        "rd_logp": float(Crippen.MolLogP(mol)),
        "rd_tpsa": float(Descriptors.TPSA(mol)),
        "rd_hbd": float(Lipinski.NumHDonors(mol)),
        "rd_hba": float(Lipinski.NumHAcceptors(mol)),
        "rd_rotb": float(Lipinski.NumRotatableBonds(mol)),
        "rd_rings": float(Lipinski.RingCount(mol)),
        "rd_heavy_atoms": float(Descriptors.HeavyAtomCount(mol)),
    }


def _merge_override(row: pd.Series, rd_vals: Dict[str, float], override_physchem: bool) -> Dict[str, float]:
    if not override_physchem:
        return rd_vals
    out = {}
    for k, v in rd_vals.items():
        if k in row.index and pd.notna(row[k]):
            try:
                out[k] = float(row[k])
            except Exception:
                out[k] = v
        else:
            out[k] = v
    return out


def featurize_smiles(
    df: pd.DataFrame,
    smiles_col: str = "smiles",
    config: FeatConfig = FeatConfig(),
    add_physchem: bool = True,
    drop_invalid: bool = True,
    override_physchem: bool = True,
    suppress_rdkit_warnings: bool = True,
) -> Tuple[pd.DataFrame, pd.Series]:
    """
    Returns:
      X (features dataframe), valid_mask (bool Series aligned to input df)

    Key guarantee:
      - Internal lists keep same length as df (one entry per row),
        then we filter by valid_mask at the end (prevents shape mismatch).
    """
    if suppress_rdkit_warnings:
        RDLogger.DisableLog("rdApp.*")

    smiles_list = df[smiles_col].astype(str).tolist()
    n = len(smiles_list)

    # Prepare keys once
    dummy = _physchem_from_mol(Chem.MolFromSmiles("CC"))
    phys_keys = list(dummy.keys())

    fps_all: List[Optional[np.ndarray]] = [None] * n
    phys_all: List[Optional[Dict[str, float]]] = [None] * n
    valid: List[bool] = [False] * n

    for i, s in enumerate(smiles_list):
        mol = Chem.MolFromSmiles(s)
        if mol is None:
            valid[i] = False
            continue

        valid[i] = True
        fps_all[i] = _morgan_fp(mol, config.fp_radius, config.fp_nbits)

        if add_physchem:
            rd_vals = _physchem_from_mol(mol)
            phys_all[i] = _merge_override(df.iloc[i], rd_vals, override_physchem)

    valid_mask = pd.Series(valid, index=df.index)

    if drop_invalid:
        keep_idx = df.index[valid_mask]
        fps_kept = [fp for fp, ok in zip(fps_all, valid) if ok]
        if len(fps_kept) == 0:
            raise ValueError("No valid SMILES found after filtering.")
        fp_mat = np.vstack(fps_kept)

        X_fp = pd.DataFrame(
            fp_mat,
            columns=[f"fp_{i}" for i in range(config.fp_nbits)],
            index=keep_idx,
        )

        if add_physchem:
            phys_kept = [p for p, ok in zip(phys_all, valid) if ok]
            X_phys = pd.DataFrame(phys_kept, index=keep_idx)
            # stable schema
            for k in phys_keys:
                if k not in X_phys.columns:
                    X_phys[k] = np.nan
            X_phys = X_phys[phys_keys]
            X = pd.concat([X_phys, X_fp], axis=1)
        else:
            X = X_fp

        return X, valid_mask

    # drop_invalid = False (keep row count, fill invalid with zeros/NaN)
    fp_mat = np.vstack([fp if fp is not None else np.zeros((config.fp_nbits,), dtype=np.int8) for fp in fps_all])
    X_fp = pd.DataFrame(fp_mat, columns=[f"fp_{i}" for i in range(config.fp_nbits)], index=df.index)

    if add_physchem:
        phys_rows = []
        for p in phys_all:
            if p is None:
                phys_rows.append({k: np.nan for k in phys_keys})
            else:
                phys_rows.append(p)
        X_phys = pd.DataFrame(phys_rows, index=df.index)
        X_phys = X_phys[phys_keys]
        X = pd.concat([X_phys, X_fp], axis=1)
    else:
        X = X_fp

    return X, valid_mask