#!/usr/bin/python3

import math
import time

from .edits import classify_edits
from .prolog.engine import test
from .prolog.util import compose, decompose, map_vars, rename_vars, stringify
from .util import PQueue, Token

# score a program (a list of lines) according to lines distribution
def score(program, lines):
    result = 1
    for line in program:
        line_normal = list(line)
        rename_vars(line_normal)
        line_normal = tuple(line_normal)
        result *= lines.get(line_normal, 0.01)

    if len(program) == 0 or result == 0:
        return 0.01
    return math.pow(result, 1/len(program))

# find a sequence of edits that fixes [code]
def fix(name, code, edits, aux_code='', timeout=30, debug=False):
    todo = PQueue()  # priority queue of candidate solutions
    done = set()     # set of already-analyzed solutions

    # Add a new candidate solution ([lines]+[rules]) to the priority queue.
    # This solution is generated by applying [step] with [cost] to [prev] task.
    def add_task(lines, rules, prev=None, step=None, cost=None):
        if prev is None:
            path = ()
            path_cost = 1.0
        else:
            path = tuple(list(prev[1]) + [step])
            path_cost = prev[2] * cost
        todo.push(((tuple(lines), tuple(rules)), path, path_cost), -path_cost)

    lines, rules = decompose(code)
    add_task(lines, rules)

    inserts, removes, changes = classify_edits(edits)
    start_time = time.monotonic()
    n_tested = 0
    while True:
        total_time = time.monotonic() - start_time
        if total_time > timeout:
            break

        task = todo.pop()
        if task == None:
            break

        (lines, rules), path, path_cost = task
        code = compose(lines, rules)
        if code in done:
            continue
        done.add(code)

        if debug:
            print('Cost {:.12f}'.format(path_cost))
            for line, (before, after) in path:
                print('line ' + str(line) + ':\t' + stringify(before) + ' → ' + stringify(after))

        # if the code is correct, we are done
        try:
            if test(name, code + '\n' + aux_code):
                return code, path, total_time, n_tested
        except:
            pass
        n_tested += 1

        # otherwise generate new solutions
        rule_no = 0
        for start, end in rules:
            rule = lines[start:end]
            rule_tokens = [t for line in rule for t in line]

            for line_idx in range(start, end):
                line = lines[line_idx]

                line_normal = list(line)
                rename_vars(line_normal)
                line_normal = tuple(line_normal)

                seen = False
                for (before, after), cost in changes.items():
                    if line_normal == before:
                        seen = True
                        mapping = map_vars(before, after, line, rule_tokens)
                        after_real = tuple([t if t.type != 'VARIABLE' else Token('VARIABLE', mapping[t.val]) for t in after])
                        new_lines = lines[:line_idx] + (after_real,) + lines[line_idx+1:]
                        new_step = ((rule_no, line_idx-start), (tuple(line), after_real))

                        add_task(new_lines, rules, prev=task, step=new_step, cost=cost)

                # if nothing could be done with this line, try removing it
                # (maybe try removing in any case?)
                if line_normal in removes.keys() or not seen:
                    new_lines = lines[:line_idx] + lines[line_idx+1:]
                    new_rules = []
                    for old_start, old_end in rules:
                        new_start, new_end = (old_start - (0 if old_start <= line_idx else 1),
                                              old_end - (0 if old_end <= line_idx else 1))
                        if new_end > new_start:
                            new_rules.append((new_start, new_end))
                    new_step = ((rule_no, line_idx-start), (tuple(line), ()))
                    new_cost = removes[line_normal] if line_normal in removes.keys() else 0.9

                    add_task(new_lines, new_rules, prev=task, step=new_step, cost=new_cost)

            # try adding a line to this rule… would need to distinguish between
            # head/body lines in transforms
            for after, cost in inserts.items():
                mapping = map_vars([], after, [], rule_tokens)
                after_real = [t if t.type != 'VARIABLE' else Token('VARIABLE', mapping[t.val]) for t in after]
                after_real = tuple(after_real)
                new_lines = lines[:end] + (after_real,) + lines[end:]
                new_rules = []
                for old_start, old_end in rules:
                    new_rules.append((old_start + (0 if old_start < end else 1),
                                      old_end + (0 if old_end < end else 1)))
                new_step = ((rule_no, end-start), ((), after_real))

                add_task(new_lines, new_rules, prev=task, step=new_step, cost=cost)
            rule_no += 1

        # try adding a new fact
        if len(rules) < 2:
            for after, cost in inserts.items():
                new_lines = lines + (after,)
                new_rules = rules + (((len(lines), len(lines)+1)),)
                new_step = ((len(new_rules)-1, 0), (tuple(), tuple(after)))

                add_task(new_lines, new_rules, prev=task, step=new_step, cost=cost)

    return '', [], total_time, n_tested