import numpy as np
from rdkit import Chem


def smiles2graph(sml, maxatoms):
    """Argument for the RD2NX function should be a valid SMILES sequence
    returns: the graph
    """
    m = Chem.MolFromSmiles(sml)
    m = Chem.AddHs(m)
    order_string = {
        Chem.rdchem.BondType.SINGLE: 1,
        Chem.rdchem.BondType.DOUBLE: 2,
        Chem.rdchem.BondType.TRIPLE: 3,
        Chem.rdchem.BondType.AROMATIC: 4,
    }
    N = len(list(m.GetAtoms()))
    nodes = np.zeros((N, maxatoms))
    for i in m.GetAtoms():
        nodes[i.GetIdx(), i.GetAtomicNum() - 1] = 1

    adj = np.zeros((N, N))
    for j in m.GetBonds():
        u = min(j.GetBeginAtomIdx(), j.GetEndAtomIdx())
        v = max(j.GetBeginAtomIdx(), j.GetEndAtomIdx())
        order = j.GetBondType()
        if order in order_string:
            order = order_string[order]
        else:
            raise Warning("Ignoring bond order" + order)
        adj[u, v] = 1
        adj[v, u] = 1
    adj += np.eye(N)
    return nodes, adj