from django.core.management.base import BaseCommand
from django.contrib.auth.models import User

from student.models import UserTestGroup

import random
import sys
import datetime
from textwrap import dedent

import json
from pytz import UTC


def group_from_value(groups, v):
    ''' Given group: (('a',0.3),('b',0.4),('c',0.3)) And random value
    in [0,1], return the associated group (in the above case, return
    'a' if v<0.3, 'b' if 0.3<=v<0.7, and 'c' if v>0.7
'''
    sum = 0
    for (g, p) in groups:
        sum = sum + p
        if sum > v:
            return g
    return g  # For round-off errors


class Command(BaseCommand):
    help = dedent("""\
        Assign users to test groups. Takes a list of groups:
        a:0.3,b:0.4,c:0.3 file.txt "Testing something"
        Will assign each user to group a, b, or c with
        probability 0.3, 0.4, 0.3. Probabilities must
        add up to 1.

        Will log what happened to file.txt.
    """)

    def handle(self, *args, **options):
        if len(args) != 3:
            print "Invalid number of options"
            sys.exit(-1)

        # Extract groups from string
        group_strs = [x.split(':') for x in args[0].split(',')]
        groups = [(group, float(value)) for group, value in group_strs]
        print "Groups", groups

        ## Confirm group probabilities add up to 1
        total = sum(zip(*groups)[1])
        print "Total:", total
        if abs(total - 1) > 0.01:
            print "Total not 1"
            sys.exit(-1)

        ## Confirm groups don't already exist
        for group in dict(groups):
            if UserTestGroup.objects.filter(name=group).count() != 0:
                print group, "already exists!"
                sys.exit(-1)

        group_objects = {}

        f = open(args[1], "a+")

        ## Create groups
        for group in dict(groups):
            utg = UserTestGroup()
            utg.name = group
            utg.description = json.dumps({"description": args[2]},
                                         {"time": datetime.datetime.now(UTC).isoformat()})
            group_objects[group] = utg
            group_objects[group].save()

        ## Assign groups
        users = list(User.objects.all())
        count = 0
        for user in users:
            if count % 1000 == 0:
                print count
            count = count + 1
            v = random.uniform(0, 1)
            group = group_from_value(groups, v)
            group_objects[group].users.add(user)
            f.write(u"Assigned user {name} ({id}) to {group}\n".format(
                name=user.username,
                id=user.id,
                group=group
            ).encode('utf-8'))

        ## Save groups
        for group in group_objects:
            group_objects[group].save()
        f.close()

# python manage.py assigngroups summary_test:0.3,skip_summary_test:0.7 log.txt "Do previews of future materials help?"
# python manage.py assigngroups skip_capacitor:0.3,capacitor:0.7 log.txt "Do we show capacitor in linearity tutorial?"