import dgl
from dgl.data.utils import download, extract_archive, get_download_dir
import rdkit.Chem as Chem
import torch
from torch.utils.data import Dataset
from typing import List, Tuple

from structures.moltree import MolTree
from structures.vocab import Vocab
import structures.mol_features as mf
from structures.mol_features import N_ATOM_FEATS as ATOM_FDIM_DEC
from structures.mol_features import N_BOND_FEATS as BOND_FDIM_DEC
from models.modules.gcn import mol2dgl_enc

_url = 'https://s3-ap-southeast-1.amazonaws.com/dgl-data-cn/dataset/jtnn.zip'

def _unpack_field(examples, field):
    try:
        retval = [e[field] for e in examples]
    except KeyError:
        retval = None
    return retval

def mol2dgl_dec(cand_batch: List[Tuple[Chem.rdchem.Mol, 
                MolTree, int]]):
    cand_graphs = []
    tree_mess_source_edges = [] # map these edges from trees to...
    tree_mess_target_edges = [] # these edges on candidate graphs
    tree_mess_target_nodes = []
    n_nodes = 0
    n_edges = 0
    atom_x = []
    bond_x = []

    for mol, mol_tree, ctr_node_id in cand_batch:
        n_atoms = mol.GetNumAtoms()
        n_bonds = mol.GetNumBonds()

        ctr_node = mol_tree.nodes_dict[ctr_node_id]
        ctr_bid = ctr_node['idx']
        g = dgl.DGLGraph()

        for i, atom in enumerate(mol.GetAtoms()):
            assert i == atom.GetIdx()
            atom_x.append(mf.get_atom_features(atom))
        g.add_nodes(n_atoms)

        bond_src = []
        bond_dst = []
        for i, bond in enumerate(mol.GetBonds()):
            a1 = bond.GetBeginAtom()
            a2 = bond.GetEndAtom()
            begin_idx = a1.GetIdx()
            end_idx = a2.GetIdx()
            features = mf.get_bond_features(bond)

            bond_src.append(begin_idx)
            bond_dst.append(end_idx)
            bond_x.append(features)
            bond_src.append(end_idx)
            bond_dst.append(begin_idx)
            bond_x.append(features)

            x_nid, y_nid = a1.GetAtomMapNum(), a2.GetAtomMapNum()
            # Tree node ID in the batch
            x_bid = mol_tree.nodes_dict[x_nid - 1]['idx'] if x_nid > 0 else -1
            y_bid = mol_tree.nodes_dict[y_nid - 1]['idx'] if y_nid > 0 else -1
            if x_bid >= 0 and y_bid >= 0 and x_bid != y_bid:
                if mol_tree.has_edge_between(x_bid, y_bid):
                    tree_mess_target_edges.append((begin_idx + n_nodes, end_idx + n_nodes))
                    tree_mess_source_edges.append((x_bid, y_bid))
                    tree_mess_target_nodes.append(end_idx + n_nodes)
                if mol_tree.has_edge_between(y_bid, x_bid):
                    tree_mess_target_edges.append((end_idx + n_nodes, begin_idx + n_nodes))
                    tree_mess_source_edges.append((y_bid, x_bid))
                    tree_mess_target_nodes.append(begin_idx + n_nodes)

        n_nodes += n_atoms
        g.add_edges(bond_src, bond_dst)
        cand_graphs.append(g)

    return cand_graphs, torch.stack(atom_x), \
            torch.stack(bond_x) if len(bond_x) > 0 else torch.zeros(0), \
            torch.LongTensor(tree_mess_source_edges), \
            torch.LongTensor(tree_mess_target_edges), \
            torch.LongTensor(tree_mess_target_nodes)
            
