import warnings
import collections
import functools
import inspect
import keras 
import json
import abc
import typing 
import numpy as np
import pandas as pd
import tensorflow as tf
import multiprocessing as mp

from pathlib import Path

from molcraft import tensors 
from molcraft import features
from molcraft import records
from molcraft import chem
from molcraft import descriptors


@keras.saving.register_keras_serializable(package='molcraft')
class GraphFeaturizer(abc.ABC):

    """Base graph featurizer.
    """

    @abc.abstractmethod
    def call(self, x: typing.Any, context: dict) -> tensors.GraphTensor:
        pass

    def get_config(self) -> dict:
        return {}
    
    @classmethod
    def from_config(cls, config: dict) -> 'GraphFeaturizer':
        return cls(**config)
    
    def save(self, filepath: str | Path, *args, **kwargs) -> None:
        save_featurizer(self, filepath, *args, **kwargs)

    @staticmethod
    def load(filepath: str | Path, *args, **kwargs) -> 'GraphFeaturizer':
        return load_featurizer(filepath, *args, **kwargs)
    
    def write_records(self, inputs: typing.Iterable, path: str | Path, **kwargs) -> None:
         records.write(
            inputs, featurizer=self, path=path, **kwargs
         )

    @staticmethod
    def read_records(path: str | Path, **kwargs) -> tf.data.Dataset:
        return records.read(
            path=path, **kwargs
        )

    def _call(self, inputs: typing.Any) -> tensors.GraphTensor:
        inputs, context = _unpack_inputs(inputs)
        if _call_kwargs(self.call):
            graph = self.call(inputs, context=context)
        else:
            graph = self.call(inputs)
        if not isinstance(graph, tensors.GraphTensor):
            graph = tensors.from_dict(graph)
        return graph

    def __call__(
        self,
        inputs: typing.Iterable,
        *,
        multiprocessing: bool = False,
        processes: int | None = None,
        device: str = '/cpu:0',
    ) -> tensors.GraphTensor:
        if not isinstance(
            inputs, (list, np.ndarray, pd.Series, pd.DataFrame, typing.Generator)
        ):
            return self._call(inputs)
        
        if isinstance(inputs, (np.ndarray, pd.Series)):
            inputs = inputs.tolist()
        elif isinstance(inputs, pd.DataFrame):
            inputs = inputs.iterrows()

        if not multiprocessing:
            outputs = [self._call(x) for x in inputs]
        else:
            with tf.device(device):
                with mp.Pool(processes) as pool:
                    outputs = pool.map(func=self._call, iterable=inputs)
        outputs = [x for x in outputs if x is not None]
        if tensors.is_scalar(outputs[0]):
            return tf.stack(outputs, axis=0)
        return tf.concat(outputs, axis=0)


