# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Mark and save training and holdout sets of snippets in the database
#
# Some of the sets are generated using gsdmm for clustering by topic (in two rounds).
#
# The sets are then written to disk for labeling via Mechanical Turk (Mturk)
#
# To run on demon; written in python 2

# +
### Mark a simple random holdout & training sets in the db
# (only needs be run once)
import pymongo

MONGODB_PORT = 25541

with pymongo.MongoClient('localhost', MONGODB_PORT) as mclient:
    mdb = mclient.snippetdb
    msnippets = mdb.snippets
    msnippets.create_index([('holdout', pymongo.ASCENDING)])

    # add a station-stratified random sample to the holdout set (40*13 = 520)
    #msnippets.update_many({"station_rand_idx": {"$lt": 40}}, { '$set': {"holdout": 1} })
    # add station-stratified random sample training set (in 4 batches of size 520 each)
    msnippets.update_many({"station_rand_idx": {"$gte": 40, "$lt": 80}}, { '$set': {"train_set": 0} })
    msnippets.update_many({"station_rand_idx": {"$gte": 80, "$lt": 120}}, { '$set': {"train_set": 1} }) 
    msnippets.update_many({"station_rand_idx": {"$gte": 120, "$lt": 160}}, { '$set': {"train_set": 2} })
    msnippets.update_many({"station_rand_idx": {"$gte": 160, "$lt": 200}}, { '$set': {"train_set": 3} })


# +
### Load clusters generated by GSDMM run and
### 0) consolidate custers, combining small & removing empty
### 1) mark a small holdout set from those
### 2) save the cluster ids to the database
### 3) mark training batches
###
# Note the clusters were formed from snippets having a station_rand_idx>200 (so not in any train_set yet)
import pymongo
from bson.objectid import ObjectId
import numpy
from collections import defaultdict
import pickle

MONGODB_PORT = 25541
NUM_HOLDOUTS_PER_CLUSTER = 7
NUM_TRAIN_PER_CLUSTER = 16

# load saved clusters from sampled data
output_stem = "snippet_topicmodel/snippet"
cluster_fn = output_stem + '_gsdmm_doc_clusters_2.pickle'

## TEMP: !!!
#cluster_fn = output_stem + '_gsdmm_doc_clusters.index_a0.1_b0.1_15its_K326.pickle'

with open(cluster_fn, 'r') as f:
    data = pickle.load(f)

# save the members of each cluster
cluster_members = defaultdict(list)
for idx, cluster_id in enumerate(data['doc_clusters']):
    cluster_members[cluster_id].append(idx)

nClustersBefore = len(cluster_members.keys())
    
# sort the indices in each cluster, and combine clusters having too few members
for cluster, idxs in cluster_members.items(): # we deliberately want a copy (so not iteritems())
    cluster_members[cluster] = sorted(idxs)
    if len(idxs) < NUM_HOLDOUTS_PER_CLUSTER + NUM_TRAIN_PER_CLUSTER:
        cluster_members[-1] += idxs
        del cluster_members[cluster]

nClustersAfter = len(cluster_members.keys())
print("Combined %d small clusters into 1, leaving %d of %d clusters." % (nClustersBefore - nClustersAfter, nClustersAfter, nClustersBefore))

# renumber the clusters to remove gaps; make the combined cluster be index 0.
cluster_members2 = {}
new_keys = range( len(cluster_members.keys()) )
for new_key, old_key in zip(new_keys, sorted(cluster_members.keys())):
    cluster_members2[new_key] = cluster_members[old_key]

cluster_members = cluster_members2
cluster_members2 = None
    
# map from idx to _id and insert cluster ids; mark holdouts and train data in mongodb
with pymongo.MongoClient('localhost', MONGODB_PORT) as mclient:
    mdb = mclient.snippetdb
    msnippets = mdb.snippets
    #msnippets.update_many({}, {"$unset": "train_set"})
    #msnippets.update_many({}, {"$unset": "recluster"})
    
    for cluster, idxs in cluster_members.iteritems():
        print("cluster %d of size %d\n" % (cluster, len(idxs)))
        for i, idx in enumerate(idxs):
            update_spec = {"cluster_2": cluster}
            
            # assign to holdout or train set as applicable
            if i < NUM_HOLDOUTS_PER_CLUSTER:
                update_spec["holdout"] = 3
            elif i < NUM_HOLDOUTS_PER_CLUSTER + NUM_TRAIN_PER_CLUSTER: 
                update_spec["train_set"] = 20 + (i-NUM_HOLDOUTS_PER_CLUSTER) / 8  # 10-12 w/ 3 docs per cluster
        
#             # label the large clusters for re-clustering (exclude what are mostly ads in small clusters)
#             if cluster != 0 and len(idxs) >= 100:
#                 update_spec["recluster"] = True
            
            query = {"_id": ObjectId(data['doc_ids'][idx])}
            
            msnippets.update_one(query, {"$set": update_spec})
            
#             if i > len(idxs) - 4:
#                 doc = msnippets.find_one(query)
#                 print("%s\n%s\n%s\n\n" % (doc['snippet_part1'], doc['snippet_part2'], doc['snippet_part3']))
                
# -

# to convert from isodate as written below: datetime.datetime.strptime(datestr,"%Y-%m-%d %H:%M:%S")
# snippets have already had double quotes replaced with single quotes
# snippets can be identified in mongodb by ccfn & ccsampnum (filename & file_idx)
def writeMTurkBatchFile(docs, filename, batch, holdout):
    with open(filename, 'w') as mtbf:
        mtbf.write('batch,idx,holdout,ccfn,ccsampnum,station,airdate,priorsample,sample,nextsample\n')
        #TODO: UTF-8 encoding; remove non- UTF-8
        for idx, doc in enumerate(docs):
            samplerecord = str(batch) + ',' + str(idx) + ',' + str(holdout).lower()  + ',' \
                    + doc['filename'] + ',' + str(doc['file_idx']) + ',' \
                    + doc['station'] + ',' + str(doc['airdatetime']) + ',"' + doc['snippet_part1'] + '","' \
                    + doc['snippet_part2'] + '","' + doc['snippet_part3'] + '"\n'
            mtbf.write(samplerecord)


# +
mturk_path = "mturk_batches/" # assumed to exist

# map from idx to _id and insert cluster ids; mark holdouts and train data in mongodb
with pymongo.MongoClient('localhost', MONGODB_PORT) as mclient:
    mdb = mclient.snippetdb
    msnippets = mdb.snippets
    
    # save holdouts
    for holdout_set in range(1,4):
        doc_cursor = msnippets.find({"holdout": holdout_set})
        writeMTurkBatchFile(doc_cursor, mturk_path + "holdouts_" + str(holdout_set) + ".csv", holdout_set, True)
    
    # save train sets
    for train_set in [0, 1, 2, 3, 10, 11, 12, 20, 21]:
        doc_cursor = msnippets.find({"train_set": train_set})
        writeMTurkBatchFile(doc_cursor, mturk_path + "train_set_" + str(train_set) + ".csv", train_set, False)
    
# -



StarcoderdataJupyterScriptsDedupFiltered