class JTNNCollator(object):
    def __init__(self, 
                 vocab: Vocab, 
                 training: bool, 
                 intermediates:bool = False):
        self.intermediates = intermediates
        self.training = training
        self.vocab = vocab
        
    @staticmethod
    def _batch_and_set(graphs: List,
                       atom_x: torch.Tensor,
                       bond_x: torch.Tensor,
                       flatten: bool) -> dgl.BatchedDGLGraph:
        if flatten:
            graphs = [g for f in graphs for g in f]
            
        graph_batch = dgl.batch(graphs)
        graph_batch.ndata['x'] = atom_x
        graph_batch.edata.update({
            'x': bond_x,
            'src_x': atom_x.new(bond_x.shape[0], atom_x.shape[1]).zero_(),
        })
        return graph_batch

    def __call__(self, examples):
        # get list of trees
        mol_trees = _unpack_field(examples, 'mol_tree')
        #try:
        #    smiles = _unpack_field(examples, 'smiles')
        #except KeyError:
        #    smiles = None
        wid = _unpack_field(examples, 'wid')
        for _wid, mol_tree in zip(wid, mol_trees):
            mol_tree.ndata['wid'] = torch.LongTensor(_wid)

        # TODO: either support pickling or get around ctypes pointers using scipy
        # batch molecule graphs
        mol_graphs = _unpack_field(examples, 'mol_graph')
        atom_x = torch.cat(_unpack_field(examples, 'atom_x_enc'))
        bond_x = torch.cat(_unpack_field(examples, 'bond_x_enc'))
        mol_graph_batch = self._batch_and_set(mol_graphs, atom_x, bond_x, False)

        result = {
                'mol_trees': mol_trees,
                'mol_graph_batch': mol_graph_batch,
                }

        if not self.training:
            return result

        # batch candidate graphs
        cand_graphs = _unpack_field(examples, 'cand_graphs')
        cand_batch_idx = []
        atom_x = torch.cat(_unpack_field(examples, 'atom_x_dec'))
        bond_x = torch.cat(_unpack_field(examples, 'bond_x_dec'))
        tree_mess_src_e = _unpack_field(examples, 'tree_mess_src_e')
        tree_mess_tgt_e = _unpack_field(examples, 'tree_mess_tgt_e')
        tree_mess_tgt_n = _unpack_field(examples, 'tree_mess_tgt_n')

        n_graph_nodes = 0
        n_tree_nodes = 0
        for i in range(len(cand_graphs)):
            tree_mess_tgt_e[i] += n_graph_nodes
            tree_mess_src_e[i] += n_tree_nodes
            tree_mess_tgt_n[i] += n_graph_nodes
            n_graph_nodes += sum(g.number_of_nodes() for g in cand_graphs[i])
            n_tree_nodes += mol_trees[i].number_of_nodes()
            cand_batch_idx.extend([i] * len(cand_graphs[i]))
        tree_mess_tgt_e = torch.cat(tree_mess_tgt_e)
        tree_mess_src_e = torch.cat(tree_mess_src_e)
        tree_mess_tgt_n = torch.cat(tree_mess_tgt_n)

        cand_graph_batch = self._batch_and_set(cand_graphs, atom_x, bond_x, True)

        # batch stereoisomers
        stereo_cand_graphs = _unpack_field(examples, 'stereo_cand_graphs')
        atom_x = torch.cat(_unpack_field(examples, 'stereo_atom_x_enc'))
        bond_x = torch.cat(_unpack_field(examples, 'stereo_bond_x_enc'))
        stereo_cand_batch_idx = []
        for i in range(len(stereo_cand_graphs)):
            stereo_cand_batch_idx.extend([i] * len(stereo_cand_graphs[i]))

        if len(stereo_cand_batch_idx) > 0:
            stereo_cand_labels = [
                    (label, length)
                    for ex in _unpack_field(examples, 'stereo_cand_label')
                    for label, length in ex
                    ]
            stereo_cand_labels, stereo_cand_lengths = zip(*stereo_cand_labels)
            stereo_cand_graph_batch = self._batch_and_set(
                    stereo_cand_graphs, atom_x, bond_x, True)
        else:
            stereo_cand_labels = []
            stereo_cand_lengths = []
            stereo_cand_graph_batch = None
            stereo_cand_batch_idx = []

        result.update({
            'cand_graph_batch': cand_graph_batch,
            'cand_batch_idx': cand_batch_idx,
            'tree_mess_tgt_e': tree_mess_tgt_e,
            'tree_mess_src_e': tree_mess_src_e,
            'tree_mess_tgt_n': tree_mess_tgt_n,
            #'smiles': smiles,
            'stereo_cand_graph_batch': stereo_cand_graph_batch,
            'stereo_cand_batch_idx': stereo_cand_batch_idx,
            'stereo_cand_labels': stereo_cand_labels,
            'stereo_cand_lengths': stereo_cand_lengths,
            })
        
        if self.intermediates:
            result.update({
                'img_grid': _unpack_field(examples, 'img_grid'),
                'smiles_img': _unpack_field(examples, 'smiles_img')
            })

        return result
    
"""
From JTNN:
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
class MoleculeDataset(Dataset):

    def __init__(self, data_file):
        with open(data_file) as f:
            self.data = [line.strip("\r\n ").split()[0] for line in f]

    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        smiles = self.data[idx]
        mol_tree = MolTree(smiles)
        mol_tree.recover()
        mol_tree.assemble()
        return mol_tree
"""
#from rdkit.Chem import Draw
from matplotlib.pyplot import imshow

