from selenium import webdriver, selenium
from util.perf import tick, tock
import util.str_util as str_util
import util.seutil as seutil
import sys
import time
import itertools
import random
import math
import scipy.stats as stats
import numpy as np

GET_FEATURES_JS = "js/get-features.js"
UNDERSCORE_JS = "js/underscore.js"

def tokenize_command(command):
    command = unicode(command, 'utf-8')
    return command.split(' ')

def untokenize_subcommand(sub):
    return ' '.join(sub)

class Action:
    # These are the features used in the numeric representation of the feature
    # vector corresponding to this Action. There are some features from get-features.js
    # and/or extend_and_norm_feature which are dropped from this representation.
    FEATURE_NAMES = [
        'text_words_edit',
        'n_children',
        #'width',
        #'height',
        'sibling_text_words_edit',
        'tagname_edit',
        'typeable',
        #'clickable',
        #'text_size',
        'has_id',
        'has_class',
        #'button_model',
        'relative_x',
        'relative_y',
        'alreadyInteracted',
    ]

    def __init__(self, element, atype, features, params=None):
        self.element = element
        self.type = atype
        self.params = params
        self.features = features

    def perform(self, driver, dry=False):

        driver.execute_script('''
            var elem = arguments[0]
            elem.setAttribute("x-WebtalkInteracted", "1")
            ''', self.element)

        if dry and self.type == 'click':
            seutil.highlight(driver, self.element, opacity=0.5)
            return

        if self.type == 'click':
            self.element.click()
        elif self.type == 'type':
            self.element.clear()
            self.element.send_keys(untokenize_subcommand(self.params))

    def as_numeric_vector(self):
        return tuple([
            self.features[k] for k in self.FEATURE_NAMES
        ])

    def __repr__(self):
        classes = self.element.get_attribute("class")
        if self.params == None:
            return '<Action %s on %s#%s.%s>' % (self.type, self.features['tagname'], self.features['id'], classes)
        return '<Action %s on %r, %s#%s.%s>' % (self.type, self.params, self.features['tagname'], self.features['id'], classes)


class State:
    def __init__(self, command, features):
        self.command = command
        self.features = features

    def _subcommands(self):
        if hasattr(self, '_subcommand_cache'):
            return self._subcommand_cache

        subcommands = []
        for window_size in range(1, len(self.command)+1):
            for i in range(0, len(self.command)-window_size+1):
                subcommands.append(self.command[i:i+window_size])

        self._subcommand_cache = subcommands
        return self._subcommand_cache

    def enumerate_actions(self):
        actions = []
        for el in self.features:
            if el['typeable'] == 1:
                for subwords in self._subcommands():
                    actions.append(Action(el['element'], 'type', el, params=subwords))
            else:
                actions.append(Action(el['element'], 'click', el))
        return actions

    def phi_dot_theta(self, action, theta, verbose=True):
        phi = action.as_numeric_vector()

        if verbose:
            print 'Product:'
            for i,f in enumerate(Action.FEATURE_NAMES):
                print '\t', phi[i]*theta[i], f, phi[i], '*', theta[i]

        return np.dot(phi, theta)

    def get_action_probs(self, actions, theta):
        '''Chooses the modal action'''
        products = {}

        normalization = 0.0
        for action in actions:
            v = math.exp(self.phi_dot_theta(action, theta))
            products[action] = v
            normalization += v

        max_v, max_action = None, []
        for action in products:
            products[action] /= normalization
            if max_v == None or products[action] >= max_v:
                if max_v == products[action]:
                    max_action.append(action)
                else:
                    max_action = [action]

                max_v = products[action]

        return random.choice(max_action), max_v, products

def start(url):
    driver = seutil.get_driver()
    driver.get(url)
    return driver


ELEMENT_DATA = json.load(open('element_sizes.json', 'rb'))

BUTTON_SIZE_KDE = stats.gaussian_kde(numpy.transpose(numpy.array(
    [x[1:] for x in ELEMENT_DATA if x[0]]
)))

ELEMENT_SIZE_KDE = stats.gaussian_kde(numpy.transpose(numpy.array(
    [x[1:] for x in ELEMENT_DATA]
)))

def extend_and_norm_feature(element, feature, command, num_elems):
    feature['element'] = element

    feature['tagname_edit'] = str_util.get_normed_dist_for_words(command, [feature['tagname']])
    feature['text_words_edit'] = str_util.get_normed_dist_for_words(command, feature['text_words'])
    feature['sibling_text_words_edit'] = str_util.get_normed_dist_for_words(command, feature['sibling_text_words'])
    feature['n_children'] = 1 - float(feature['n_children']) / num_elems

    w,h = feature['width'], feature['height']
    mx, my, sx, sy = 54.611, 25.206, 43.973, 6.467

    #likelihood_button = (stats.norm.cdf(h+1, loc=my, scale=sy) - stats.norm.cdf(h, loc=my, scale=sy)) * (stats.norm.cdf(w+1, loc=mx, scale=sx) - stats.norm.cdf(w, loc=mx, scale=sx))
    likelihood_button = BUTTON_SIZE_KDE.integrate_box([w,h], [w+1,h+1])
    prior = 0.02
    marginal = ELEMENT_SIZE_KDE.integrate_box([w,h], [w+1,h+1])

    feature['button_model'] = (likelihood_button * prior) / marginal

    # relative x and y can  be more than 1 because things can be beyond the edge of the window
    # so nudge things to be between -1 and 1
    feature['relative_x'] = np.arctan(1 * (feature['relative_x'] + 0.5)) / (np.pi / 2)
    feature['relative_y'] = np.arctan(1 * feature['relative_y']) / (np.pi / 2)

    # new on page?
    # position (relative to last action?)
    # color
    # relative tab index
    # text specificity
    return feature

def extract(driver, command):
    driver.execute_script(open(UNDERSCORE_JS).read())
    features, tree = driver.execute_script(open(GET_FEATURES_JS).read())
    features = [extend_and_norm_feature(f[0], f[1], command, len(features)) for f in features]

    return features, tree

def build_state(driver, command):
    features, tree = extract(driver, command)
    return State(command, features)



@@ -241,16 +241,28 @@
 py as np
+%0Aimport json
 %0A%0AGET_FE
@@ -1030,17 +1030,16 @@
         
-#
 'button_
@@ -4447,36 +4447,33 @@
 s.gaussian_kde(n
-umpy
+p
 .transpose(numpy
@@ -4460,36 +4460,33 @@
 e(np.transpose(n
-umpy
+p
 .array(%0A    %5Bx%5B1
@@ -4563,20 +4563,17 @@
 an_kde(n
-umpy
+p
 .transpo
@@ -4576,20 +4576,17 @@
 nspose(n
-umpy
+p
 .array(%0A

Add button model back