class monomer_ROG(Potential):
    def __init__(self, args, features, potential_scale, DEVICE, min_dist=15):

        super().__init__()

        self.use_features=True

        self.features = features
        self.args = args

        # grad_test is used to retain features in the graph
        if args['sampler']!='grad_test':
            sys.exit('This potential uses grad_test sampler. Please, set --sampler grad_test')

        self.min_dist = min_dist
        self.potential_scale = potential_scale
        self.DEVICE = DEVICE

    def get_gradients(self, features):

        Ca = features['xyz'][:,1] # [L,3]
        centroid = torch.mean(Ca, dim=0, keepdim=True) # [1,3]
        dgram = torch.cdist(Ca[None,...].contiguous(), centroid[None,...].contiguous(), p=2) # [1,L,1,3]
        dgram = torch.maximum(self.min_dist * torch.ones_like(dgram.squeeze(0)), dgram.squeeze(0)) # [L,1,3]
        rad_of_gyration = torch.sqrt( torch.sum(torch.square(dgram)) / Ca.shape[0] ) # [1]
        rad_of_gyration.backward()            

        # Get gradients from msa_masked
        self.gradients = features['msa_masked'][:,0,:22].grad

        # reset gradients
        for key in features:
            try:
                features[key].grad=None
            except AttributeError:
                pass

        return -self.gradients*self.potential_scale

### ADD NEW POTENTIALS INTO LIST DOWN BELOW ###
POTENTIALS = {'aa_bias':AACompositionalBias, 'charge':ChargeBias, 'hydrophobic':HydrophobicBias, 'PSSM':PSSMbias, 
              'monomer_ROG':monomer_ROG}

class monomer_ROG(Potential):

    def __init__(self, args, features, potential_scale, DEVICE, model, min_dist=15):

        super().__init__()

        self.features = features
        for key in features:
            try:
                features[key].requires_grad_(False)
            except:
                pass
        self.args = args

        self.min_dist = min_dist
        self.potential_scale = potential_scale
        self.DEVICE = DEVICE
        self.model=model
        self.model.eval()

    def get_gradients(self, seq):

        seq=seq.clone().detach().requires_grad_(True)

        msa_masked=self.features['msa_masked'].clone()
        msa_full=self.features['msa_full'].clone()
        msa_masked[0,0,0,:,:21] = seq
        msa_masked[0,0,0,:,22:43] = seq
        msa_full[0,0,0,:,:21] = seq
        in_seq = torch.argmax(seq, dim=-1)[None,None]

        _, _, xyz, _, _ = self.model(msa_masked[:,0], 
                                     msa_full[:,0], 
                                     in_seq[:,0],
                                     self.features['xyz_prev'], 
                                     self.features['idx_pdb'], 
                                     t1d=self.features['t1d'], 
                                     t2d=self.features['t2d'],
                                     xyz_t=self.features['xyz_t'], 
                                     alpha_t=self.features['alpha_t'],
                                     return_raw=True)

        Ca = xyz[0,:,1,:] # [L,3]
        centroid = torch.mean(Ca, dim=0, keepdim=True) # [1,3]
        dgram = torch.cdist(Ca[None,...].contiguous(), centroid[None,...].contiguous(), p=2) # [1,L,1,3]
        dgram = torch.maximum(self.min_dist * torch.ones_like(dgram.squeeze(0)), dgram.squeeze(0)) # [L,1,3]
        rad_of_gyration = torch.sqrt( torch.sum(torch.square(dgram)) / Ca.shape[0] ) # [1]
        rad_of_gyration.backward()  

        # Get gradients from msa_masked
        self.gradients = seq.grad

        # reset gradient
        seq.grad=None

        return -self.gradients*self.potential_scale

### ADD NEW POTENTIALS INTO LIST DOWN BELOW ###
POTENTIALS = {'aa_bias':AACompositionalBias, 'charge':ChargeBias, 'hydrophobic':HydrophobicBias, 'PSSM':PSSMbias, 
              'monomer_ROG':monomer_ROG}