@keras.saving.register_keras_serializable(package='molcraft')
class MolGraphFeaturizer(GraphFeaturizer):

    """Molecular graph featurizer.

    Converts SMILES or InChI strings to a molecular graph.

    The molecular graph may encode a single molecule or a batch of molecules.

    Example:

    >>> import molcraft 
    >>> 
    >>> featurizer = molcraft.featurizers.MolGraphFeaturizer(
    ...     atom_features=[
    ...         molcraft.features.AtomType(),
    ...         molcraft.features.NumHydrogens(),
    ...         molcraft.features.Degree(),
    ...     ],
    ...     bond_features=[
    ...         molcraft.features.BondType(),
    ...     ],
    ...     super_node=False,
    ...     self_loops=False,
    ... )
    >>> 
    >>> graph = featurizer(["N[C@@H](C)C(=O)O", "N[C@@H](CS)C(=O)O"])
    >>> graph
    GraphTensor(
        context={
            'size': <tf.Tensor: shape=[2], dtype=int32>
        },
        node={
            'feature': <tf.Tensor: shape=[13, 129], dtype=float32>
        },
        edge={
            'source': <tf.Tensor: shape=[22], dtype=int32>,
            'target': <tf.Tensor: shape=[22], dtype=int32>,
            'feature': <tf.Tensor: shape=[22, 5], dtype=float32>
        }
    )

    Args:
        atom_features:
            A list of `features.Feature` encoded as the node features.
        bond_features:
            A list of `features.Feature` encoded as the edge features.
        molecule_features:
            A list of `descriptors.Descriptor` encoded as the context feature.
        super_node:
            A boolean specifying whether to include a super node.
        self_loops:
            A boolean specifying whether self loops exist.
        include_hydrogens:
            A boolean specifying whether hydrogens should be encoded as nodes.
        wildcards:
            A boolean specifying whether wildcards exist. If True, wildcard labels will
            be encoded in the graph and separately embedded in `layers.NodeEmbedding`.
    """

    def __init__(
        self,
        atom_features: list[features.Feature] | str = 'auto',
        bond_features: list[features.Feature] | str | None = 'auto',
        molecule_features: list[descriptors.Descriptor] | str | None = None,
        super_node: bool = False,
        self_loops: bool = False,
        include_hydrogens: bool = False,
        wildcards: bool = False,
    ) -> None:
        use_default_atom_features = (
            atom_features == 'auto' or atom_features == 'default'
        )
        if use_default_atom_features:
            atom_features = [features.AtomType(), features.Degree()]
            if not include_hydrogens:
                atom_features += [features.NumHydrogens()]

        use_default_bond_features = (
            bond_features == 'auto' or bond_features == 'default'
        )
        if use_default_bond_features:
            bond_features = [features.BondType()]

        use_default_molecule_features = (
            molecule_features == 'auto' or molecule_features == 'default'
        )
        if use_default_molecule_features:
            molecule_features = [
                descriptors.MolWeight(),
                descriptors.TotalPolarSurfaceArea(),
                descriptors.LogP(),
                descriptors.MolarRefractivity(),
                descriptors.NumHeavyAtoms(),
                descriptors.NumHeteroatoms(),
                descriptors.NumHydrogenDonors(),
                descriptors.NumHydrogenAcceptors(),
                descriptors.NumRotatableBonds(),
                descriptors.NumRings(),
            ]

        self._atom_features = atom_features
        self._bond_features = bond_features
        self._molecule_features = molecule_features
        self._include_hydrogens = include_hydrogens
        self._wildcards = wildcards
        self._self_loops = self_loops
        self._super_node = super_node

    def call(
        self, 
        mol: str | chem.Mol | tuple, 
        context: dict | None = None
    ) -> tensors.GraphTensor:

        if isinstance(mol, str):
            mol = chem.Mol.from_encoding(
                mol, explicit_hs=self._include_hydrogens
            )
        elif isinstance(mol, chem.RDKitMol):
            mol = chem.Mol.cast(mol)

        data = {'context': {}, 'node': {}, 'edge': {}}
        
        data['context']['size'] = np.asarray(mol.num_atoms)

        if self._molecule_features is not None:
            data['context']['feature'] = np.concatenate(
                [f(mol) for f in self._molecule_features], axis=-1
            )

        if context:
            data['context'].update(context)

        data['node']['feature'] = np.concatenate(
            [f(mol) for f in self._atom_features], axis=-1
        )

        if self._wildcards:
            wildcard_labels = np.asarray([
                (atom.label or 0) + 1 if atom.symbol == "*" else 0
                for atom in mol.atoms
            ])
            data['node']['wildcard'] = wildcard_labels
            data['node']['feature'] = np.where(
                wildcard_labels[:, None],
                np.zeros_like(data['node']['feature']),
                data['node']['feature']
            )

        data['edge']['source'], data['edge']['target'] = mol.adjacency(
            fill='full', sparse=True, self_loops=self._self_loops
        )
        
        if self._bond_features is not None:
            bond_features = np.concatenate(
                [f(mol) for f in self._bond_features], axis=-1
            )
            if self._self_loops:
                bond_features = np.pad(bond_features, [(0, 1), (0, 0)])

            bond_indices = [
                mol.get_bond_between_atoms(i, j).index if (i != j) else -1
                for (i, j) in zip(data['edge']['source'], data['edge']['target'])
            ]

            data['edge']['feature'] = bond_features[bond_indices]

        if self._super_node:
            data = _add_super_node(data)

        return tensors.GraphTensor(**_convert_dtypes(data))
    
    def get_config(self):
        config = super().get_config()
        config.update({
            'atom_features': keras.saving.serialize_keras_object(
                self._atom_features
            ),
            'bond_features': keras.saving.serialize_keras_object(
                self._bond_features
            ),
            'molecule_features': keras.saving.serialize_keras_object(
                self._molecule_features
            ),
            'super_node': self._super_node,
            'self_loops': self._self_loops,
            'include_hydrogens': self._include_hydrogens,
            'wildcards': self._wildcards,
        })
        return config

    @classmethod
    def from_config(cls, config: dict):
        config['atom_features'] = keras.saving.deserialize_keras_object(
            config['atom_features']
        )
        config['bond_features'] = keras.saving.deserialize_keras_object(
            config['bond_features']
        )
        config['molecule_features'] = keras.saving.deserialize_keras_object(
            config['molecule_features']
        )
        return cls(**config)
    

