<reponame>RosettaCommons/RFDesign
import torch
import sys
import random
from icecream import ic 
import numpy as np
from kinematics import xyz_to_t2d
from operator import itemgetter
#####################################
# Misc functions for mask generation
#####################################
def get_ss(ss_string):
    """
    Converts secondary structure string into a list of indices of where secondary structure changes (plus start and end coordinates of the sequence)
    """
    output = [0]
    for i in range(1,len(ss_string)):
        if ss_string[i] != ss_string[i-1]:
            output.append(i)
    if output[-1] != len(ss_string)-1:
        output.append(len(ss_string)-1)
    return output

def get_masks(ss, min_length, max_length, min_flank, max_flank):
    """
    Takes list of indices where secondary structure changes, and outputs coordinates for a masked region which masks out at least one whole secondary structure region, and allows flanking.
    """
    flank_width = random.randint(min_flank, max_flank)
    try:       
        options=[]
        for i in range(min_length, max_length+1):
            temp=[]
            for j in range(flank_width,ss[-1]-flank_width-i):
                if np.digitize([j,j+i], ss)[1] - np.digitize([j,j+i], ss)[0] > 1:
                    temp.append((j, j+i))
            options.append(temp)
        return random.choice(random.choice(options)), flank_width
        
    except:
        try:
            options=[]
            for i in range(min_length, max_length+1):
                temp=[]
                for j in range(flank_width,ss[-1]-flank_width-i):
                    if np.digitize([j,j+i], ss)[1] - np.digitize([j,j+i], ss)[0] > 0:
                        temp.append((j, j+i))
                options.append(temp)
            return random.choice(random.choice(options)), flank_width
        except:
            try:
                length = random.randint(min_length, max_length+1)
                start = random.randint(flank_width, ss[-1]-length-flank_width)
                return (start, start+length), flank_width
            except:
                return (4,7),3 #temporary fix to prevent crashing on small proteins

def get_double_masks(ss, min_length, max_length):
    '''
    Takes list of indices where secondary structure changes, and outputs coordinates for two masked regions which mask out one whole secondary structure region. No flanking.
    '''
    proportion = random.uniform(0.3,0.7)
    split_idx = int(np.round(ss[-1] * proportion,0))
    ss_1 = [i for i in ss if i < split_idx]
    ss_1.append(split_idx)
    ss_2 = [split_idx]
    ss_2.extend([i for i in ss if i > split_idx])
    split_1,_ = get_masks(ss_1, min_length, max_length,0,0)
    split_2,_ = get_masks(ss_2, min_length, max_length,0,0)
    ic(ss_1)
    ic(ss_2)
    return split_1, split_2

def get_double_masks_proximal(xyz_t,min_length,max_length,min_gap):
    temp_xyzt = torch.clone(xyz_t.unsqueeze(dim=0))
    t2d = xyz_to_t2d(temp_xyzt)[0,0,:,:,:37].numpy()
    t2d_flat = np.argmax(t2d, axis=-1)
    t2d_close = np.where(t2d_flat>11,0,1) #set anything less than 8A to 1, and more than to 0
    L = np.shape(t2d_close)[0]
    
    temp = []
    try:
        length1 = random.randint(np.minimum(min_length,L-min_gap),np.minimum(max_length,L-min_gap))
        length2 = random.randint(np.minimum(min_length,L-min_gap),np.minimum(max_length,L-min_gap))
        for j in range(0,L-length1-length2-min_gap):
            for k in range(min_gap+length2+j,L-length1):
                temp.append([j,k,np.mean(t2d_close[j:j+length1,k:k+length2])])
        sorted_temp = sorted(temp, key=itemgetter(2),reverse=True)
        choice = random.randint(0,int(0.05*len(sorted_temp)))
        output = sorted_temp[choice]
    except: #this is a temporary fix to handle small proteins
        length1 = 2
        length2 = 2
        output = [4,10,0.0] 
    return output, length1, length2

    

def get_rand_ss_idx(ss):
    """
    Takes list of indices where secondary structure changes, and outputs a single coordinate from within each ss element
    """
    output = []
    for i in range(len(ss)-1):
        if ss[i+1] - ss[i] > 4: #don't sample if ss region is smaller than 5 residues
            output.append(random.randint(ss[i],ss[i+1]-1))
    return output

