#!/usr/bin/python3

import os
import pickle
import sys

import django
from termcolor import colored

from .action import parse
from .edits import classify_edits, edit_graph, get_edits_from_traces
from .graph import graphviz
from .monkey import fix
from prolog.engine import test
from prolog.util import Token,  compose, decompose, stringify
from .util import indent

# Load django models.
os.environ['DJANGO_SETTINGS_MODULE'] = 'webmonkey.settings'
django.setup()
from django.apps import apps as tutor_apps
from django.contrib.auth.models import User
from tutor.models import Attempt, Problem, get_aux_code

# Get problem id from commandline.
if len(sys.argv) < 2:
    print('usage: ' + sys.argv[0] + ' <pid>')
    sys.exit(1)
pid = int(sys.argv[1])

problem = Problem.objects.get(pk=pid)
aux_code = get_aux_code(user=User.objects.get(pk=1), problem=problem)

attempts = Attempt.objects.filter(problem=problem)

# Load hint database stored in edits.pickle.
edits = tutor_apps.get_app_config('tutor').edits[problem.pk]
lines = tutor_apps.get_app_config('tutor').lines[problem.pk]
submissions = tutor_apps.get_app_config('tutor').submissions[problem.pk]
queries = tutor_apps.get_app_config('tutor').queries[problem.pk]

# Find incorrect submissions.
incorrect_all = []
for submission, count in sorted(submissions.items()):
    if not test(problem.name, submission + '\n' + aux_code):
        # This incorrect submission appeared in [count] attempts.
        incorrect_all += [submission]*count
incorrect = set(incorrect_all)

# Load current status (programs for which a hint was found).
try:
    done = pickle.load(open('status-'+str(problem.pk)+'.pickle', 'rb'))
except:
    done = []

def print_hint(solution, steps, fix_time, n_tested):
    if solution:
        print(colored('Hint found! Tested {} programs in {:.1f} s.'.format(n_tested, fix_time), 'green'))
        print(colored(' Edits', 'blue'))
        for step_type, line, (before, after) in steps:
            print('  {}: {} {} → {}'.format(line, step_type, stringify(before), stringify(after)))
        print(colored(' Final version', 'blue'))
        print(indent(compose(*decompose(solution)), 2))
    else:
        print(colored('Hint not found! Tested {} programs in {:.1f} s.'.format(n_tested, fix_time), 'red'))

# Test fix() on incorrect student submissions.
if len(sys.argv) >= 3 and sys.argv[2] == 'test':
    timeout = int(sys.argv[3]) if len(sys.argv) >= 4 else 10

    print('Fixing {}/{} programs (timeout={})…'.format(
        len([p for p in incorrect if p not in done]), len(incorrect), timeout))

    for i, program in enumerate(sorted(incorrect)):
        if program in done:
            continue
        print(colored('Analyzing program {0}/{1}…'.format(i+1, len(incorrect)), 'yellow'))
        print(indent(compose(*decompose(program)), 2))

        solution, steps, fix_time, n_tested = fix(problem.name, program, edits, lines, aux_code=aux_code, timeout=timeout)
        if solution:
            done.append(program)
        print_hint(solution, steps, fix_time, n_tested)
        print()

        pickle.dump(done, open('status-'+str(problem.pk)+'.pickle', 'wb'))

    print('Found hints for ' + str(len(done)) + ' of ' + str(len(incorrect)) + ' incorrect programs')

# Print info for this problem.
elif len(sys.argv) >= 3 and sys.argv[2] == 'info':
    # With no additional arguments, print some stats.
    if len(sys.argv) == 3:
        print('Problem {} ({}): {} edits in {} traces, fixed {}/{} ({}/{} unique)'.format(
            problem.pk, colored(problem.name, 'yellow'),
            colored(str(len(edits)), 'yellow'), colored(str(len([a.trace for a in attempts])), 'yellow'),
            colored(str(len([p for p in incorrect_all if p in done])), 'yellow'),
            colored(str(len(incorrect_all)), 'yellow'),
            colored(str(len(set(done))), 'yellow'),
            colored(str(len(set(incorrect))), 'yellow')))
    else:
        if sys.argv[3] == 'users':
            print(' '.join([str(a.user.pk) for a in attempts]))
        # Print all observed edits and their costs.
        elif sys.argv[3] == 'edits':
            inserts, removes, changes = classify_edits(edits)
            print('Inserts')
            for after, cost in sorted(inserts.items(), key=lambda x: x[1]):
                print(' {:.4f}\t{}'.format(cost, stringify(after)))
            print('Removes')
            for before, cost in sorted(removes.items(), key=lambda x: x[1]):
                print(' {:.4f}\t{}'.format(cost, stringify(before)))
            print('Changes')
            for (before, after), cost in sorted(changes.items(), key=lambda x: x[1]):
                print(' {:.4f}\t{} → {}'.format(cost,
                                               stringify(before if before else [Token('INVALID', 'ε')]),
                                               stringify(after if after else [Token('INVALID', 'ε')])))
        # Print all student submissions not (yet) corrected.
        elif sys.argv[3] == 'unsolved':
            for p in sorted(incorrect):
                if p in done:
                    continue
                print(indent(compose(*decompose(p)), 2))
                print()
        # Print all student queries and their counts.
        elif sys.argv[3] == 'queries':
            for query, count in queries.most_common():
                print('  ' + str(count) + '\t' + query)

# Print the edit graph in graphviz dot syntax.
elif len(sys.argv) == 4 and sys.argv[2] == 'graph':
    uid = int(sys.argv[3])
    user = User.objects.get(pk=uid)
    actions = parse(Attempt.objects.get(problem=problem, user=user).trace)

    nodes, submissions, queries = edit_graph(actions)

    def position(node):
        return (node.data[1]*150, node.data[0]*-60)

    def label(node):
        return stringify(node.data[2])

    def node_attr(node):
        if node.ein and node.data[2] == node.ein[0].data[2]:
            return 'color="gray", shape="point"'
        return ''

    def edge_attr(a, b):
        if a.data[2] == b.data[2]:
            return 'arrowhead="none"'
        return ''

    graphviz_str = graphviz(nodes, pos=position, label=label,
                            node_attr=node_attr, edge_attr=edge_attr)
    print(graphviz_str)

# Run interactive loop.
else:
    while True:
        # Read the program from stdin.
        print('Enter program, end with empty line:')
        code = ''
        try:
            while True:
                line = input()
                if not line:
                    break
                code += line + '\n'
        except EOFError:
            break

        # Try finding a fix.
        print(colored('Analyzing program…', 'yellow'))
        solution, steps, fix_time, n_tested = fix(problem.name, code, edits, aux_code=aux_code, debug=True)
        print_hint(solution, steps, fix_time, n_tested)