@keras.saving.register_keras_serializable(package='molcraft')
class MolGraphFeaturizer3D(MolGraphFeaturizer):

    """3D Molecular graph featurizer.

    Converts SMILES or InChI strings to a 3d molecular graph.

    The molecular graph may encode a single molecule or a batch of molecules.

    Example:

    >>> import molcraft 
    >>> 
    >>> featurizer = molcraft.featurizers.MolGraphFeaturizer3D(
    ...     atom_features=[
    ...         molcraft.features.AtomType(),
    ...         molcraft.features.NumHydrogens(),
    ...         molcraft.features.Degree(),
    ...     ],
    ...     radius=5.0,
    ...     random_seed=42,
    ... )
    >>> 
    >>> graph = featurizer(["N[C@@H](C)C(=O)O", "N[C@@H](CS)C(=O)O"])
    >>> graph
    GraphTensor(
        context={
            'size': <tf.Tensor: shape=[2], dtype=int32>
        },
        node={
            'feature': <tf.Tensor: shape=[13, 129], dtype=float32>,
            'coordinate': <tf.Tensor: shape=[13, 3], dtype=float32>
        },
        edge={
            'source': <tf.Tensor: shape=[72], dtype=int32>,
            'target': <tf.Tensor: shape=[72], dtype=int32>,
            'feature': <tf.Tensor: shape=[72, 12], dtype=float32>
        }
    )
        
    Args:
        atom_features:
            A list of `features.Feature` encoded as the node features.
        pair_features:
            A list of `features.PairFeature` encoded as the edge features.
        molecule_features:
            A list of `descriptors.Descriptor` encoded as the context feature.
        super_node:
            A boolean specifying whether to include a super node.
        self_loops:
            A boolean specifying whether self loops exist.
        include_hydrogens:
            A boolean specifying whether hydrogens should be encoded as nodes.
        wildcards:
            A boolean specifying whether wildcards exist. If True, wildcard labels will
            be encoded in the graph and separately embedded in `layers.NodeEmbedding`.
        radius:
            A floating point value specifying maximum edge length. 
        random_seed:
            An integer specifying the random seed for the conformer generation.
    """

    def __init__(
        self,
        atom_features: list[features.Feature] | str = 'auto',
        pair_features: list[features.PairFeature] | str = 'auto',
        molecule_features: features.Feature | str | None = None,
        super_node: bool = False,
        self_loops: bool = False,
        include_hydrogens: bool = False,
        wildcards: bool = False,
        radius: int | float | None = 6.0,
        random_seed: int | None = None,
        **kwargs,
    ) -> None:
        kwargs.pop('bond_features', None)
        super().__init__(
            atom_features=atom_features,
            bond_features=None,
            molecule_features=molecule_features,
            super_node=super_node,
            self_loops=self_loops,
            include_hydrogens=include_hydrogens,
            wildcards=wildcards,
        )

        use_default_pair_features = (
            pair_features == 'auto' or pair_features == 'default'
        )
        if use_default_pair_features:
            pair_features = [features.PairDistance()]

        self._pair_features = pair_features 
        self._radius = float(radius) if radius else None
        self._random_seed = random_seed

    def call(
        self, 
        mol: str | chem.Mol | tuple, 
        context: dict | None = None
    ) -> tensors.GraphTensor:

        if isinstance(mol, str):
            mol = chem.Mol.from_encoding(
                mol, explicit_hs=True
            )
        elif isinstance(mol, chem.RDKitMol):
            mol = chem.Mol.cast(mol)

        if mol.num_conformers == 0:
            mol = chem.embed_conformers(
                mol, num_conformers=1, random_seed=self._random_seed
            )

        if not self._include_hydrogens:
            mol = chem.remove_hs(mol)

        data = {'context': {}, 'node': {}, 'edge': {}}

        data['context']['size'] = np.asarray(mol.num_atoms)

        if self._molecule_features is not None:
            data['context']['feature'] = np.concatenate(
                [f(mol) for f in self._molecule_features], axis=-1
            )

        if context:
            data['context'].update(context)

        conformer = mol.get_conformer()

        data['node']['feature'] = np.concatenate(
            [f(mol) for f in self._atom_features], axis=-1
        )
        data['node']['coordinate'] = conformer.coordinates

        if self._wildcards:
            wildcard_labels = np.asarray([
                (atom.label or 0) + 1 if atom.symbol == "*" else 0
                for atom in mol.atoms
            ])
            data['node']['wildcard'] = wildcard_labels
            data['node']['feature'] = np.where(
                wildcard_labels[:, None],
                np.zeros_like(data['node']['feature']),
                data['node']['feature']
            )

        adjacency_matrix = conformer.adjacency(
            fill='full', radius=self._radius, sparse=False, self_loops=self._self_loops,
        )

        data['edge']['source'], data['edge']['target'] = np.where(adjacency_matrix)
        
        if self._pair_features is not None:
            pair_features = np.concatenate(
                [f(mol) for f in self._pair_features], axis=-1
            )
            pair_keep = adjacency_matrix.reshape(-1).astype(bool)
            data['edge']['feature'] = pair_features[pair_keep]

        if self._super_node:
            data = _add_super_node(data)
            data['node']['coordinate'] = np.concatenate(
                [data['node']['coordinate'], conformer.centroid[None]], axis=0
            )

        return tensors.GraphTensor(**_convert_dtypes(data))
    
    @property 
    def random_seed(self) -> int | None:
        return self._random_seed
    
    @random_seed.setter
    def random_seed(self, value: int) -> None:
        self._random_seed = value 

    def get_config(self):
        config = super().get_config()
        config['radius'] = self._radius
        config['pair_features'] = keras.saving.serialize_keras_object(
            self._pair_features
        )
        config['random_seed'] = self._random_seed
        return config

    @classmethod
    def from_config(cls, config: dict):
        config['pair_features'] = keras.saving.deserialize_keras_object(
            config['pair_features']
        )
        return super().from_config(config)