def move_coords(xyz_t, translate_mask, dist_sd=1):
    """ Moves coordinates specified (False) in translate_mask.
        Moves by value picked from normal distribution (0 +/- sd = dist_sd)
    """
    B, N, L,res,coord = xyz_t.size()
    rand_mut = torch.stack((torch.randn(B,N,L,1).expand(B,N,L,3)*dist_sd,torch.randn(B,N,L,1).expand(B,N,L,3)*dist_sd,torch.randn(B,N,L,1).expand(B,N,L,3)*dist_sd),dim=-1)
    rand_mut[:,:,translate_mask.squeeze(),:,:] = 0
    xyz_t = torch.add(xyz_t, rand_mut)
    return xyz_t

#####################################
# Main mask generator function
#####################################

def generate_masks(msa, t1d, xyz_t, msa_mask_input, msa_full, task, ss, loader_params, mask_seq=None):
    '''
    Generates masks for different RF training tasks.
    separate into input masks and loss masks.
    all input masks are outputted as 1/True = unmasked, 0/False = masked.
    output masks are 1/True if scored, or 0/False if not.
    input masks are msa, t1d, str, t2d, t1d_coord_mask.
    loss masks are lddt, dih, bond, c6d, token.
    '''
    
    B,N,L = msa.size()
    t2d_b = xyz_t.shape[0]
    hal_mask_low = loader_params['HAL_MASK_LOW']
    hal_mask_high = loader_params['HAL_MASK_HIGH']
    flank_low = loader_params['FLANK_LOW']
    flank_high = loader_params['FLANK_HIGH']
    str2seqfull_low = loader_params['STR2SEQ_FULL_LOW']
    str2seqfull_high = loader_params['STR2SEQ_FULL_HIGH']
    constant_recycle = loader_params['CONSTANT_RECYCLE']

    #generate generic masks (blank masks)
    #input_masks
    t1d_mask = torch.full_like(t1d, True).bool()
    t1d_coord_mask = torch.full_like(t1d, False).bool()
    t1d_conf_mask = torch.full_like(t1d, 0.9) #scale confidences for non-fixedbbdes/hal tasks
    t2d_mask = torch.ones(t2d_b,L,L,44).bool() ### new t2ds are only 44 in last dimension (37+6+1)
    translate_mask = torch.ones(L).bool()
    msa_masktok_mask =  torch.ones((B,N,L), device=msa.device).bool() #mask for converting residues to mask token
    str_mask = torch.full_like(xyz_t, True).bool()
    
    #loss masks
    c6d_mask = torch.ones(B,37,L,L).bool() #[B,37,L,L]
    str_out_mask = torch.ones(L).bool()
    bond_mask = torch.ones(B,N,L).bool()
    dih_mask = lddt_mask = torch.ones(N,L).bool()
    
    
    
    if task == 'seq2str':
        '''
        Classic structure prediction task.
        '''
        #loss masks
        
        token_mask = msa_mask_input.clone()
        
    elif task == 'hal':
        '''
        This is Joe's hallucination task, where a contiguous region is masked, along with flanks, and two residues either end are given (but not their angles).
        The two floating points are randomly moved some distance.
        Scored on everything but the flank regions (may want to change this to only score central inpainted region
        '''
        ss_idx = get_ss(ss)
        splice, flank_width = get_masks(ss_idx, hal_mask_low, hal_mask_high, flank_low, flank_high)
        hal_mask_len = splice[1]-splice[0] 
        #input masks
        msa_masktok_mask[:,:,splice[0]-flank_width:splice[1]+flank_width] = False # flanks and central
        
        str_mask[:,splice[0]-flank_width:splice[1]+flank_width,:,:] = False # flanks and central
        str_mask[:,splice[0]-1,:,:] = True # immediate two flanking residues
        str_mask[:,splice[1],:,:] = True
        
        translate_mask[splice[0]-1] = False #imeediate two flanking residues
        translate_mask[splice[1]] = False
        
        t1d_mask[:,splice[0]-flank_width:splice[1]+flank_width,:] = False #flanks and central 
        # make mask to add to t1d to give 'fake' confidence in two floating coordinates
        t1d_coord_mask[:,splice[0]-1,:] = True # two immediate flanking residues
        t1d_coord_mask[:,splice[1],:] = True
        
        t1d_conf_mask = torch.full_like(t1d, 1) #set confidences to 1

        t2d_mask[:,:,splice[0]-1, 37:] = False #also masks last plane of last dimension
        t2d_mask[:,:,splice[1], 37:] = False
        t2d_mask[:,splice[0]-1, :, 37:] = False
        t2d_mask[:,splice[1], :, 37:] = False
        
        #loss masks
        randtemp = random.randint(0,10)
        
        token_mask = msa_mask_input.clone()
        token_mask[:,:,splice[0]:splice[1]] = True # central only
        bond_mask[:,:,splice[0]-flank_width:splice[0]] = False 
        bond_mask[:,:,splice[1]:splice[1]+flank_width] = False # flanks only
        dih_mask = lddt_mask = dih_mask * bond_mask[:,0,:] # flanks only
        c6d_mask[:,:,~dih_mask[0,:],:] = False
        c6d_mask[:,:,:,~dih_mask[0,:]] = False # flanks only
        
        str_out_mask[splice[0]-flank_width:splice[0]] = False
        str_out_mask[splice[1]:splice[1]+flank_width] = False #score on everything but flanks
    
    elif task == 'str2seq':
        '''
        This is Joe's str2seq task, where a contiguous region is masked, along with flanks.
        Everything, but the flanked regions, is scored
        ### This may need to change to upweight the central region, by only scoring on the central region.
        '''
        ss_idx = get_ss(ss)
        splice, flank_width = get_masks(ss_idx, hal_mask_low, hal_mask_high, flank_low, flank_high)
        hal_mask_len = splice[1]-splice[0] 
        
        #input masks
        msa_masktok_mask[:,:,splice[0]-flank_width:splice[1]+flank_width] = False # flanks and central
        
        str_mask[:,splice[0]-flank_width:splice[0],:,:] = False # flanks only
        str_mask[:,splice[1]:splice[1]+flank_width,:,:] = False
        
        t1d_mask[:,splice[0]-flank_width:splice[1]+flank_width,:] = False #flanks and central
        
        t1d_conf_mask = torch.full_like(t1d, 1) #set confidences to 1

        #loss masks
        token_mask = msa_mask_input.clone()
        token_mask[:,:,splice[0]:splice[1]] = True # central only
        
        bond_mask[:,:,splice[0]-flank_width:splice[0]] = False 
        bond_mask[:,:,splice[1]:splice[1]+flank_width] = False # flanks only

        dih_mask = lddt_mask = dih_mask * bond_mask[:,0,:] # flanks only
        c6d_mask[:,:,~dih_mask[0,:],:] = False
        c6d_mask[:,:,:,~dih_mask[0,:]] = False # flanks only
        
        str_out_mask[splice[0]-flank_width:splice[0]] = False
        str_out_mask[splice[1]:splice[1]+flank_width] = False #score on everything but flanks
    
    elif task == 'single_ss_hal':
        '''
        Here, a random amino acid is selected from each secondary structure element > 4 residues in length and unmasked.
        Angles of these residues ARE masked, so network just sees points in space.
        Scored on everything
        '''
        ss_idx = get_ss(ss)
        unmasked_coords = get_rand_ss_idx(ss_idx)
        #input masks
        msa_mask = torch.zeros((B,N,L), device=msa.device).bool() #mask out whole sequence
        msa_masktok_mask[:,:,:] = False      
        str_mask = ~str_mask
        str_mask[:,unmasked_coords,:,:] = True # give position of one amino acid in each ss element
        
        t1d_mask = ~t1d_mask
        t1d_mask[:,unmasked_coords,:] = True #### need to check what 1d features are
        
        t1d_conf_mask = torch.full_like(t1d, 1) #set confidences to 1
        
        
        t2d_mask[:,:,unmasked_coords, 37:] = False # mask out torsion angles
        t2d_mask[:,unmasked_coords, :, 37:] = False
        
        #loss masks
        token_mask = torch.full_like(msa_mask,   True).bool()
    
    elif task == 'str2seq_full':
        '''
        This is David's str2seq task, where most (default 90-100%) of sequence is masked.
        Score on everything.
        '''
        rand_prop = random.uniform(str2seqfull_low, str2seqfull_high) #random proportion masked, between two extremes  
        temp = torch.rand(msa_masktok_mask.shape, device=msa.device) < rand_prop #make mask
        if constant_recycle is True:
            temp = temp[:1,:,:].repeat(B,1,1)
        msa_masktok_mask[temp] = False
        temp2=temp.squeeze()
        t1d_mask[temp2] = False
        #loss masks
        token_mask = torch.full_like(msa_mask_input,   False).bool()
        token_mask[temp] = True

    elif task == 'packing_hal':
        '''
        This is a new task, where two regions are masked out so the network has to learn to do two-sided packing.
        '''
        #### v2, splices proximal parts of the protein
        
        splice, length1, length2 = get_double_masks_proximal(xyz_t,min_length=hal_mask_low,max_length=hal_mask_high,min_gap=20)
 
        #input masks
        msa_masktok_mask[:,:,splice[0]:splice[0]+length1] = False
        msa_masktok_mask[:,:,splice[1]:splice[1]+length2] = False

        str_mask[:,splice[0]:splice[0]+length1,:,:] = False
        str_mask[:,splice[1]:splice[1]+length2,:,:] = False

        t1d_mask[:,splice[0]:splice[0]+length1,:] = False
        t1d_mask[:,splice[1]:splice[1]+length2,:] = False
        
        t1d_conf_mask = torch.full_like(t1d, 1) #set confidences to 1
        #loss masks

        token_mask = msa_mask_input.clone()
        token_mask[:,:,splice[0]:splice[0]+length1] = True
        token_mask[:,:,splice[1]:splice[1]+length2] = True


    elif task == 'packing_termini':
        '''
        This is a new task, which masks the N and C termini, in the hope this will help the model pack a protein, without positional information about the termini.
        '''
        #choose length of N and C terminus mask
        split = int(np.round(L/2,0))
        if split < 10:
            term_N_len = 0
        elif L-split < 10:
            term_C_len = 0
        else:
            maxlen=min(40,split-10)
            term_N_len = random.randint(10,maxlen+1)
            maxlen=min(40,L-split-10)
            term_C_len = random.randint(10,maxlen+1)

        #input masks
        msa_masktok_mask[:,:,:term_N_len] = False
        msa_masktok_mask[:,:,L-term_C_len-1:] = False

        str_mask[:,:term_N_len,:,:] = False
        str_mask[:,L-term_C_len-1:,:,:] = False

        t1d_mask[:,:term_N_len,:] = False
        t1d_mask[:,L-term_C_len-1:,:] = False

        t1d_conf_mask = torch.full_like(t1d, 1) #set confidences to 1

        #loss masks

        token_mask = msa_mask_input.clone()
        token_mask[:,:,:term_N_len] = True
        token_mask[:,:,L-term_C_len-1:] = True

    elif task == 'msa_hal':
        '''
        This is Joe's hallucination task, where a contiguous region is masked, along with flanks, and two residues either end are given (but not their angles).
        This is adapted for msa inputs.
        The two floating points are randomly moved some distance.
        Scored on everything but the flank regions (may want to change this to only score central inpainted region
        '''
        ss_idx = get_ss(ss)
        splice, flank_width = get_masks(ss_idx, hal_mask_low, hal_mask_high, flank_low, flank_high)
        hal_mask_len = splice[1]-splice[0]
        #input masks
        msa_masktok_mask[:,:,splice[0]-flank_width:splice[1]+flank_width] = False # flanks and central

        str_mask[:,splice[0]-flank_width:splice[1]+flank_width,:,:] = False # flanks and central
        str_mask[:,splice[0]-1,:,:] = True # immediate two flanking residues
        str_mask[:,splice[1],:,:] = True

        translate_mask[splice[0]-1] = False #imeediate two flanking residues
        translate_mask[splice[1]] = False

        t1d_mask[:,splice[0]-flank_width:splice[1]+flank_width,:] = False #flanks and central
        # make mask to add to t1d to give 'fake' confidence in two floating coordinates
        t1d_coord_mask[:,splice[0]-1,:] = True # two immediate flanking residues
        t1d_coord_mask[:,splice[1],:] = True

        t1d_conf_mask = torch.full_like(t1d, 1) #set confidences to 1

        t2d_mask[:,:,splice[0]-1, 37:] = False #also masks last plane of last dimension
        t2d_mask[:,:,splice[1], 37:] = False
        t2d_mask[:,splice[0]-1, :, 37:] = False
        t2d_mask[:,splice[1], :, 37:] = False

        #loss masks
        randtemp = random.randint(0,10)

        token_mask = msa_mask_input.clone()
        token_mask[:,:,splice[0]:splice[1]] = True # central only
        bond_mask[:,:,splice[0]-flank_width:splice[0]] = False
        bond_mask[:,:,splice[1]:splice[1]+flank_width] = False # flanks only
        dih_mask = lddt_mask = dih_mask * bond_mask[0,:,:] # flanks only. This line differs from 'hal' task masks
        c6d_mask[:,:,~dih_mask[0,:],:] = False
        c6d_mask[:,:,:,~dih_mask[0,:]] = False # flanks only

        str_out_mask[splice[0]-flank_width:splice[0]] = False
        str_out_mask[splice[1]:splice[1]+flank_width] = False #score on everything but flanks

    elif task in ['stub_fixedbb', 'stub_centroid', 'stub_fixedseq']:
        '''
        This is the stub design task for a fixed backbone.
        '''
        ic(mask_seq)
        ic(mask_seq.dtype)
        token_mask = msa_mask_input.clone()
        if task in ['stub_fixedbb', 'stub_centroid']:
            msa_masktok_mask[:,:,mask_seq] = False
            t1d_mask[:,mask_seq,:] = False
            token_mask[:,:,mask_seq] = True
    
        if task in ['stub_fixedseq', 'stub_centroid']:
            '''
            This is the stub design task where only the centroid of the stub is given.
            '''
            str_mask[:,mask_seq,:,:] = False
            # N, L,res,coord = xyz_t.size()
            ic(xyz_t.shape)
            stub_start = torch.min(mask_seq)
            stub_len = len(mask_seq)
            mask_seq_continguous = torch.arange(stub_start, stub_start+stub_len, dtype=mask_seq.dtype)
            torch.testing.assert_close(mask_seq, mask_seq_continguous)
            stub_center_idx = mask_seq_continguous[stub_len // 2]
            # C_BETA = 2
            # stub_centroid = xyz_t[:,C_BETA,:]
            # stub_centroid = stub_crop.c_alpha_centroid(xyz_t[0, mask_seq])
            # CENTROID_NOISE = 0
            # stub_centroid += stub_crop.spherical_sample(CENTROID_NOISE)
            # xyz_t[0, stub_center_idx, stub_crop.C_ALPHA_INDEX] = stub_centroid
            str_mask[:, stub_center_idx, :] = True
            # ic('str_mask idxs:', (~str_mask).nonzero())

    else:
        sys.exit(f'Masks cannot be generated for the {task} task!')
    
    #horrible logic, but mask2s mask out the whole column (one hot encodings and indel rows in msa or msa_full), and then masks bring back the mask token in planes 21 and 42.

    msa_masked_mask = torch.ones(B,N,L,46).bool()
    msa_masked_mask2 = msa_masktok_mask.unsqueeze(-1)
    msa_masked_mask2 = msa_masked_mask2.repeat(1,1,1,46)
    msa_masked_mask[:,:,:,20] = msa_masktok_mask
    msa_masked_mask[:,:,:,42] = msa_masktok_mask
    
    #make msa_full_mask compatible with all tasks
    full_N = msa_full.size()[1]
    msa_full_mask2 = msa_masked_mask2[:,:1,:,:23].repeat(1,full_N,1,1)
    msa_full_mask = msa_masked_mask[:,:1,:,:23].repeat(1,full_N,1,1)
    
    
    mask_dict = {'t1d_mask':t1d_mask,
                't1d_coord_mask':t1d_coord_mask,
                't1d_conf_mask':t1d_conf_mask,
                't2d_mask':t2d_mask,
                'str_mask':str_mask,
                'msa_masktok_mask':msa_masktok_mask,
                'msa_masked_mask':msa_masked_mask,
                'msa_masked_mask2':msa_masked_mask2,
                'msa_full_mask':msa_full_mask,
                'msa_full_mask2':msa_full_mask2,
                'c6d_mask':c6d_mask,
                'lddt_mask':lddt_mask,
                'dih_mask':dih_mask,
                'bond_mask':bond_mask,
                'token_mask':token_mask,
                'translate_mask':translate_mask,
                'str_out_mask':str_out_mask
                }
    
    return mask_dict
