"""
Module with functions used in parrot-preprocess

.............................................................................
idptools-parrot was developed by the Holehouse lab
     Original release ---- 2020

Question/comments/concerns? Raise an issue on github:
https://github.com/idptools/parrot

Licensed under the MIT license. 
"""


# .................................................................
#
#
def are_lists_identical(l1,l2):
    """
    Quick function that asks if two lists are identical or not.

    Parameters
    -----------
    l1 : list
        First list

    l2 : list
        Second list

    Returns
    ----------
    bool
        Returns True if they are identical, or False if not
    
    """
    if len(l1) != len(l2):
        return False
    for i in range(len(l1)):
        if l1[i] != l2[i]:
            return False
    return True


# .................................................................
#
#
def recode_classes(lines):
    """
    Function which takes a set of data read in (lines) and converts it into a PARROT-compliant set of 
    lines files where classes start at 0. This works for both residue and sequence datatypes.

    Parameters
    -----------
    lines : list
        The lines list is a list of lists, where each sublist corresponds to a line in a parsed
        datafile. If the dataype is sequence then each sublist has three elements (id, sequence, class/score),
        while if the datatype is residues then each sublists has n elements (id, sequence, per-residue scores or 
        classes for each residue).


    Returns
    -------------
    list
        Returns a revised lines list, where sequence or residue classes are guarenteed to start at 0 and increase
        non-monnotonically.
    
    """
    
    # initialize empty vars
    classes = set([])
    class_mapping = {}


    # first figure out the number of classes
    for line in lines:
        classes = classes.union(set(line[2:]))

    n_classes = len(classes)

    if n_classes > 100:
        print('WARNING: 100 individual classes found. This is probably not a good classification problem...')

    # next build a dictionary that maps each unique class (sorted) back to a monotonically increasing
    # number of class labels starting at 0
    classes = list(classes)
    classes.sort()

    idx = 0
    for i in classes:
        class_mapping[i] = idx
        idx = idx + 1

    # Finally, build a new lines list, where the sequence or residues classes are mapped
    # via the class_mapping dictionary 
    new_lines = []
    for line in lines:

        # convert the class
        tmp = [line[0], line[1]]       
        tmp.extend([str(class_mapping[i]) for i in line[2:]])

        new_lines.append(tmp)
        
    return (new_lines, class_mapping)



# .................................................................
#
#
def remove_sequence_duplicates(lines, verbose=False):
    """
    Function that removes any line where the same sequence appears twice. Because in such
    and instance we cannot know which of the two sequences is the 'ground truth', this method
    removes both sequences from the dataset

    Parameters
    -------------
    lines : list
        Lines list file generated by parrot.process_input_data.__read_tsv(<filename>).

    verbose : bool
        Flag which tells the function how loud to be. Default = False.

    Returns
    ------------
    list
        Returns an updated lines file that has had all the lines with the same sequence
        removed - i.e. the resulting file is guarenteed to have unique sequences, although
        not necessarily unique IDs or datasets.

    """
    if verbose:
        print('-------------------------------------------------------------------')
        print('Performing duplicate sequence removal --remove-sequence-duplicates ')
        print('-------------------------------------------------------------------')


    seen = {}
    nonuniq = set([])


    duplicate_count = 0
    for line in lines:

        if line[1] not in seen:
            seen[line[1]] = line
        else:
            duplicate_count = duplicate_count + 1
            nonuniq.add(line[1])

    
    for x in nonuniq:
        seen.pop(x)

    new_lines = []
    for x in seen:
        new_lines.append(seen[x])

    if verbose:
        print(f'   Found {duplicate_count} sequence duplicates in this file. All copies of these sequences have been removed.\n')
        print(f'   NOTE this may be more aggressive than you intended - to only remove one of the duplicate copies of duplicate lines, use the --remove-line-dup flag instead')
        

        print(f'   Also found {seq_duplicate_count} ID-sequence duplicates in this file (same ID, same sequence, different data.\n   If any of these were full duplicates (likely) they will have been removed.\n')
        print(f'   Found {seq_duplicate_count_full} sequence duplicates in this file (same sequence with different ID and different data).\n   If any of these were full duplicates (likely) they will have been removed.\n')

    if verbose:
        print(f'Found {duplicate_count} sequence duplicates. These have been removed.')        

    return new_lines
                


# .................................................................
#
#
def remove_line_duplicates(lines, verbose=False):
    """
    Function that removes true duplicates and warns about potential psuedoduplicates if
    verbose is set to true.

    Parameters
    -------------
    lines : list
        Lines list file generated by parrot.process_input_data.__read_tsv(<filename>).

    verbose : bool
        Flag which tells the function how loud to be. Default = False.

    Returns
    ------------
    list
        Returns an updated lines file that has had all the true duplicate lines removed (i.e.
        lines where the id, sequence, and data were all identical. 
    

    """

    if verbose:
        print('---------------------------------------------------------------')
        print('Performing duplicate checking/removal --remove-line-duplicates ')
        print('---------------------------------------------------------------')

    found = {}
    all_sequences = set([])
    duplicate_count = 0
    seq_duplicate_count = 0
    seq_duplicate_count_full = 0

    # for each line
    for i in lines:

        # this part ONLY keeps track of sequences. Note these are not removed
        # because which one should you remove!
        if i[1] in all_sequences:
            seq_duplicate_count_full = seq_duplicate_count_full + 1
        else:
            all_sequences.add(i[1])

        # if the ID code is has not yet been seen
        if i[0] not in found:
            found[i[0]] = [i]

        # else if the ID code (first column in datafile) HAS been seen
        else:
            
            # set the unique flag to true
            this_is_a_duplicate = False

            # for EVERY time it was seen (hits is a previous line)
            for hits in found[i[0]]:

                # if the sequence ALSO MATCHED
                if hits[1] == i[1]:
                    seq_duplicate_count = seq_duplicate_count + 1
                    
                    if are_lists_identical(hits[2:], i[2:]):
                        duplicate_count = duplicate_count + 1
                        #
                        this_is_a_duplicate = True

            # if we get here then this was 
            if this_is_a_duplicate is False:
                found[i[0]].append(i)

    if verbose:
        print(f'   Found {duplicate_count} full duplicate lines in this file. These have been removed.\n')
        print(f'   Also found {seq_duplicate_count} ID-sequence duplicates in this file (same ID, same sequence, different data.\n   If any of these were full duplicates (likely) they will have been removed.\n')
        print(f'   Found {seq_duplicate_count_full} sequence duplicates in this file (same sequence with different ID and different data).\n   If any of these were full duplicates (likely) they will have been removed.\n')

    new_lines = []
    for id_key in found:
        for x in found[id_key]:
            new_lines.append(x)

    return new_lines
        
                




