import torch


def convergence_loop(query_update_func, P, queries, W_k=None):
    """
    Ensures that each query in the batch has converged or hits the 
    max number of convergence steps allowed. 
    
    args::
    query_update_func - partial function that has all parameters specific to each algorithm already set. 
    It just needs the query as input. 
    
    P - sdm_params
    queries - vectors in the matrix. 
    W_k - if given used to project the queries in their latent space. 
    """
    
    bits_changed_sum = 100000 # some large number
    i = 0
    # will keep going as long as none of them are satisfied. 
    while bits_changed_sum > 0 and i<P.max_conv_steps: 
        
        '''plt.imshow(queries[:,0].cpu().numpy().reshape(28,28))
        plt.show()'''
        
        if W_k is not None: 
            pQs = W_k@queries
            # L2 norm all
            pQs = pQs/torch.norm(pQs, dim=0, keepdim=True)
            new_queries = query_update_func(pQs, full_dim_queries=queries)
        else: 
            new_queries = query_update_func(queries)
        if P.binary_data:
            bits_changed_sum = (new_queries!=queries).sum()
        else: 
            # (P.n_base*P.batch_size)-
            bits_changed_sum = torch.logical_or(queries-P.epsilon > new_queries , 
                             queries+P.epsilon < new_queries).sum()
        queries = new_queries
        if not P.binary_data and P.renorm_cont_queries:
            # renorming the query for all Attention operations here.
            # dont want to do this have learnt a projection matrix as 
            # returned query here is in image space not L2 norm latent space. 
            queries = queries/torch.norm(queries, dim=0, keepdim=True)
        i+=1
    return queries