%load_ext autoreload
%autoreload 2

import numpy as np
import scipy as sp
import pandas as pd

import dill
from pprint import pprint

with open('resource/asnlib/publicdata/acm-v9-sample-100k.dill', 'rb') as fp:
    raw_records = dill.load(fp)

print(f"=== Success: Loaded {len(raw_records):,} raw records. ===")
print(f"\nExample: Records 0, 7, and 15:\n")
pprint([raw_records[k] for k in [0, 7, 15]])

with open('resource/asnlib/publicdata/STOPWORDS.dill', 'rb') as fp:
    STOPWORDS = dill.load(fp)

print("A sample of stopwords:")
sorted(list(STOPWORDS))[:7]

### Solution - Exercise 0  

def get_record(lines: list) -> dict:
    
    # INPUT: lines is a list of strings
    
    # GOAL: Return 'record', a dictionary of lists. Convert prefix at beginning of each string in lines to correct key. 
    # List values should be sorted.
    
    # STRATEGY:
    # 1. Initialize record dictionary to hold our results
    # 2. Create a mappings dictionary to map each prefix to its corresponding key
    # 3. Iterate over prefixes
        # 4. Iterate over lines
            # 5. If the raw attribute string in lines starts with the prefix, add that raw attribute to a list, 
            #    getting rid of newline char at the end
        # 6. After iterating over all lines, if the list for that prefix is not empty, sort the list and add it to our record dictionary with the correct key
    # 7. Return record dictionary     
    
    
    # SOLUTION:
    record = {}
    mappings_dict = {'#*': 'title', '#@': 'authors', '#t': 'year', '#c': 'venue', '#index': 'id', '#%': 'refs', '#!': 'abstract'}
    
    for prefix in mappings_dict.keys():
        list_of_attributes = []
        for raw_attribute in lines:
            if raw_attribute.startswith(prefix):
            # Alternative: instead of using .startswith()
            #if prefix == raw_attribute[:len(prefix)]:
                list_of_attributes.append(raw_attribute[len(prefix):].rstrip())
        if len(list_of_attributes) > 0:
            record[mappings_dict[prefix]] = sorted(list_of_attributes)
    
    return record

### Demo function call

pprint(get_record(raw_records[15]))

### Test Cell - Exercise 0  

from cse6040_devkit.tester_fw.testers import Tester
from yaml import safe_load

with open('resource/asnlib/publicdata/assignment_config.yaml') as f:
    ex_conf = safe_load(f)['exercises']['get_record']['config']

ex_conf['func'] = get_record

tester = Tester(ex_conf, key=b'4oTuD3jYU_dfQvrxIzSmKJQBDeiWMaODi2nmk0sQk1o=', path='resource/asnlib/publicdata/')
for _ in range(75):
    try:
        tester.run_test()
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
    except:
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
        raise

###
### AUTOGRADER TEST - DO NOT REMOVE
###

print('Passed! Please submit.')

with open('resource/asnlib/publicdata/records.dill', 'rb') as fp:
    records = dill.load(fp)

print(f"Sample post-processed records:")
for k in [0, 7, 15]:
    print(f"\n=== get_record(raw_records[{k}]) ===")
    pprint(records[k])

### Solution - Exercise 1  

def clean_record(record: dict) -> dict:
    
    # INPUT: record is a dictionary of lists
    
    # GOAL: Return cleaned dictionary. 
    # 1. If the input is missing any of the following keys, return None: ['id', 'title', 'year', 'venue', 'abstract'].
    # 2. For the 'id' and 'year' keys, assume there is only one value in the list and convert it into an integer 
    # 3. If there is a 'refs' key, convert each element of its list to an integer and store this list of integers.
    # 4. For any other keys that are present in record, assume there is only one value and store that value as-is 
    
    # STRATEGY:
    # 1. First check to make sure all required keys exist in 'record'
    # 2. Initialize empty dictionary 'cleaned' to hold our results
    # 3. Iterate over key/value pairs in record dictionary
            # 4. If key is 'id' or 'year', convert 0th entry in list to an integer and store in 'cleaned'
            # 5. If key is 'refs', convert each element of list to integer and store list in 'cleaned'
            # 6. Anything else, store 0th entry in list in 'cleaned'
    # 7. Return 'cleaned' dictionary
        
    # SOLUTION:
    required_keys = ['id', 'title', 'year', 'venue', 'abstract']
    
    for key in required_keys:
        if key not in record:
            return None
    
    cleaned = {}
    
    for key, val in record.items():
        if key == 'id' or key == 'year':
            cleaned[key] = int(val[0])
        elif key == 'refs':
            cleaned[key] = [int(e) for e in val]
        else:
            cleaned[key] = val[0]
    
    return cleaned

