## Copyright (c) 2017 Robert Bosch GmbH
## All rights reserved.
##
## This source code is licensed under the MIT license found in the
## LICENSE file in the root directory of this source tree.

from graph_transformer import sqdist
import torch
import torch.nn.functional as F
import numpy as np

def tta(x_atom, x_atom_pos, x_bond, x_bond_dist, x_triplet, x_triplet_angle, x_quad, x_quad_angle, args):
    elem_drop = args.elem_drop
    if elem_drop == 0.0:
        return x_atom, x_atom_pos, x_bond, x_bond_dist, x_triplet, x_triplet_angle, x_quad, x_quad_angle
    dev = x_atom.device
    bsz = x_atom.shape[0]
    N = x_atom.shape[1]
    M = x_bond.shape[1]
    P = x_triplet.shape[1]
    Q = x_quad.shape[1] if args.use_quad else 0
    atom_mask = (torch.zeros(bsz, N, 1).bernoulli_(1-elem_drop) / (1-elem_drop)).to(dev)
    bond_mask = (torch.zeros(bsz, M, 1).bernoulli_(1-elem_drop) / (1-elem_drop)).to(dev)
    trip_mask = (torch.zeros(bsz, P, 1).bernoulli_(1-elem_drop) / (1-elem_drop)).to(dev)
    x_atom = x_atom * atom_mask.long()
    x_atom_pos = x_atom_pos * atom_mask
    x_bond = x_bond * bond_mask.long()
    x_bond_dist = x_bond_dist * bond_mask[:,:,0]
    x_triplet = x_triplet * trip_mask.long()
    x_triplet_angle = x_triplet_angle * trip_mask[:,:,0]
    if args.use_quad:
        quad_mask = torch.zeros(bsz, Q, 1).bernoulli_(1-elem_drop) / (1-elem_drop)
        x_quad = x_quad * quad_mask
        x_quad_angle = x_quad_angle * quad_mask
    return x_atom, x_atom_pos, x_bond, x_bond_dist, x_triplet, x_triplet_angle, x_quad, x_quad_angle


def subgraph_filter(x_atom, x_atom_pos, x_bond, x_bond_dist, x_triplet, x_triplet_angle, cutout):
    D = sqdist(x_atom_pos[:,:,:3], x_atom_pos[:,:,:3])
    x_atom, x_atom_pos, x_bond, x_bond_dist, x_triplet, x_triplet_angle = \
        x_atom.clone().detach(), x_atom_pos.clone().detach(), x_bond.clone().detach(), x_bond_dist.clone().detach(), x_triplet.clone().detach(), x_triplet_angle.clone().detach()
    bsz = x_atom.shape[0]
    bonds_mask = torch.ones(bsz, x_bond.shape[1], 1).to(x_atom.device)
    for mol_id in range(bsz):
        if np.random.uniform(0,1) > cutout:
            continue
        #assert not args.use_quad, "Quads are NOT cut out yet"
        atom_dists = D[mol_id]
        atoms = x_atom[mol_id, :, 0]
        n_valid_atoms = (atoms > 0).sum().item()
        if n_valid_atoms < 10:
            continue
        idx_to_drop = np.random.randint(n_valid_atoms-1)
        dist_row = atom_dists[idx_to_drop]
        neighbor_to_drop = torch.argmin((dist_row[dist_row>0])[:n_valid_atoms-1]).item()
        if neighbor_to_drop >= idx_to_drop:
            neighbor_to_drop += 1
        x_atom[mol_id, idx_to_drop] = 0
        x_atom[mol_id, neighbor_to_drop] = 0
        x_atom_pos[mol_id, idx_to_drop] = 0
        x_atom_pos[mol_id, neighbor_to_drop] = 0
        bond_pos_to_drop = (x_bond[mol_id, :, 3] == idx_to_drop) | (x_bond[mol_id, :, 3] == neighbor_to_drop) \
                         | (x_bond[mol_id, :, 4] == idx_to_drop) | (x_bond[mol_id, :, 4] == neighbor_to_drop)
        trip_pos_to_drop = (x_triplet[mol_id, :, 2] == idx_to_drop) | (x_triplet[mol_id, :, 2] == neighbor_to_drop) \
                         | (x_triplet[mol_id, :, 3] == idx_to_drop) | (x_triplet[mol_id, :, 3] == neighbor_to_drop) \
                         | (x_triplet[mol_id, :, 4] == idx_to_drop) | (x_triplet[mol_id, :, 4] == neighbor_to_drop)
        x_bond[mol_id, bond_pos_to_drop] = 0
        x_bond_dist[mol_id, bond_pos_to_drop] = 0
        bonds_mask[mol_id, bond_pos_to_drop] = 0
        x_triplet[mol_id, trip_pos_to_drop] = 0
        x_triplet_angle[mol_id, trip_pos_to_drop] = 0
    return x_atom, x_atom_pos, x_bond, x_bond_dist, x_triplet, x_triplet_angle, bonds_mask