class JTNNDataset(Dataset):
    def __init__(self, data, vocab, training=True, intermediates=False):
        self.dir = get_download_dir()
        self.gpu = torch.cuda.is_available()
        self.intermediates = intermediates
        self.training = training
        self.zip_file_path='{}/jtnn.zip'.format(self.dir)
        print('Loading data...')
        download(_url, path=self.zip_file_path)
        
        extract_archive(self.zip_file_path, '{}/jtnn'.format(self.dir))
        data_file = '{}/jtnn/{}.txt'.format(self.dir, data)
        with open(data_file) as f:
            self.data = [line.strip("\r\n ").split()[0] for line in f]
        
        self.vocab_file = '{}/jtnn/{}.txt'.format(self.dir, vocab)
        print('Loading finished.')
        print('\tNum samples:', len(self.data))
        print('\tVocab file:', self.vocab_file)
        self.vocab = Vocab([x.strip("\r\n ") for x in open(self.vocab_file)])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        smiles = self.data[idx]
        mol_tree = MolTree(smiles)
        
        # prebuild the molecule graph
        mol_graph, \
        atom_x_enc, \
        bond_x_enc, \
        wid = mol_tree.encode(recssemble=True, vocab=self.vocab)

        result = {
                'mol_tree': mol_tree,  # MolTree
                'mol_graph': mol_graph,  #
                'atom_x_enc': atom_x_enc,  # torch.Tensor
                'bond_x_enc': bond_x_enc,  # torch.Tensor
                'wid': wid,  # List[int]
                }

        if not self.training:
            return result

        # prebuild the candidate graph list
        cands = []
        for node_id, node in mol_tree.nodes_dict.items():
            # fill in ground truth
            if node['label'] not in node['cands']:
                node['cands'].append(node['label'])
                node['cand_mols'].append(node['label_mol'])

            if node['is_leaf'] or len(node['cands']) == 1:
                continue
                
            """
            cand of type: Chem.rdchem.Mol (Kekulized)
            mol_tree of type: MolTree
            node_id of type: int
            """
            cands.extend([(cand, mol_tree, node_id) 
                         for cand in node['cand_mols']])
          
        smile_img = None
        img_grid = None
        if len(cands) > 0:
            cand_graphs, atom_x_dec, bond_x_dec, tree_mess_src_e, \
                    tree_mess_tgt_e, tree_mess_tgt_n = mol2dgl_dec(cands)
            #print('%'*50)
            #print('Candidates:')
            #for cand in cands:
            #    print(Chem.MolToSmiles(cand[0]))
            #print('%'*50)
            
            
            #smile_img = Draw.MolToImage(Chem.MolFromSmiles(smiles))
            #imshow(smile_img)
            #img_grid = Draw.MolsToGridImage([cand[0] for cand in cands][0:12],
            #                                molsPerRow=3,
            #                                subImgSize=(300,200),
            #                                legends=smiles)
            #imshow(img_grid)
        else:
            cand_graphs = []
            atom_x_dec = torch.zeros(0, ATOM_FDIM_DEC)
            bond_x_dec = torch.zeros(0, BOND_FDIM_DEC)
            tree_mess_src_e = torch.zeros(0, 2).long()
            tree_mess_tgt_e = torch.zeros(0, 2).long()
            tree_mess_tgt_n = torch.zeros(0).long()
                

        # prebuild the stereoisomers
        cands = mol_tree.stereo_cands
        if len(cands) > 1:
            if mol_tree.smiles3D not in cands:
                cands.append(mol_tree.smiles3D)

            stereo_graphs = [mol2dgl_enc(c) for c in cands]
            stereo_cand_graphs, stereo_atom_x_enc, stereo_bond_x_enc = \
                    zip(*stereo_graphs)
            stereo_atom_x_enc = torch.cat(stereo_atom_x_enc)
            stereo_bond_x_enc = torch.cat(stereo_bond_x_enc)
            stereo_cand_label = [(cands.index(mol_tree.smiles3D), len(cands))]
        else:
            stereo_cand_graphs = []
            stereo_atom_x_enc = torch.zeros(0, atom_x_enc.shape[1])
            stereo_bond_x_enc = torch.zeros(0, bond_x_enc.shape[1])
            stereo_cand_label = []

            
        result.update({
            'cand_graphs': cand_graphs,
            'atom_x_dec': atom_x_dec,
            'bond_x_dec': bond_x_dec,
            'tree_mess_src_e': tree_mess_src_e,
            'tree_mess_tgt_e': tree_mess_tgt_e,
            'tree_mess_tgt_n': tree_mess_tgt_n,
            'stereo_cand_graphs': stereo_cand_graphs,
            'stereo_atom_x_enc': stereo_atom_x_enc,
            'stereo_bond_x_enc': stereo_bond_x_enc,
            'stereo_cand_label': stereo_cand_label,
            'smiles': smiles
            })
        """
        if self.intermediates:
            result.update({
                'smiles_img': smile_img,
                'img_grid': img_grid
            })
"""
        return result