def save_featurizer(
    featurizer: GraphFeaturizer, 
    filepath: str | Path, 
    overwrite: bool = True, 
    **kwargs
) -> None:
    filepath = Path(filepath)
    if filepath.suffix != '.json':
        raise ValueError(
            'Invalid `filepath` extension for saving a `GraphFeaturizer`. '
            'A `GraphFeaturizer` should be saved as a JSON file.'
        )
    if not filepath.parent.exists():
        filepath.parent.mkdir(parents=True, exist_ok=True)
    if filepath.exists() and not overwrite:
        return 
    serialized_featurizer = keras.saving.serialize_keras_object(featurizer)
    with open(filepath, 'w') as f:
        json.dump(serialized_featurizer, f, indent=4)

def load_featurizer(
    filepath: str | Path,
    **kwargs
) -> GraphFeaturizer:
    filepath = Path(filepath)
    if filepath.suffix != '.json':
        raise ValueError(
            'Invalid `filepath` extension for loading a `GraphFeaturizer`. '
            'A `GraphFeaturizer` should be saved as a JSON file.'
        )
    if not filepath.exists():
        return 
    with open(filepath, 'r') as f:
        config = json.load(f)
    return keras.saving.deserialize_keras_object(config)

def _add_super_node(
    data: dict[str, dict[str, np.ndarray]]
) -> dict[str, dict[str, np.ndarray]]:

    data['context']['size'] += 1

    num_nodes = data['node']['feature'].shape[0]
    num_edges = data['edge']['source'].shape[0]
    super_node_index = num_nodes

    add_self_loops = np.any(
        data['edge']['source'] == data['edge']['target']
    )
    if add_self_loops:
        data['edge']['source'] = np.append(
            data['edge']['source'], super_node_index
        )
        data['edge']['target'] = np.append(
            data['edge']['target'], super_node_index
        )

    data['node']['feature'] = np.pad(data['node']['feature'], [(0, 1), (0, 0)])
    data['node']['super'] = np.asarray([False] * num_nodes + [True])
    if 'wildcard' in data['node']:
        data['node']['wildcard'] = np.pad(data['node']['wildcard'], [(0, 1)])

    node_indices = list(range(num_nodes))
    super_node_indices = [super_node_index] * num_nodes

    data['edge']['source'] = np.append(
        data['edge']['source'], node_indices + super_node_indices
    )
    data['edge']['target'] = np.append(
        data['edge']['target'], super_node_indices + node_indices
    )
    
    total_num_edges = data['edge']['source'].shape[0]
    num_super_edges = (total_num_edges - num_edges) 
    data['edge']['super'] = np.asarray(
        [False] * num_edges + [True] * num_super_edges
    )

    if 'feature' in data['edge']:
        data['edge']['feature'] = np.pad(
            data['edge']['feature'], [(0, num_super_edges), (0, 0)]
        )
        
    return data

