<commit_before><commit_msg>Update subject subdocuments with magic and wonder.<commit_after>#!/usr/bin/env python
import json
import names # Use 'pip install names'
import random
import pymongo

'''
Updates the DB for a dev instance, making the subject subdocuments rich and interesting. 
'''

# Connect to the DB
db = pymongo.MongoClient('mongodb://docker.local.flywheel.io:9001/scitran').get_default_database()

# Update DB: Set group names to use the group '_id' (Should be done by the UI)
for g in db.groups.find({},['_id']):
  group_name = g['_id'].capitalize() + ' Group'
  db.groups.update_one({'_id': g['_id']},{'$set':{'name': group_name}})

# Get a list of the subjects/sessions
subjects = list(db.sessions.find({}, ['subject']))

# Get unique subject codes (some have multiple sessions)
codes = set([s['subject']['code'] for s in subjects])
total_sessions = len(codes)

# Update the subject subdocument only once per unique subject code
idx = 0
for c in codes:
    # Set half to male, rest female, using gender specific names
    if idx <= total_sessions * 0.5:
        first_name = names.get_first_name(gender='male')
        sex = 'male'
    else:
        first_name = names.get_first_name(gender='female')
        sex = 'female'
    last_name = names.get_last_name()
    
    # Set the age (in seconds) from a normal distrubition
    age = int(random.normalvariate(35, 10) * 31536000)
    
    # Build some fun subject metadata for each subject 
    metadata = {}
    ses = ['High', 'Middle', 'Low']
    metadata['SES'] = ses[random.randrange(0, len(ses))]
    metadata['IQ'] = int(random.normalvariate(100, 15)) 
    metadata['Education'] = int(random.normalvariate(14, 2))
    metadata['Verbal Reasoning'] = {}
    metadata['Verbal Reasoning']['Initial Score'] = random.randrange(130, 170)
    metadata['Verbal Reasoning']['Post Score'] = random.randrange(130, 170)
    metadata['Quantitative Reasoning'] = {}
    metadata['Quantitative Reasoning']['Initial Score'] = random.randrange(130, 170)
    metadata['Quantitative Reasoning']['Post Score'] = random.randrange(130, 170)
    metadata['Analytical Writing'] = {}
    metadata['Analytical Writing']['Initial Score'] = random.randrange(2, 6)
    metadata['Analytical Writing']['Post Score'] = random.randrange(2, 6)
    
    # Update the DB 
    db.sessions.update_many({'subject.code': c},{'$set':{'subject.lastname': last_name}})
    db.sessions.update_many({'subject.code': c},{'$set':{'subject.firstname': first_name}})
    db.sessions.update_many({'subject.code': c},{'$set':{'subject.sex': sex}})
    db.sessions.update_many({'subject.code': c},{'$set':{'subject.age': age}})
    db.sessions.update_many({'subject.code': c},{'$set':{'subject.metadata': metadata}})
    
    # Increment the index
    idx +=1
{'fraction_non_alphanumeric': 0.16007194244604317, 'fraction_numerical': 0.02050359712230216, 'mean_word_length': 4.442270058708415, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6166600', 'n_tokens_mistral': 924, 'n_tokens_neox': 855, 'n_words': 226}