### Demo function call

pprint(clean_record(records[0])) # None!
pprint(clean_record(records[17])) # Valid

### Test Cell - Exercise 1  

from cse6040_devkit.tester_fw.testers import Tester
from yaml import safe_load

with open('resource/asnlib/publicdata/assignment_config.yaml') as f:
    ex_conf = safe_load(f)['exercises']['clean_record']['config']

ex_conf['func'] = clean_record

tester = Tester(ex_conf, key=b'BhZT_oWZZqOjWxKNFXA5A7uL44XtRE6X7vsh1IcmmOs=', path='resource/asnlib/publicdata/')
for _ in range(75):
    try:
        tester.run_test()
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
    except:
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
        raise

###
### AUTOGRADER TEST - DO NOT REMOVE
###

print('Passed! Please submit.')

with open('resource/asnlib/publicdata/papers.dill', 'rb') as fp:
    cleaned_records = dill.load(fp)

print("Examples of cleaned records:")
for k in [0, len(cleaned_records)//2, len(cleaned_records)-1]:
    record = cleaned_records[k]
    print(f"\n=== Record {k} ===")
    pprint(record)

### Solution - Exercise 2  

def records_to_metadf(records: list) -> pd.DataFrame:
    
    # INPUT: records is a list of cleaned records (dictionaries)
    
    # GOAL: Return metadf data frame with columns: id, title, year, venue, abstract
    
    # STRATEGY: 
    # 1. Create metadf dataframe from 'records' (list of dictionaries)
    # 2. Filter metadf down to the columns we care about ('id', 'title', 'year', 'venue', 'abstract')
    # 3. Return   
    
    # SOLUTION:
    # See section 'Create DataFrame from List of Dictionaries'
    # https://www.geeksforgeeks.org/different-ways-to-create-pandas-dataframe/
    
    metadf = pd.DataFrame(records)
    columns_we_want = ['id', 'title', 'year', 'venue', 'abstract']
    
    return metadf[columns_we_want]

### Demo function call

demo_cleaned_records = [cleaned_records[k] for k in [11222, 11239, 12329]]
demo_metadf = records_to_metadf(demo_cleaned_records)
display(demo_metadf)

### Test Cell - Exercise 2  

from cse6040_devkit.tester_fw.testers import Tester
from yaml import safe_load

with open('resource/asnlib/publicdata/assignment_config.yaml') as f:
    ex_conf = safe_load(f)['exercises']['records_to_metadf']['config']

ex_conf['func'] = records_to_metadf

tester = Tester(ex_conf, key=b'hgW1Cvvo9Kg4sTImuGlhMic09mEdkjjQQFPMCiw5gaE=', path='resource/asnlib/publicdata/')
for _ in range(75):
    try:
        tester.run_test()
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
    except:
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
        raise

###
### AUTOGRADER TEST - DO NOT REMOVE
###

print('Passed! Please submit.')

with open('resource/asnlib/publicdata/metadf.dill', 'rb') as fp:
    metadf = dill.load(fp)

# For Ex. 3, break dependence with a working Ex. 2
demo_metadf = metadf.loc[[11222, 11239, 12329]].reset_index(drop=True)

print("Examples of metadata records:")
display(metadf.sample(5))

### Solution - Exercise 3  

def gen_corpus(metadf, stopwords):
    
    # INPUT: 
    # 1. metadf - a data frame 
    # 2. stopwords - a set of words we should exclude
    
    # GOAL:
    # Return 'pseudodf' a data frame with columns:
    #     'id': The document ID, taken as-is from the input, metadf['id']
    #     'title': The document title, taken as-is from metadf['title']
    #     'pseudodoc': A "pseudo-document," which is a string built from the 'title', 'venue', and 'abstract' as follows:
                # 1. get unique tokens from lowercase of title, venue, and abstract (consecutive sequence of a-z only)
                # 2. remove stopwords
                # 3. concatenate only the unique tokens, in sorted order, and separated by spaces
    
    # STRATEGY:
    # 1. Make a copy of metadf. Call this 'df'
    # 2. Create an empty list to hold our pseudodoc column data
    # 3. Iterate over rows in df
            # 4. Combine lowercase of columns 'title', 'venue', and 'abstract' (separate by space) into one big string
            # 5. Find list of words containing a-z characters
            # 6. Find unique words in Step 5 list (use set?)
            # 7. Remove all stopwords from Step 6 and sort
            # 8. Join Step 7 list of words back together into one string using single space character
            # 9. Add to 'pseudodoc' list
    # 10. Add 'pseudodoc' list as a column to our df
    # 11. Filter df to columns we care about (id, title, pseudodoc) and return 
    
    # SOLUTION:
    import re
    
    df = metadf.copy()
    pseudodoc = []
    
    for index, row in df.iterrows():
        all_words_str = row['title'].lower() + ' ' + row['venue'].lower() + ' ' + row['abstract'].lower()
        word_list = re.findall('[a-z]+', all_words_str)
        
        # Alternative way if you don't want to use regex:
#         new_words_str = ''
#         for char in all_words_str:
#             if char in 'abcdefghijklmnopqrstuvwxyz':
#                 new_words_str += char
#             else:
#                 new_words_str += ' '
        
#         word_list = new_words_str.split()

        unique_words = set(word_list)
        final_list_of_words = sorted(unique_words.difference(stopwords))
        pseudodoc.append(' '.join(final_list_of_words))
    
    df['pseudodoc'] = pseudodoc
    return df[['id', 'title', 'pseudodoc']]
    


### Demo function call
demo_corpus = gen_corpus(demo_metadf, STOPWORDS)
display(demo_corpus)

### Test Cell - Exercise 3  

from cse6040_devkit.tester_fw.testers import Tester
from yaml import safe_load

with open('resource/asnlib/publicdata/assignment_config.yaml') as f:
    ex_conf = safe_load(f)['exercises']['gen_corpus']['config']

ex_conf['func'] = gen_corpus

tester = Tester(ex_conf, key=b'o0kicBPd2HinEbuciVMQHNjfCLJj4eVAWPBqKCfSeNw=', path='resource/asnlib/publicdata/')
for _ in range(75):
    try:
        tester.run_test()
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
    except:
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
        raise

###
### AUTOGRADER TEST - DO NOT REMOVE
###

print('Passed! Please submit.')

with open('resource/asnlib/publicdata/corpus.dill', 'rb') as fp:
    corpusdf = dill.load(fp)

print(f"A sample of the pseudo-document corpus:")
display(corpusdf.sample(4))

with open('resource/asnlib/publicdata/X.dill', 'rb') as fp:
    X = dill.load(fp)
    
print(f"The shape of `X`: {X.shape}")
print(f"- Smallest value: {X.min()}")
print(f"- Largest value: {X.max()}")
print(f"- Mean value: {X.mean()}")

with open('resource/asnlib/publicdata/kmeans.dill', 'rb') as fp:
    _, km_sizes, km_centers, km_labels = dill.load(fp)

print(f"- Cluster sizes:", km_sizes)
print(f"- Cluster labels:", km_labels)
print(f"- Cluster centers, {km_centers.shape[0]:,} x {km_centers.shape[1]}, where each column is a centroid:")
pprint(km_centers)

### Solution - Exercise 4  

def build_cluster_matrix(labels: np.ndarray, max_label: int) -> sp.sparse.coo_matrix:
    
    # INPUT: 
    # 1. labels: A 1-D array where labels[i] is the cluster ID assigned to document i
    # 2. max_label: The largest cluster ID plus 1. That is, 0 <= labels[i] < max_label for all i

    # GOAL: Return Scipy coo matrix, with len(labels) rows and max_label columns. 
    # It should only have entries equal to 1.0 at positions (i, j) where j == labels[i]; 
    # all other entries should be regarded as zero values and, therefore, not stored explicitly.

    # STRATEGY:
    # 1. Create a 2-D array of zeros with len(labels) rows and max_label columns
    # 2. Iterate over labels and set the corresponding i, j locations to 1.0
    # 3. Convert to a coo matrix and return
    
    # SOLUTION:   
    full_arr = np.zeros((len(labels), max_label))
    
    for i, j in enumerate(labels):
        full_arr[i,j] = 1.0
    
    return sp.sparse.coo_matrix(full_arr)

    # -------------------------------------------------------------
    
    # Alternative solution building out data, row, and column arrays (instead of populating 1.0s in giant array):
    
    # Create coo matrix using this syntax: coo_matrix((data, (i, j)), shape=(M, N))
    # Construct from three arrays:
    # 1. data[:] the entries of the matrix, in any order
    # 2. i[:] the row indices of the matrix entries
    # 3. j[:] the column indices of the matrix entries
    
    data = []
    i_list = []
    j_list = []
    
    for i, j in enumerate(labels):
        data.append(1.0)
        i_list.append(i)
        j_list.append(j)
    
    return sp.sparse.coo_matrix((np.array(data), (np.array(i_list), np.array(j_list))), shape=(len(labels), max_label))



### Demo function call
demo_R = build_cluster_matrix(np.array([1, 2, 0, 0, 2, 2, 2, 3, 0, 3]), 4)
print(demo_R)

### Test Cell - Exercise 4  

from cse6040_devkit.tester_fw.testers import Tester
from yaml import safe_load

with open('resource/asnlib/publicdata/assignment_config.yaml') as f:
    ex_conf = safe_load(f)['exercises']['build_cluster_matrix']['config']

ex_conf['func'] = build_cluster_matrix

tester = Tester(ex_conf, key=b'cub7HYh_bdtYjSXeD5UPThVliqosSf9XWu4LLxZjlOs=', path='resource/asnlib/publicdata/')
for _ in range(75):
    try:
        tester.run_test()
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
    except:
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
        raise

###
### AUTOGRADER TEST - DO NOT REMOVE
###

print('Passed! Please submit.')

with open('resource/asnlib/publicdata/R.dill', 'rb') as fp:
    R_km = dill.load(fp)
    
R_km

### Solution - Exercise 5  

def get_top_tokens(cid: int, labels: np.ndarray, corpusdf: pd.DataFrame, k=10) -> set:
    
    # INPUT: 
    # 1. cid: The ID of a cluster to analyze
    # 2. labels: 1-D array. Cluster assignments, where document i was assigned to cluster labels[i]
    # 3. corpusdf: A pseudo-document dataframe, having the columns 'id', 'title', and 'pseudodoc'
    # 4. k: integer - the number of tokens to return
    
    # GOAL:
    # Return top_tokens, a Python set of the k most frequent tokens
    
    # STRATEGY:
    # 1. Use labels to identify the documents belonging to cluster cid
    # 2. Select their corresponding pseudo-documents corpusdf. To match a document ID i to its pseudo-document, 
    #    note that i is the index value of corpusdf (and, in particular, not the 'id' column!)
    # 3. For each unique token, count the number of pseudo-documents in which it appeared.
    # 4. Return up to the k most frequently occurring tokens as a Python set (so we need to sort).
    #    In the case of ties, consider tokens in ascending order.
    
    # SOLUTION:
    # Create a big list of all of the tokens from documents associated with our desired cluster id cid:
    token_list = []
    
    for i, cluster_id in enumerate(labels):
        if cluster_id == cid:
            pseudodoc = corpusdf['pseudodoc'].loc[i]
            token_list += pseudodoc.split()
    
    # Now count occurrences of all the tokens in token_list:
    from collections import Counter
    
    counter_dict = Counter(token_list)
    
    # Sort by count value in descending order, then break ties by key value in ascending order:
    sorted_list_of_tuples = sorted(counter_dict.items(), key=lambda x: (-x[1], x[0]))
    
    # Grab top k elements and get only the keys (not the values):
    top_k_list_of_tuples = sorted_list_of_tuples[:k]
    top_k_list = [x[0] for x in top_k_list_of_tuples]
    
    # Return as a set:
    return set(top_k_list)

### Demo function call
with open('resource/asnlib/publicdata/demo_args_get_top_tokens.dill', 'rb') as fp:
    demo_cid, demo_labels, demo_corpusdf, demo_k = dill.load(fp)
print(get_top_tokens(demo_cid, demo_labels, demo_corpusdf, demo_k))

### Test Cell - Exercise 5  

from cse6040_devkit.tester_fw.testers import Tester
from yaml import safe_load

with open('resource/asnlib/publicdata/assignment_config.yaml') as f:
    ex_conf = safe_load(f)['exercises']['get_top_tokens']['config']

ex_conf['func'] = get_top_tokens

tester = Tester(ex_conf, key=b'BBpXHA04_pc-q0tx10aYqwXtRUfiJeKp90vRPTLzC5k=', path='resource/asnlib/publicdata/')
for _ in range(75):
    try:
        tester.run_test()
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
    except:
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
        raise

###
### AUTOGRADER TEST - DO NOT REMOVE
###

print('Passed! Please submit.')

with open('resource/asnlib/publicdata/top_tokens.dill', 'rb') as fp:
    km_top_tokens = dill.load(fp)
    
km_top_tokens

### Solution - Exercise 6  

def find_distinctive_tokens(token_sets: dict, max_occur=3) -> dict:
    # INPUT: 
    # 1. token_sets is a dictionary of sets, where token_sets[i] are the top tokens appearing in cluster i
    # 2. max_occur is integer, maximum number of allowable occurrences of a token across clusters for that token 
    #    to be considered "distinctive" of that cluster.
    
    # GOAL:
    # Return 'distinctive': A new dictionary mapping cluster key to a set of distinctive tokens
    # So we're filtering out the common tokens, and only keeping the distinctive ones
    
    # STRATEGY:
    # 1. Get a list of all of the tokens in the 'token_sets' dictionary (aka combine all of the sets into one big list)
    # 2. Count up occurrences of each token in token list created in Step 1. Let's create a Counter dictionary.
    # 3. Initialize a 'distinctive' default dictionary of sets to hold our new dictionary results
    # 4. Iterate over key, value pairs in 'token_sets' dictionary:
            # 5. Iterate over tokens in that key's set of values:
                    # 6. If token has less than or equal to max_occur in our Counter dictionary, he's good 
                    #    and we add him to our new 'distinctive' dictionary set for that key
    # 7. Return 'distinctive' dictionary
    
    # SOLUTION:
    from collections import Counter
    from collections import defaultdict
    
    list_of_tokens = [val for subset in token_sets.values() for val in subset]
    token_count_dict = Counter(list_of_tokens)
    
    distinctive = defaultdict(set)
    
    for cluster_id, set_of_tokens in token_sets.items():
        for token in set_of_tokens:
            if token_count_dict[token] <= max_occur:
                distinctive[cluster_id].add(token)
    
    return dict(distinctive)

### Demo function call
print(find_distinctive_tokens(km_top_tokens))

### Test Cell - Exercise 6  

from cse6040_devkit.tester_fw.testers import Tester
from yaml import safe_load

with open('resource/asnlib/publicdata/assignment_config.yaml') as f:
    ex_conf = safe_load(f)['exercises']['find_distinctive_tokens']['config']

ex_conf['func'] = find_distinctive_tokens

tester = Tester(ex_conf, key=b'5VeeLsiYTK9AUwJ8sN5jiO4xU3sWzBJooV29Q6R7Erk=', path='resource/asnlib/publicdata/')
for _ in range(75):
    try:
        tester.run_test()
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
    except:
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
        raise

###
### AUTOGRADER TEST - DO NOT REMOVE
###

print('Passed! Please submit.')

### Solution - Exercise 7  

def lsa_svd_cluster__FREE(X: np.ndarray, num_clusters=10, num_components=10):
    """**FREE EXERCISE**: Use the SVD and K-means to compute a clustering.
This method is sometimes referred to as _latent semantic analysis_ (LSA).

**Background**: In the data-matrix factorization view of K-means, we compute $X \sim R C^T$ and use $R$ to identify clusters.
That suggests we could try _other_ factorization methods, including the singular value decomposition (SVD) from Notebook 15.

In this strategy, we compute an $d$-truncated SVD, $X \\approx U_d \Sigma_d V_d^T$, where we choose $d$ to be "small" compared to the number of columns of $X$.
In so doing, the matrix $Y \equiv U_d \Sigma_d$ acts as lower-dimensional representation of $X$, where each row of $Y$ is a $d$-dimensional version of the corresponding row of $X$.
We can then run K-means on $Y$.

Here, we are giving you some code to carry out this analysis.
This code uses scikit-learn's [TruncatedSVD](https://scikit-learn.org/stable/modules/generated/sklearn.decomposition.TruncatedSVD.html#) and [K-means](https://scikit-learn.org/stable/modules/generated/sklearn.cluster.KMeans.html#) implementations.
Inspect the results and then move on to the next (and last) part and exercise.

**Inputs**:
- `X`: An `n`-by-`m` data matrix where rows are data points.
- `num_clusters`: The desired number of clusters to identify.
- `num_components`: The truncated dimension, $d$.

**Return:** `labels`, an array of cluster labels.
The `i`-th row of `X` is "assigned" to a cluster whose ID is `labels[i]`.
These labels will lie in the range 0 to `num_clusters-1`.
    """
    # Step 1: Compute the (truncated) SVD of X: X ~ U S V^T
    from sklearn.decomposition import TruncatedSVD
    svd = TruncatedSVD(n_components=num_components, n_iter=20, random_state=1_234)
    Y = svd.fit_transform(X)
    
    # Step 2: Run K-means on U
    from sklearn.cluster import KMeans
    kmeans = KMeans(n_clusters=num_clusters, max_iter=100, n_init=1, random_state=5_678).fit(Y)
    labels = kmeans.predict(Y)
    return labels

### Demo function call
pass # bug? this line gets omitted in the output
with open('resource/asnlib/publicdata/X.dill', 'rb') as fp:
    X = dill.load(fp)
    print(f"Data matrix (`X`) shape: {X.shape[0]:,} x {X.shape[1]:,}")
lsa_labels = lsa_svd_cluster__FREE(X, num_clusters=10)
print(f"Cluster assignments: {lsa_labels}")

print(f"\nDistinctive top tokens for LSA:")
with open('resource/asnlib/publicdata/lsa_tokens.dill', 'rb') as fp:
    lsa_top_tokens, lsa_distinctive = dill.load(fp)
pprint(lsa_distinctive)

### Test Cell - Exercise 7  


print('Passed! Please submit.')

with open('resource/asnlib/publicdata/W.dill', 'rb') as fp:
    W = dill.load(fp)
    
print(W.shape)

### Solution - Exercise 8  

def get_nmf_clusters(W: np.ndarray) -> (np.ndarray, np.ndarray, np.ndarray):
    
    # INPUT: W, an n-by-m matrix produced by an NMF computation.
    
    # GOAL: Return 3 1-D arrays:
    # labels: Cluster assignments: for each row i of W, labels[i] is the cluster (column) assignment.
    # ids: Cluster IDs: a Numpy array containing the values from 0 to m-1.
    # sizes: Cluster sizes: sizes[c] is the number of rows of W that are assigned to cluster c.
    
    # STRATEGY:
    # First, build out 'labels':
    # 1. Find column locations of max values for each row in W
    
    # Now, let's build out 'ids':
    # 2. Create 1-D array with values ranging from 0 to m-1
    
    # Finally, build out 'sizes':
    # 3. For each cluster c, count how many rows were assigned to it and store that as sizes[c]
    
    # SOLUTION:
    # Build labels using Numpy argmax:
    # https://numpy.org/doc/stable/reference/generated/numpy.argmax.html
    labels = np.argmax(W, axis=1)
    
    # Build ids and sizes using Numpy unique
    # Google Search: numpy count occurrences of all values
    # https://www.tutorialspoint.com/how-to-count-the-frequency-of-unique-values-in-numpy-array
    ids, sizes = np.unique(labels, return_counts=True)
    
    return labels, ids, sizes
    
    # ---------------------------------------------------------
    
    # Alternative solution not using Numpy unique (using Counter instead):
    W_shape = np.shape(W)
    number_of_clusters = W_shape[1]
    
    ids = np.arange(number_of_clusters)
    labels = np.argmax(W, axis=1)
    
    from collections import Counter
    
    cluster_counter_dict = Counter(labels)
    sizes_list = []
    
    for i in range(number_of_clusters):
        sizes_list.append(cluster_counter_dict[i])
        
    sizes = np.array(sizes_list)
    
    return labels, ids, sizes
    


### Demo function call
with open('resource/asnlib/publicdata/demo_get_nmf_clusters_W.dill', 'rb') as fp:
    demo_W = dill.load(fp)
print("Demo input W:")
pprint(demo_W)
print("\nYour output:")
print(get_nmf_clusters(demo_W))

### Test Cell - Exercise 8  

from cse6040_devkit.tester_fw.testers import Tester
from yaml import safe_load

with open('resource/asnlib/publicdata/assignment_config.yaml') as f:
    ex_conf = safe_load(f)['exercises']['get_nmf_clusters']['config']

ex_conf['func'] = get_nmf_clusters

tester = Tester(ex_conf, key=b'SGeoIKLnOAbhAQCj3ymEM2J0nfjbk6CylhmBcXTw3s4=', path='resource/asnlib/publicdata/')
for _ in range(75):
    try:
        tester.run_test()
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
    except:
        (input_vars, original_input_vars, returned_output_vars, true_output_vars) = tester.get_test_vars()
        raise

###
### AUTOGRADER TEST - DO NOT REMOVE
###

print('Passed! Please submit.')