def _convert_dtypes(data: dict[str, dict[str, np.ndarray]]) -> np.ndarray:
    for outer_key, inner_dict in data.items():
        for inner_key, inner_value in inner_dict.items():
            if inner_key in ['source', 'target', 'size']:
                data[outer_key][inner_key] = inner_value.astype(np.int32)
            elif np.issubdtype(inner_value.dtype, np.integer):
                data[outer_key][inner_key] = inner_value.astype(np.int32)
            elif np.issubdtype(inner_value.dtype, np.floating):
                data[outer_key][inner_key] = inner_value.astype(np.float32)
    return data

def _unpack_inputs(inputs) -> tuple:
    if isinstance(inputs, np.ndarray):
        inputs = tuple(inputs.tolist())
    elif isinstance(inputs, list):
        inputs = tuple(inputs)
    if not isinstance(inputs, tuple):
        mol, context = inputs, {}
    elif isinstance(inputs[0], int) and isinstance(inputs[1], pd.Series):
        index, series = inputs
        mol = series.values[0]
        context = dict(
            zip(
                map(_snake_case, series.index[1:]), 
                map(np.asarray, series.values[1:])
            )
        )
        context['index'] = np.asarray(index)
    else:
        mol, *context = inputs
        context = dict(zip(['label', 'sample_weight'], map(np.asarray, context)))
    return mol, context
    
def _snake_case(x: str) -> str:
    return '_'.join(x.lower().split())

def _call_kwargs(func) -> bool:
    signature = inspect.signature(func)
    return any(
        (param.kind == inspect.Parameter.VAR_KEYWORD) or (param.name == 'context')
        for param in signature.parameters.values()
    )