import django

from enum import Enum
from typing import List, Tuple, Optional, Dict, Any
from collections import OrderedDict

from django.db import models
from django.db.models import sql
from django.db.models.constants import LOOKUP_SEP
from django.core.exceptions import SuspiciousOperation

from .fields import HStoreField
from .expressions import HStoreColumn
from .datastructures import ConditionalJoin

class ConflictAction(Enum):
    """Possible actions to take on a conflict."""

    NOTHING = 'NOTHING'
    UPDATE = 'UPDATE'

class PostgresQuery(sql.Query):
    def rename_annotations(self, annotations) -> None:
        """Renames the aliases for the specified annotations:

            .annotate(myfield=F('somestuf__myfield'))
            .rename_annotations(myfield='field')

        Arguments:
            annotations:
                The annotations to rename. Mapping the
                old name to the new name.
        """

        for old_name, new_name in annotations.items():
            annotation = self.annotations.get(old_name)

            if not annotation:
                raise SuspiciousOperation((
                    'Cannot rename annotation "{old_name}" to "{new_name}", because there'
                    ' is no annotation named "{old_name}".'
                ).format(old_name=old_name, new_name=new_name))

            self._annotations = OrderedDict(
                [(new_name, v) if k == old_name else (k, v) for k, v in self._annotations.items()])

            if django.VERSION < (2, 0):
                self.set_annotation_mask(
                    (new_name if v == old_name else v for v in (self.annotation_select_mask or [])))

    def add_join_conditions(self, conditions: Dict[str, Any]) -> None:
        """Adds an extra condition to an existing JOIN.

        This allows you to for example do:

            INNER JOIN othertable ON (mytable.id = othertable.other_id AND [extra conditions])

        This does not work if nothing else in your query doesn't already generate the
        initial join in the first place.
        """

        alias = self.get_initial_alias()
        opts = self.get_meta()

        for name, value in conditions.items():
            parts = name.split(LOOKUP_SEP)
            join_info = self.setup_joins(parts, opts, alias, allow_many=True)
            self.trim_joins(join_info[1], join_info[3], join_info[4])

            target_table = join_info[3][-1]
            field = join_info[1][-1]
            join = self.alias_map.get(target_table)

            if not join:
                raise SuspiciousOperation((
                    'Cannot add an extra join condition for "%s", there\'s no'
                    ' existing join to add it to.'
                ) % target_table)

            # convert the Join object into a ConditionalJoin object, which
            # allows us to add the extra condition
            if not isinstance(join, ConditionalJoin):
                self.alias_map[target_table] = ConditionalJoin.from_join(join)
                join = self.alias_map[target_table]

            join.add_condition(field, value)

    def add_fields(self, field_names: List[str], allow_m2m: bool=True) -> bool:
        """
        Adds the given (model) fields to the select set. The field names are
        added in the order specified.

        This overrides the base class's add_fields method. This is called by
        the .values() or .values_list() method of the query set. It instructs
        the ORM to only select certain values. A lot of processing is neccesarry
        because it can be used to easily do joins. For example, `my_fk__name` pulls
        in the `name` field in foreign key `my_fk`.

        In our case, we want to be able to do `title__en`, where `title` is a HStoreField
        and `en` a key. This doesn't really involve a join. We iterate over the specified
        field names and filter out the ones that refer to HStoreField and compile it into
        an expression which is added to the list of to be selected fields using `self.add_select`.
        """

        alias = self.get_initial_alias()
        opts = self.get_meta()
        cols = []
        for name in field_names:
            parts = name.split(LOOKUP_SEP)

            # it cannot be a special hstore thing if there's no __ in it
            if len(parts) > 1:
                column_name, hstore_key = parts[:2]
                is_hstore, field = self._is_hstore_field(column_name)
                if is_hstore:
                    cols.append(
                        HStoreColumn(self.model._meta.db_table or self.model.name, field, hstore_key)
                    )
                    continue

            join_info = self.setup_joins(parts, opts, alias, allow_many=allow_m2m)
            targets, final_alias, joins = self.trim_joins(
                join_info[1], join_info[3], join_info[4]
            )

            for target in targets:
                cols.append(target.get_col(final_alias))
        if cols:
            self.set_select(cols)

    def _is_hstore_field(self, field_name: str) -> Tuple[bool, Optional[models.Field]]:
        """Gets whether the field with the specified name is a
        HStoreField.

        Returns
            A tuple of a boolean indicating whether the field
            with the specified name is a HStoreField, and the
            field instance.
        """

        field_instance = None
        for field in self.model._meta.local_concrete_fields:
            if field.name == field_name or field.column == field_name:
                field_instance = field
                break

        return isinstance(field_instance, HStoreField), field_instance

class PostgresInsertQuery(sql.InsertQuery):
    """Insert query using PostgreSQL."""

    def __init__(self, *args, **kwargs):
        """Initializes a new instance :see:PostgresInsertQuery."""

        super(PostgresInsertQuery, self).__init__(*args, **kwargs)

        self.conflict_target = []
        self.conflict_action = ConflictAction.UPDATE

        self.update_fields = []

    def values(self, objs: List, insert_fields: List, update_fields: List=[]):
        """Sets the values to be used in this query.

        Insert fields are fields that are definitely
        going to be inserted, and if an existing row
        is found, are going to be overwritten with the
        specified value.

        Update fields are fields that should be overwritten
        in case an update takes place rather than an insert.
        If we're dealing with a INSERT, these will not be used.

        Arguments:
            objs:
                The objects to apply this query to.

            insert_fields:
                The fields to use in the INSERT statement

            update_fields:
                The fields to only use in the UPDATE statement.
        """

        self.insert_values(insert_fields, objs, raw=False)
        self.update_fields = update_fields


from ..fields import HStoreField

class HStoreRequiredSchemaEditorMixin:
    sql_hstore_required_create = (
        'ALTER TABLE {table} '
        'ADD CONSTRAINT {name} '
        'CHECK (({field}->\'{key}\') '
        'IS NOT NULL)'
    )

    sql_hstore_required_rename = (
        'ALTER TABLE {table} '
        'RENAME CONSTRAINT '
        '{old_name} '
        'TO '
        '{new_name}'
    )

    sql_hstore_required_drop = (
        'ALTER TABLE {table} '
        'DROP CONSTRAINT IF EXISTS {name}'
    )

    def create_model(self, model):
        """Ran when a new model is created."""

        for field in model._meta.local_fields:
            if not isinstance(field, HStoreField):
                continue

            self.add_field(model, field)

    def delete_model(self, model):
        """Ran when a model is being deleted."""

        for field in model._meta.local_fields:
            if not isinstance(field, HStoreField):
                continue

            self.remove_field(model, field)

    def alter_db_table(self, model, old_db_table, new_db_table):
        """Ran when the name of a model is changed."""

        for field in model._meta.local_fields:
            if not isinstance(field, HStoreField):
                continue

            for key in self._iterate_required_keys(field):
                self._rename_hstore_required(
                    old_db_table,
                    new_db_table,
                    field,
                    field,
                    key
                )

    def add_field(self, model, field):
        """Ran when a field is added to a model."""

        for key in self._iterate_required_keys(field):
            self._create_hstore_required(
                model._meta.db_table,
                field,
                key
            )

    def remove_field(self, model, field):
        """Ran when a field is removed from a model."""

        for key in self._iterate_required_keys(field):
            self._drop_hstore_required(
                model._meta.db_table,
                field,
                key
            )

    def alter_field(self, model, old_field, new_field, strict=False):
        """Ran when the configuration on a field changed."""

        is_old_field_hstore = isinstance(old_field, HStoreField)
        is_new_field_hstore = isinstance(new_field, HStoreField)

        if not is_old_field_hstore and not is_new_field_hstore:
            return

        old_required = getattr(old_field, 'required', []) or []
        new_required = getattr(new_field, 'required', []) or []

        # handle field renames before moving on
        if str(old_field.column) != str(new_field.column):
            for key in self._iterate_required_keys(old_field):
                self._rename_hstore_required(
                    model._meta.db_table,
                    model._meta.db_table,
                    old_field,
                    new_field,
                    key
                )

        # drop the constraints for keys that have been removed
        for key in old_required:
            if key not in new_required:
                self._drop_hstore_required(
                    model._meta.db_table,
                    old_field,
                    key
                )

        # create new constraints for keys that have been added
        for key in new_required:
            if key not in old_required:
                self._create_hstore_required(
                    model._meta.db_table,
                    new_field,
                    key
                )

    def _create_hstore_required(self, table_name, field, key):
        """Creates a REQUIRED CONSTRAINT for the specified hstore key."""

        name = self._required_constraint_name(
            table_name, field, key)

        sql = self.sql_hstore_required_create.format(
            name=self.quote_name(name),
            table=self.quote_name(table_name),
            field=self.quote_name(field.column),
            key=key
        )
        self.execute(sql)

    def _rename_hstore_required(self, old_table_name, new_table_name,
                                old_field, new_field, key):
        """Renames an existing REQUIRED CONSTRAINT for the specified
        hstore key."""

        old_name = self._required_constraint_name(
            old_table_name, old_field, key)
        new_name = self._required_constraint_name(
            new_table_name, new_field, key)

        sql = self.sql_hstore_required_rename.format(
            table=self.quote_name(new_table_name),
            old_name=self.quote_name(old_name),
            new_name=self.quote_name(new_name)
        )
        self.execute(sql)

    def _drop_hstore_required(self, table_name, field, key):
        """Drops a REQUIRED CONSTRAINT for the specified hstore key."""

        name = self._required_constraint_name(
            table_name, field, key)

        sql = self.sql_hstore_required_drop.format(
            table=self.quote_name(table_name),
            name=self.quote_name(name)
        )
        self.execute(sql)

    @staticmethod
    def _required_constraint_name(table: str, field, key):
        """Gets the name for a CONSTRAINT that applies
        to a single hstore key.

        Arguments:
            table:
                The name of the table the field is
                a part of.

            field:
                The hstore field to create a
                UNIQUE INDEX for.

            key:
                The name of the hstore key
                to create the name for.

        Returns:
            The name for the UNIQUE index.
        """

        return '{table}_{field}_required_{postfix}'.format(
            table=table,
            field=field.column,
            postfix=key
        )

    @staticmethod
    def _iterate_required_keys(field):
        """Iterates over the keys marked as "required"
        in the specified field.

        Arguments:
            field:
                The field of which key's to
                iterate over.
        """

        required_keys = getattr(field, 'required', None)
        if not required_keys:
            return

        for key in required_keys:
            yield key


import django
from django.db.models.indexes import Index

class ConditionalUniqueIndex(Index):
    """
    Creates a partial unique index based on a given condition.

    Useful, for example, if you need unique combination of foreign keys, but you might want to include
    NULL as a valid value. In that case, you can just use:
    >>> class Meta:
    ...    indexes = [
    ...        ConditionalUniqueIndex(fields=['a', 'b', 'c'], condition='"c" IS NOT NULL'),
    ...        ConditionalUniqueIndex(fields=['a', 'b'], condition='"c" IS NULL')
    ...    ]
    """

    sql_create_index = "CREATE UNIQUE INDEX %(name)s ON %(table)s (%(columns)s)%(extra)s WHERE %(condition)s"

    def __init__(self, condition: str, fields=[], name=None):
        """Initializes a new instance of :see:ConditionalUniqueIndex."""

        super().__init__(fields=fields, name=name)
        self.condition = condition

    def create_sql(self, model, schema_editor, using=''):
        """Creates the actual SQL used when applying the migration."""
        if django.VERSION >= (2, 0):
            statement = super().create_sql(model, schema_editor, using)
            statement.template = self.sql_create_index
            statement.parts['condition'] = self.condition
            return statement
        else:
            sql_create_index = self.sql_create_index
            sql_parameters = {
                **Index.get_sql_create_template_values(self, model, schema_editor, using),
                'condition': self.condition
            }
            return sql_create_index % sql_parameters

    def deconstruct(self):
        """Serializes the :see:ConditionalUniqueIndex for the migrations file."""
        path = '%s.%s' % (self.__class__.__module__, self.__class__.__name__)
        path = path.replace('django.db.models.indexes', 'django.db.models')
        return path, (), {'fields': self.fields, 'name': self.name, 'condition': self.condition}


import os
import distutils.cmd
import subprocess

from setuptools import find_packages, setup

class BaseCommand(distutils.cmd.Command):
    user_options = []

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

def create_command(text, commands):
    """Creates a custom setup.py command."""

    class CustomCommand(BaseCommand):
        description = text

        def run(self):
            for cmd in commands:
                subprocess.check_call(cmd)

    return CustomCommand

with open(os.path.join(os.path.dirname(__file__), 'README.rst'), encoding='utf-8') as readme:
    README = readme.read().split('h1>\n\n', 2)[1]

setup(
    name='django-postgres-extra',
    version='1.21a16',
    packages=find_packages(),
    include_package_data=True,
    license='MIT License',
    description='Bringing all of PostgreSQL\'s awesomeness to Django.',
    long_description=README,
    url='https://github.com/SectorLabs/django-postgres-extra',
    author='Sector Labs',
    author_email='open-source@sectorlabs.ro',
    keywords=['django', 'postgres', 'extra', 'hstore', 'ltree'],
    classifiers=[
        'Environment :: Web Environment',
        'Framework :: Django',
        'Intended Audience :: Developers',
        'License :: OSI Approved :: MIT License',
        'Operating System :: OS Independent',
        'Programming Language :: Python',
        'Programming Language :: Python :: 3.5',
        'Topic :: Internet :: WWW/HTTP',
        'Topic :: Internet :: WWW/HTTP :: Dynamic Content',
    ],
    cmdclass={
        'lint': create_command(
            'Lints the code',
            [['flake8', 'setup.py', 'psqlextra', 'tests']],
        ),
    },
)


import importlib
import logging

from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from django.db import ProgrammingError
from django.db.backends.postgresql.base import \
    DatabaseWrapper as Psycopg2DatabaseWrapper

from .hstore_unique import HStoreUniqueSchemaEditorMixin
from .hstore_required import HStoreRequiredSchemaEditorMixin

logger = logging.getLogger(__name__)

def _get_backend_base():
    """Gets the base class for the custom database back-end.

    This should be the Django PostgreSQL back-end. However,
    some people are already using a custom back-end from
    another package. We are nice people and expose an option
    that allows them to configure the back-end we base upon.

    As long as the specified base eventually also has
    the PostgreSQL back-end as a base, then everything should
    work as intended.
    """
    base_class_name = getattr(
        settings,
        'POSTGRES_EXTRA_DB_BACKEND_BASE',
        'django.db.backends.postgresql'
    )

    base_class_module = importlib.import_module(base_class_name + '.base')
    base_class = getattr(base_class_module, 'DatabaseWrapper', None)

    if not base_class:
        raise ImproperlyConfigured((
            '\'%s\' is not a valid database back-end.'
            ' The module does not define a DatabaseWrapper class.'
            ' Check the value of POSTGRES_EXTRA_DB_BACKEND_BASE.'
        ) % base_class_name)

    if isinstance(base_class, Psycopg2DatabaseWrapper):
        raise ImproperlyConfigured((
            '\'%s\' is not a valid database back-end.'
            ' It does inherit from the PostgreSQL back-end.'
            ' Check the value of POSTGRES_EXTRA_DB_BACKEND_BASE.'
        ) % base_class_name)

    return base_class

def _get_schema_editor_base():
    """Gets the base class for the schema editor.

    We have to use the configured base back-end's
    schema editor for this."""

    return _get_backend_base().SchemaEditorClass

class SchemaEditor(_get_schema_editor_base()):
    """Custom schema editor, see mixins for implementation."""

    post_processing_mixins = [
        HStoreUniqueSchemaEditorMixin(),
        HStoreRequiredSchemaEditorMixin()
    ]

    def __init__(self, connection, collect_sql=False, atomic=True):
        super(SchemaEditor, self).__init__(
            connection, collect_sql, atomic)

        self.base = super(SchemaEditor, self)

        for mixin in self.post_processing_mixins:
            mixin.execute = self.execute
            mixin.quote_name = self.quote_name

    def create_model(self, model):
        """Ran when a new model is created."""

        super().create_model(model)

        for mixin in self.post_processing_mixins:
            mixin.create_model(model)

    def delete_model(self, model):
        """Ran when a model is being deleted."""

        for mixin in self.post_processing_mixins:
            mixin.delete_model(model)

        super().delete_model(model)

    def alter_db_table(self, model, old_db_table, new_db_table):
        """Ran when the name of a model is changed."""

        super(SchemaEditor, self).alter_db_table(
            model, old_db_table, new_db_table
        )

        for mixin in self.post_processing_mixins:
            mixin.alter_db_table(
                model,
                old_db_table,
                new_db_table
            )

    def add_field(self, model, field):
        """Ran when a field is added to a model."""

        super(SchemaEditor, self).add_field(model, field)

        for mixin in self.post_processing_mixins:
            mixin.add_field(model, field)

    def remove_field(self, model, field):
        """Ran when a field is removed from a model."""

        for mixin in self.post_processing_mixins:
            mixin.remove_field(model, field)

        super(SchemaEditor, self).remove_field(model, field)

    def alter_field(self, model, old_field, new_field, strict=False):
        """Ran when the configuration on a field changed."""

        super(SchemaEditor, self).alter_field(
            model, old_field, new_field, strict
        )

        for mixin in self.post_processing_mixins:
            mixin.alter_field(
                model, old_field, new_field, strict
            )

class DatabaseWrapper(_get_backend_base()):
    """Wraps the standard PostgreSQL database back-end.

    Overrides the schema editor with our custom
    schema editor and makes sure the `hstore`
    extension is enabled."""

    SchemaEditorClass = SchemaEditor

    def prepare_database(self):
        """Ran to prepare the configured database.

        This is where we enable the `hstore` extension
        if it wasn't enabled yet."""

        super().prepare_database()
        with self.cursor() as cursor:
            try:
                cursor.execute('CREATE EXTENSION IF NOT EXISTS hstore')
            except ProgrammingError:  # permission denied
                logger.warning(
                    'Failed to create "hstore" extension. '
                    'Tables with hstore columns may fail to migrate. '
                    'If hstore is needed, make sure you are connected '
                    'to the database as a superuser '
                    'or add the extension manually.',
                    exc_info=True)


from typing import List, Tuple, Union

from django.contrib.postgres.fields import HStoreField as DjangoHStoreField
from django.db.models.expressions import Expression
from django.db.models.fields import Field

class HStoreField(DjangoHStoreField):
    """Improved version of Django's :see:HStoreField that
    adds support for database-level constraints.

    Notes:
        - For the implementation of uniqueness, see the
          custom database back-end.
    """

    def __init__(self, *args,
                 uniqueness: List[Union[str, Tuple[str, ...]]]=None,
                 required: List[str]=None, **kwargs):
        """Initializes a new instance of :see:HStoreField."""

        super(HStoreField, self).__init__(*args, **kwargs)

        self.uniqueness = uniqueness
        self.required = required

    def get_prep_value(self, value):
        """Override the base class so it doesn't cast all values
        to strings.

        psqlextra supports expressions in hstore fields, so casting
        all values to strings is a bad idea."""

        value = Field.get_prep_value(self, value)

        if isinstance(value, dict):
            prep_value = {}
            for key, val in value.items():
                if isinstance(val, Expression):
                    prep_value[key] = val
                elif val is not None:
                    prep_value[key] = str(val)
                else:
                    prep_value[key] = val

            value = prep_value

        if isinstance(value, list):
            value = [str(item) for item in value]

        return value

    def deconstruct(self):
        """Gets the values to pass to :see:__init__ when
        re-creating this object."""

        name, path, args, kwargs = super(
            HStoreField, self).deconstruct()

        if self.uniqueness is not None:
            kwargs['uniqueness'] = self.uniqueness

        if self.required is not None:
            kwargs['required'] = self.required

        return name, path, args, kwargs


from django.core.exceptions import SuspiciousOperation
from django.db.models.sql.compiler import SQLInsertCompiler, SQLUpdateCompiler

from psqlextra.expressions import HStoreValue

class PostgresReturningUpdateCompiler(SQLUpdateCompiler):
    """Compiler for SQL UPDATE statements that return
    the primary keys of the affected rows."""

    def execute_sql(self, _result_type):
        sql, params = self.as_sql()
        sql += self._form_returning()

        with self.connection.cursor() as cursor:
            cursor.execute(sql, params)
            primary_keys = cursor.fetchall()

        return primary_keys

    def as_sql(self):
        self._prepare_query_values()
        return super().as_sql()

    def _prepare_query_values(self):
        """Extra prep on query values by converting
        dictionaries into :see:HStoreValue expressions.

        This allows putting expressions in a dictionary.
        The :see:HStoreValue will take care of resolving
        the expressions inside the dictionary."""

        new_query_values = []
        for field, model, val in self.query.values:
            if isinstance(val, dict):
                val = HStoreValue(val)

            new_query_values.append((
                field,
                model,
                val
            ))

        self.query.values = new_query_values

    def _form_returning(self):
        """Builds the RETURNING part of the query."""

        qn = self.connection.ops.quote_name
        return ' RETURNING %s' % qn(self.query.model._meta.pk.attname)

class PostgresInsertCompiler(SQLInsertCompiler):
    """Compiler for SQL INSERT statements."""

    def __init__(self, *args, **kwargs):
        """Initializes a new instance of :see:PostgresInsertCompiler."""

        super().__init__(*args, **kwargs)
        self.qn = self.connection.ops.quote_name

    def as_sql(self, return_id=False):
        """Builds the SQL INSERT statement."""

        queries = [
            self._rewrite_insert(sql, params, return_id)
            for sql, params in super().as_sql()
        ]

        return queries

    def execute_sql(self, return_id=False):
        # execute all the generate queries
        with self.connection.cursor() as cursor:
            rows = []
            for sql, params in self.as_sql(return_id):
                cursor.execute(sql, params)
                rows.extend(cursor.fetchall())

        # create a mapping between column names and column value
        return [
            {
                column.name: row[column_index]
                for column_index, column in enumerate(cursor.description) if row
            }
            for row in rows
        ]

    def _rewrite_insert(self, sql, params, return_id=False):
        """Rewrites a formed SQL INSERT query to include
        the ON CONFLICT clause.

        Arguments:
            sql:
                The SQL INSERT query to rewrite.

            params:
                The parameters passed to the query.

            returning:
                What to put in the `RETURNING` clause
                of the resulting query.

        Returns:
            A tuple of the rewritten SQL query and new params.
        """

        returning = self.qn(self.query.model._meta.pk.attname) if return_id else '*'

        if self.query.conflict_action.value == 'UPDATE':
            return self._rewrite_insert_update(sql, params, returning)
        elif self.query.conflict_action.value == 'NOTHING':
            return self._rewrite_insert_nothing(sql, params, returning)

        raise SuspiciousOperation((
            '%s is not a valid conflict action, specify '
            'ConflictAction.UPDATE or ConflictAction.NOTHING.'
        ) % str(self.query.conflict_action))

    def _rewrite_insert_update(self, sql, params, returning):
        """Rewrites a formed SQL INSERT query to include
        the ON CONFLICT DO UPDATE clause."""

        update_columns = ', '.join([
            '{0} = EXCLUDED.{0}'.format(self.qn(field.column))
            for field in self.query.update_fields
        ])

        # build the conflict target, the columns to watch
        # for conflicts
        conflict_target = self._build_conflict_target()

        index_predicate = self.query.index_predicate

        sql_template = (
            '{insert} ON CONFLICT {conflict_target} DO UPDATE '
            'SET {update_columns} RETURNING {returning}'
        )

        if index_predicate:
            sql_template = (
                '{insert} ON CONFLICT {conflict_target} WHERE {index_predicate} DO UPDATE '
                'SET {update_columns} RETURNING {returning}'
            )

        return (
            sql_template.format(
                insert=sql,
                conflict_target=conflict_target,
                update_columns=update_columns,
                returning=returning,
                index_predicate=index_predicate,
            ),
            params
        )

    def _rewrite_insert_nothing(self, sql, params, returning):
        """Rewrites a formed SQL INSERT query to include
        the ON CONFLICT DO NOTHING clause."""

        # build the conflict target, the columns to watch
        # for conflicts
        conflict_target = self._build_conflict_target()

        where_clause = ' AND '.join([
            '{0} = %s'.format(self._format_field_name(field_name))
            for field_name in self.query.conflict_target
        ])

        where_clause_params = [
            self._format_field_value(field_name)
            for field_name in self.query.conflict_target
        ]

        params = params + tuple(where_clause_params)

        # this looks complicated, and it is, but it is for a reason... a normal
        # ON CONFLICT DO NOTHING doesn't return anything if the row already exists
        # so we do DO UPDATE instead that never executes to lock the row, and then
        # select from the table in case we're dealing with an existing row..
        return (
            (
                'WITH insdata AS ('
                '{insert} ON CONFLICT {conflict_target} DO UPDATE'
                ' SET {pk_column} = NULL WHERE FALSE RETURNING {returning})'
                ' SELECT * FROM insdata UNION ALL'
                ' SELECT {returning} FROM {table} WHERE {where_clause} LIMIT 1;'
            ).format(
                insert=sql,
                conflict_target=conflict_target,
                pk_column=self.qn(self.query.model._meta.pk.column),
                returning=returning,
                table=self.query.objs[0]._meta.db_table,
                where_clause=where_clause
            ),
            params
        )

    def _build_conflict_target(self):
        """Builds the `conflict_target` for the ON CONFLICT
        clause."""

        conflict_target = []

        if not isinstance(self.query.conflict_target, list):
            raise SuspiciousOperation((
                '%s is not a valid conflict target, specify '
                'a list of column names, or tuples with column '
                'names and hstore key.'
            ) % str(self.query.conflict_target))

        def _assert_valid_field(field_name):
            field_name = self._normalize_field_name(field_name)
            if self._get_model_field(field_name):
                return

            raise SuspiciousOperation((
                '%s is not a valid conflict target, specify '
                'a list of column names, or tuples with column '
                'names and hstore key.'
            ) % str(field_name))

        for field_name in self.query.conflict_target:
            _assert_valid_field(field_name)

            # special handling for hstore keys
            if isinstance(field_name, tuple):
                conflict_target.append(
                    '(%s->\'%s\')' % (
                        self._format_field_name(field_name),
                        field_name[1]
                    )
                )
            else:
                conflict_target.append(
                    self._format_field_name(field_name))

        return '(%s)' % ','.join(conflict_target)

    def _get_model_field(self, name: str):
        """Gets the field on a model with the specified name.

        Arguments:
            name:
                The name of the field to look for.

                This can be both the actual field name, or
                the name of the column, both will work :)

        Returns:
            The field with the specified name or None if
            no such field exists.
        """

        field_name = self._normalize_field_name(name)

        # 'pk' has special meaning and always refers to the primary
        # key of a model, we have to respect this de-facto standard behaviour
        if field_name == 'pk' and self.query.model._meta.pk:
            return self.query.model._meta.pk

        for field in self.query.model._meta.local_concrete_fields:
            if field.name == field_name or field.column == field_name:
                return field

        return None

    def _format_field_name(self, field_name) -> str:
        """Formats a field's name for usage in SQL.

        Arguments:
            field_name:
                The field name to format.

        Returns:
            The specified field name formatted for
            usage in SQL.
        """

        field = self._get_model_field(field_name)
        return self.qn(field.column)

    def _format_field_value(self, field_name) -> str:
        """Formats a field's value for usage in SQL.

        Arguments:
            field_name:
                The name of the field to format
                the value of.

        Returns:
            The field's value formatted for usage
            in SQL.
        """

        field_name = self._normalize_field_name(field_name)
        field = self._get_model_field(field_name)

        return SQLInsertCompiler.prepare_value(
            self,
            field,
            # Note: this deliberately doesn't use `pre_save_val` as we don't
            # want things like auto_now on DateTimeField (etc.) to change the
            # value. We rely on pre_save having already been done by the
            # underlying compiler so that things like FileField have already had
            # the opportunity to save out their data.
            getattr(self.query.objs[0], field.attname)
        )

    def _normalize_field_name(self, field_name) -> str:
        """Normalizes a field name into a string by
        extracting the field name if it was specified
        as a reference to a HStore key (as a tuple).

        Arguments:
            field_name:
                The field name to normalize.

        Returns:
            The normalized field name.
        """

        if isinstance(field_name, tuple):
            field_name, _ = field_name

        return field_name


from ..fields import HStoreField

class HStoreUniqueSchemaEditorMixin:
    sql_hstore_unique_create = (
        'CREATE UNIQUE INDEX IF NOT EXISTS '
        '{name} ON {table} '
        '({columns})'
    )

    sql_hstore_unique_rename = (
        'ALTER INDEX '
        '{old_name} '
        'RENAME TO '
        '{new_name}'
    )

    sql_hstore_unique_drop = (
        'DROP INDEX IF EXISTS {name}'
    )

    def create_model(self, model):
        """Ran when a new model is created."""

        for field in model._meta.local_fields:
            if not isinstance(field, HStoreField):
                continue

            self.add_field(model, field)

    def delete_model(self, model):
        """Ran when a model is being deleted."""

        for field in model._meta.local_fields:
            if not isinstance(field, HStoreField):
                continue

            self.remove_field(model, field)

    def alter_db_table(self, model, old_db_table, new_db_table):
        """Ran when the name of a model is changed."""

        for field in model._meta.local_fields:
            if not isinstance(field, HStoreField):
                continue

            for keys in self._iterate_uniqueness_keys(field):
                self._rename_hstore_unique(
                    old_db_table,
                    new_db_table,
                    field,
                    field,
                    keys
                )

    def add_field(self, model, field):
        """Ran when a field is added to a model."""

        for keys in self._iterate_uniqueness_keys(field):
            self._create_hstore_unique(
                model,
                field,
                keys
            )

    def remove_field(self, model, field):
        """Ran when a field is removed from a model."""

        for keys in self._iterate_uniqueness_keys(field):
            self._drop_hstore_unique(
                model,
                field,
                keys
            )

    def alter_field(self, model, old_field, new_field, strict=False):
        """Ran when the configuration on a field changed."""

        is_old_field_hstore = isinstance(old_field, HStoreField)
        is_new_field_hstore = isinstance(new_field, HStoreField)

        if not is_old_field_hstore and not is_new_field_hstore:
            return

        old_uniqueness = getattr(old_field, 'uniqueness', []) or []
        new_uniqueness = getattr(new_field, 'uniqueness', []) or []

        # handle field renames before moving on
        if str(old_field.column) != str(new_field.column):
            for keys in self._iterate_uniqueness_keys(old_field):
                self._rename_hstore_unique(
                    model._meta.db_table,
                    model._meta.db_table,
                    old_field,
                    new_field,
                    keys
                )

        # drop the indexes for keys that have been removed
        for keys in old_uniqueness:
            if keys not in new_uniqueness:
                self._drop_hstore_unique(
                    model,
                    old_field,
                    self._compose_keys(keys)
                )

        # create new indexes for keys that have been added
        for keys in new_uniqueness:
            if keys not in old_uniqueness:
                self._create_hstore_unique(
                    model,
                    new_field,
                    self._compose_keys(keys)
                )

    def _create_hstore_unique(self, model, field, keys):
        """Creates a UNIQUE constraint for the specified hstore keys."""

        name = self._unique_constraint_name(
            model._meta.db_table, field, keys)
        columns = [
            '(%s->\'%s\')' % (field.column, key)
            for key in keys
        ]
        sql = self.sql_hstore_unique_create.format(
            name=self.quote_name(name),
            table=self.quote_name(model._meta.db_table),
            columns=','.join(columns)
        )
        self.execute(sql)

    def _rename_hstore_unique(self, old_table_name, new_table_name,
                              old_field, new_field, keys):
        """Renames an existing UNIQUE constraint for the specified
        hstore keys."""

        old_name = self._unique_constraint_name(
            old_table_name, old_field, keys)
        new_name = self._unique_constraint_name(
            new_table_name, new_field, keys)

        sql = self.sql_hstore_unique_rename.format(
            old_name=self.quote_name(old_name),
            new_name=self.quote_name(new_name)
        )
        self.execute(sql)

    def _drop_hstore_unique(self, model, field, keys):
        """Drops a UNIQUE constraint for the specified hstore keys."""

        name = self._unique_constraint_name(
            model._meta.db_table, field, keys)
        sql = self.sql_hstore_unique_drop.format(name=self.quote_name(name))
        self.execute(sql)

    @staticmethod
    def _unique_constraint_name(table: str, field, keys):
        """Gets the name for a UNIQUE INDEX that applies
        to one or more keys in a hstore field.

        Arguments:
            table:
                The name of the table the field is
                a part of.

            field:
                The hstore field to create a
                UNIQUE INDEX for.

            key:
                The name of the hstore key
                to create the name for.

                This can also be a tuple
                of multiple names.

        Returns:
            The name for the UNIQUE index.
        """
        postfix = '_'.join(keys)
        return '{table}_{field}_unique_{postfix}'.format(
            table=table,
            field=field.column,
            postfix=postfix
        )

    def _iterate_uniqueness_keys(self, field):
        """Iterates over the keys marked as "unique"
        in the specified field.

        Arguments:
            field:
                The field of which key's to
                iterate over.
        """

        uniqueness = getattr(field, 'uniqueness', None)
        if not uniqueness:
            return

        for keys in uniqueness:
            composed_keys = self._compose_keys(keys)
            yield composed_keys

    @staticmethod
    def _compose_keys(constraint):
        """Turns a string into a list of string or
        returns it as a list."""
        if isinstance(constraint, str):
            return [constraint]

        return constraint


from typing import List, Tuple, Any

from django.db.models.sql.datastructures import Join

class ConditionalJoin(Join):
    """A custom JOIN statement that allows attaching
    extra conditions."""

    def __init__(self, *args, **kwargs):
        """Initializes a new instance of :see:ConditionalJoin."""

        super().__init__(*args, **kwargs)
        self.join_type = 'LEFT OUTER JOIN'
        self.extra_conditions = []

    def add_condition(self, field, value: Any) -> None:
        """Adds an extra condition to this join.

        Arguments:
            field:
                The field that the condition will apply to.

            value:
                The value to compare.
        """

        self.extra_conditions.append((field, value))

    def as_sql(self, compiler, connection) -> Tuple[str, List[Any]]:
        """Compiles this JOIN into a SQL string."""

        sql, params = super().as_sql(compiler, connection)
        qn = compiler.quote_name_unless_alias

        # generate the extra conditions
        extra_conditions = ' AND '.join([
            '{}.{} = %s'.format(
                qn(self.table_name),
                qn(field.column)
            )
            for field, value in self.extra_conditions
        ])

        # add to the existing params, so the connector will
        # actually nicely format the value for us
        for _, value in self.extra_conditions:
            params.append(value)

        # rewrite the sql to include the extra conditions
        rewritten_sql = sql.replace(')', ' AND {})'.format(extra_conditions))
        return rewritten_sql, params

    @classmethod
    def from_join(cls, join: Join) -> 'ConditionalJoin':
        """Creates a new :see:ConditionalJoin from the
        specified :see:Join object.

        Arguments:
            join:
                The :see:Join object to create the
                :see:ConditionalJoin object from.

        Returns:
            A :see:ConditionalJoin object created from
            the :see:Join object.
        """

        return cls(
            join.table_name,
            join.parent_alias,
            join.table_alias,
            join.join_type,
            join.join_field,
            join.nullable
        )


from __future__ import division, with_statement, print_function, absolute_import

import csv
import os.path
import math
import sys

import perf
import six
import statistics

NO_VERSION = "<not set>"

def format_result(bench):
    mean = bench.mean()
    if bench.get_nvalue() >= 2:
        args = bench.format_values((mean, bench.stdev()))
        return 'Mean +- std dev: %s +- %s' % args
    else:
        return bench.format_value(mean)

# A table of 95% confidence intervals for a two-tailed t distribution, as a
# function of the degrees of freedom. For larger degrees of freedom, we
# approximate. While this may look less elegant than simply calculating the
# critical value, those calculations suck. Look at
# http://www.math.unb.ca/~knight/utility/t-table.htm if you need more values.
_T_DIST_95_CONF_LEVELS = [0, 12.706, 4.303, 3.182, 2.776,
                          2.571, 2.447, 2.365, 2.306, 2.262,
                          2.228, 2.201, 2.179, 2.160, 2.145,
                          2.131, 2.120, 2.110, 2.101, 2.093,
                          2.086, 2.080, 2.074, 2.069, 2.064,
                          2.060, 2.056, 2.052, 2.048, 2.045,
                          2.042]

def tdist95conf_level(df):
    """Approximate the 95% confidence interval for Student's T distribution.

    Given the degrees of freedom, returns an approximation to the 95%
    confidence interval for the Student's T distribution.

    Args:
        df: An integer, the number of degrees of freedom.

    Returns:
        A float.
    """
    df = int(round(df))
    highest_table_df = len(_T_DIST_95_CONF_LEVELS)
    if df >= 200:
        return 1.960
    if df >= 100:
        return 1.984
    if df >= 80:
        return 1.990
    if df >= 60:
        return 2.000
    if df >= 50:
        return 2.009
    if df >= 40:
        return 2.021
    if df >= highest_table_df:
        return _T_DIST_95_CONF_LEVELS[highest_table_df - 1]
    return _T_DIST_95_CONF_LEVELS[df]

def pooled_sample_variance(sample1, sample2):
    """Find the pooled sample variance for two samples.

    Args:
        sample1: one sample.
        sample2: the other sample.

    Returns:
        Pooled sample variance, as a float.
    """
    deg_freedom = len(sample1) + len(sample2) - 2
    mean1 = statistics.mean(sample1)
    squares1 = ((x - mean1) ** 2 for x in sample1)
    mean2 = statistics.mean(sample2)
    squares2 = ((x - mean2) ** 2 for x in sample2)

    return (math.fsum(squares1) + math.fsum(squares2)) / float(deg_freedom)

def tscore(sample1, sample2):
    """Calculate a t-test score for the difference between two samples.

    Args:
        sample1: one sample.
        sample2: the other sample.

    Returns:
        The t-test score, as a float.
    """
    if len(sample1) != len(sample2):
        raise ValueError("different number of values")
    error = pooled_sample_variance(sample1, sample2) / len(sample1)
    diff = statistics.mean(sample1) - statistics.mean(sample2)
    return diff / math.sqrt(error * 2)

def is_significant(sample1, sample2):
    """Determine whether two samples differ significantly.

    This uses a Student's two-sample, two-tailed t-test with alpha=0.95.

    Args:
        sample1: one sample.
        sample2: the other sample.

    Returns:
        (significant, t_score) where significant is a bool indicating whether
        the two samples differ significantly; t_score is the score from the
        two-sample T test.
    """
    deg_freedom = len(sample1) + len(sample2) - 2
    critical_value = tdist95conf_level(deg_freedom)
    t_score = tscore(sample1, sample2)
    return (abs(t_score) >= critical_value, t_score)

def significant_msg(base, changed):
    if base.get_nvalue() < 2 or changed.get_nvalue() < 2:
        return "(benchmark only contains a single value)"

    avg_base = base.mean()
    avg_changed = changed.mean()

    msg = "Not significant"
    significant = False

    # Due to inherent measurement imprecisions, variations of less than 1%
    # are automatically considered insignificant. This helps present
    # a clear picture to the user.
    if abs(avg_base - avg_changed) > (avg_base + avg_changed) * 0.01:
        base_times = base.get_values()
        changed_times = changed.get_values()

        significant, t_score = is_significant(base_times, changed_times)
        if significant:
            msg = "Significant (t=%.2f)" % t_score

    return msg

def format_table(base_label, changed_label, results):
    table = [("Benchmark", base_label, changed_label, "Change", "Significance")]

    for (bench_name, result) in results:
        format_value = result.base.format_value
        avg_base = result.base.mean()
        avg_changed = result.changed.mean()
        delta_avg = quantity_delta(result.base, result.changed)
        msg = significant_msg(result.base, result.changed)
        table.append((bench_name,
                      # Limit the precision for conciseness in the table.
                      format_value(avg_base),
                      format_value(avg_changed),
                      delta_avg,
                      msg))

    # Columns with None values are skipped
    skipped_cols = set()
    col_widths = [0] * len(table[0])
    for row in table:
        for col, val in enumerate(row):
            if val is None:
                skipped_cols.add(col)
                continue
            col_widths[col] = max(col_widths[col], len(val))

    outside_line = "+"
    header_sep_line = "+"
    for col, width in enumerate(col_widths):
        if col in skipped_cols:
            continue
        width += 2  # Compensate for the left and right padding spaces.
        outside_line += "-" * width + "+"
        header_sep_line += "=" * width + "+"

    output = [outside_line]
    for row_i, row in enumerate(table):
        output_row = []
        for col_i, val in enumerate(row):
            if col_i in skipped_cols:
                continue
            output_row.append("| " + val.ljust(col_widths[col_i]) + " ")
        output.append("".join(output_row) + "|")
        if row_i > 0:
            output.append(outside_line)

    output.insert(2, "".join(header_sep_line))
    return "\n".join(output)

class BenchmarkResult(object):
    """An object representing data from a succesful benchmark run."""

    def __init__(self, base, changed):
        name = base.get_name()
        name2 = changed.get_name()
        if name2 != name:
            raise ValueError("not the same benchmark: %s != %s"
                             % (name, name2))

        if base.get_nvalue() != changed.get_nvalue():
            raise RuntimeError("base and changed don't have "
                               "the same number of values")

        self.base = base
        self.changed = changed

    def __str__(self):
        if self.base.get_nvalue() > 1:
            values = (self.base.mean(), self.base.stdev(),
                      self.changed.mean(), self.changed.stdev())
            text = "%s +- %s -> %s +- %s" % self.base.format_values(values)

            msg = significant_msg(self.base, self.changed)
            delta_avg = quantity_delta(self.base, self.changed)
            return ("Mean +- std dev: %s: %s\n%s"
                    % (text, delta_avg, msg))
        else:
            format_value = self.base.format_value
            base = self.base.mean()
            changed = self.changed.mean()
            delta_avg = quantity_delta(self.base, self.changed)
            return ("%s -> %s: %s"
                    % (format_value(base),
                       format_value(changed),
                       delta_avg))

def quantity_delta(base, changed):
    old = base.mean()
    new = changed.mean()
    is_time = (base.get_unit() == 'second')

    if old == 0 or new == 0:
        return "incomparable (one result was zero)"
    if new > old:
        if is_time:
            return "%.2fx slower" % (new / old)
        else:
            return "%.2fx larger" % (new / old)
    elif new < old:
        if is_time:
            return "%.2fx faster" % (old / new)
        else:
            return "%.2fx smaller" % (old / new)
    else:
        return "no change"

def display_suite_metadata(suite, title=None):
    metadata = suite.get_metadata()
    empty = True
    for key, fmt in (
        ('performance_version', "Performance version: %s"),
        ('python_version', "Python version: %s"),
        ('platform', "Report on %s"),
        ('cpu_count', "Number of logical CPUs: %s"),
    ):
        if key not in metadata:
            continue

        empty = False
        if title:
            print(title)
            print("=" * len(title))
            print()
            title = None

        text = fmt % metadata[key]
        print(text)

    dates = suite.get_dates()
    if dates:
        print("Start date: %s" % dates[0].isoformat(' '))
        print("End date: %s" % dates[1].isoformat(' '))
        empty = False

    if not empty:
        print()

def display_benchmark_suite(suite):
    display_suite_metadata(suite)

    for bench in suite.get_benchmarks():
        print("### %s ###" % bench.get_name())
        print(format_result(bench))
        print()

def cmd_show(options):
    suite = perf.BenchmarkSuite.load(options.filename)
    display_benchmark_suite(suite)

def get_labels(filename1, filename2):
    # Find a short label to identify two filenames:
    # the two labels must be different
    name1 = os.path.basename(filename1)
    name2 = os.path.basename(filename2)
    if name1 != name2:
        return (name1, name2)

    return (filename1, filename2)

def compare_results(options):
    base_label, changed_label = get_labels(options.baseline_filename,
                                           options.changed_filename)

    base_suite = perf.BenchmarkSuite.load(options.baseline_filename)
    changed_suite = perf.BenchmarkSuite.load(options.changed_filename)

    results = []
    common = set(base_suite.get_benchmark_names()) & set(
        changed_suite.get_benchmark_names())
    for name in sorted(common):
        base_bench = base_suite.get_benchmark(name)
        changed_bench = changed_suite.get_benchmark(name)
        result = BenchmarkResult(base_bench, changed_bench)
        results.append(result)

    hidden = []
    shown = []
    for result in results:
        name = result.base.get_name()

        significant = significant_msg(result.base, result.changed)
        if significant or options.verbose:
            shown.append((name, result))
        else:
            hidden.append((name, result))

    display_suite_metadata(base_suite, title=base_label)
    display_suite_metadata(changed_suite, title=changed_label)

    if options.output_style == "normal":
        for index, item in enumerate(shown):
            if index:
                print()
            name, result = item
            print("###", name, "###")
            print(result)

    elif options.output_style == "table":
        if shown:
            print(format_table(base_label, changed_label, shown))
    else:
        raise ValueError("Invalid output_style: %r" % options.output_style)

    if hidden:
        print()
        print("The following not significant results are hidden, "
              "use -v to show them:")
        print(", ".join(name for (name, result) in hidden) + ".")

    only_base = set(base_suite.get_benchmark_names()) - common
    if only_base:
        print()
        print("Skipped %s benchmarks only in %s: %s"
              % (len(only_base), base_label,
                 ', '.join(sorted(only_base))))

    only_changed = set(changed_suite.get_benchmark_names()) - common
    if only_changed:
        print()
        print("Skipped %s benchmarks only in %s: %s"
              % (len(only_changed), changed_label,
                 ', '.join(sorted(only_changed))))

    version1 = base_suite.get_metadata().get('performance_version', NO_VERSION)
    version2 = changed_suite.get_metadata().get('performance_version', NO_VERSION)
    if version1 != version2 or (version1 == version2 == NO_VERSION):
        print()
        print("ERROR: Performance versions are different: %s != %s"
              % (version1, version2))
        sys.exit(1)

    return results

def format_csv(value):
    abs_value = abs(value)
    # keep at least 3 significant digits, but also try to avoid too many zeros
    if abs_value >= 1.0:
        return "%.2f" % value
    elif abs_value >= 1e-3:
        return "%.5f" % value
    elif abs_value >= 1e-6:
        return "%.8f" % value
    else:
        return "%.11f" % value

def write_csv(results, filename):
    if six.PY3:
        fp = open(filename, "w", newline='', encoding='ascii')
    else:
        fp = open(filename, "wb")
    with fp:
        writer = csv.writer(fp)
        writer.writerow(['Benchmark', 'Base', 'Changed'])
        for result in results:
            name = result.base.get_name()
            base = result.base.mean()
            changed = result.changed.mean()
            row = [name, format_csv(base), format_csv(changed)]
            writer.writerow(row)

def cmd_compare(options):
    results = compare_results(options)

    if options.csv:
        write_csv(results, options.csv)


import collections
from collections import defaultdict
from fractions import Fraction

import perf
from six.moves import xrange

def topoSort(roots, getParents):
    """Return a topological sorting of nodes in a graph.

    roots - list of root nodes to search from
    getParents - function which returns the parents of a given node
    """

    results = []
    visited = set()

    # Use iterative version to avoid stack limits for large datasets
    stack = [(node, 0) for node in roots]
    while stack:
        current, state = stack.pop()
        if state == 0:
            # before recursing
            if current not in visited:
                visited.add(current)
                stack.append((current, 1))
                stack.extend((parent, 0) for parent in getParents(current))
        else:
            # after recursing
            assert(current in visited)
            results.append(current)
    return results

def getDamages(L, A, D, B, stab, te):
    x = (2 * L) // 5
    x = ((x + 2) * A * B) // (D * 50) + 2
    if stab:
        x += x // 2
    x = int(x * te)
    return [(x * z) // 255 for z in range(217, 256)]

def getCritDist(L, p, A1, A2, D1, D2, B, stab, te):
    p = min(p, Fraction(1))
    norm = getDamages(L, A1, D1, B, stab, te)
    crit = getDamages(L * 2, A2, D2, B, stab, te)

    dist = defaultdict(Fraction)
    for mult, vals in zip([1 - p, p], [norm, crit]):
        mult /= len(vals)
        for x in vals:
            dist[x] += mult
    return dist

def plus12(x):
    return x + x // 8

stats_t = collections.namedtuple('stats_t', ['atk', 'df', 'speed', 'spec'])
NOMODS = stats_t(0, 0, 0, 0)

fixeddata_t = collections.namedtuple(
    'fixeddata_t', ['maxhp', 'stats', 'lvl', 'badges', 'basespeed'])
halfstate_t = collections.namedtuple(
    'halfstate_t', ['fixed', 'hp', 'status', 'statmods', 'stats'])

def applyHPChange(hstate, change):
    hp = min(hstate.fixed.maxhp, max(0, hstate.hp + change))
    return hstate._replace(hp=hp)

def applyBadgeBoosts(badges, stats):
    return stats_t(*[(plus12(x) if b else x) for x, b in zip(stats, badges)])

attack_stats_t = collections.namedtuple(
    'attack_stats_t', ['power', 'isspec', 'stab', 'te', 'crit'])
attack_data = {
    'Ember': attack_stats_t(40, True, True, 0.5, False),
    'Dig': attack_stats_t(100, False, False, 1, False),
    'Slash': attack_stats_t(70, False, False, 1, True),
    'Water Gun': attack_stats_t(40, True, True, 2, False),
    'Bubblebeam': attack_stats_t(65, True, True, 2, False),
}

def _applyActionSide1(state, act):
    me, them, extra = state

    if act == 'Super Potion':
        me = applyHPChange(me, 50)
        return {(me, them, extra): Fraction(1)}

    mdata = attack_data[act]
    aind = 3 if mdata.isspec else 0
    dind = 3 if mdata.isspec else 1
    pdiv = 64 if mdata.crit else 512
    dmg_dist = getCritDist(me.fixed.lvl, Fraction(me.fixed.basespeed, pdiv),
                           me.stats[aind], me.fixed.stats[aind], them.stats[
                               dind], them.fixed.stats[dind],
                           mdata.power, mdata.stab, mdata.te)

    dist = defaultdict(Fraction)
    for dmg, p in dmg_dist.items():
        them2 = applyHPChange(them, -dmg)
        dist[me, them2, extra] += p
    return dist

def _applyAction(state, side, act):
    if side == 0:
        return _applyActionSide1(state, act)
    else:
        me, them, extra = state
        dist = _applyActionSide1((them, me, extra), act)
        return {(k[1], k[0], k[2]): v for k, v in dist.items()}

class Battle(object):

    def __init__(self):
        self.successors = {}
        self.min = defaultdict(float)
        self.max = defaultdict(lambda: 1.0)
        self.frozen = set()

        self.win = 4, True
        self.loss = 4, False
        self.max[self.loss] = 0.0
        self.min[self.win] = 1.0
        self.frozen.update([self.win, self.loss])

    def _getSuccessorsA(self, statep):
        st, state = statep
        for action in ['Dig', 'Super Potion']:
            yield (1, state, action)

    def _applyActionPair(self, state, side1, act1, side2, act2, dist, pmult):
        for newstate, p in _applyAction(state, side1, act1).items():
            if newstate[0].hp == 0:
                newstatep = self.loss
            elif newstate[1].hp == 0:
                newstatep = self.win
            else:
                newstatep = 2, newstate, side2, act2
            dist[newstatep] += p * pmult

    def _getSuccessorsB(self, statep):
        st, state, action = statep
        dist = defaultdict(Fraction)
        for eact, p in [('Water Gun', Fraction(64, 130)),
                        ('Bubblebeam', Fraction(66, 130))]:
            priority1 = state[0].stats.speed + \
                10000 * (action == 'Super Potion')
            priority2 = state[1].stats.speed + 10000 * (action == 'X Defend')

            if priority1 > priority2:
                self._applyActionPair(state, 0, action, 1, eact, dist, p)
            elif priority1 < priority2:
                self._applyActionPair(state, 1, eact, 0, action, dist, p)
            else:
                self._applyActionPair(state, 0, action, 1, eact, dist, p / 2)
                self._applyActionPair(state, 1, eact, 0, action, dist, p / 2)

        return {k: float(p) for k, p in dist.items() if p > 0}

    def _getSuccessorsC(self, statep):
        st, state, side, action = statep
        dist = defaultdict(Fraction)
        for newstate, p in _applyAction(state, side, action).items():
            if newstate[0].hp == 0:
                newstatep = self.loss
            elif newstate[1].hp == 0:
                newstatep = self.win
            else:
                newstatep = 0, newstate
            dist[newstatep] += p
        return {k: float(p) for k, p in dist.items() if p > 0}

    def getSuccessors(self, statep):
        try:
            return self.successors[statep]
        except KeyError:
            st = statep[0]
        if st == 0:
            result = list(self._getSuccessorsA(statep))
        else:
            if st == 1:
                dist = self._getSuccessorsB(statep)
            elif st == 2:
                dist = self._getSuccessorsC(statep)
            result = sorted(dist.items(), key=lambda t: (-t[1], t[0]))
        self.successors[statep] = result
        return result

    def getSuccessorsList(self, statep):
        if statep[0] == 4:
            return []
        temp = self.getSuccessors(statep)
        if statep[0] != 0:
            temp = list(zip(*temp))[0] if temp else []
        return temp

    def evaluate(self, tolerance=0.15):
        badges = 1, 0, 0, 0

        starfixed = fixeddata_t(59, stats_t(40, 44, 56, 50), 11, NOMODS, 115)
        starhalf = halfstate_t(starfixed, 59, 0, NOMODS,
                               stats_t(40, 44, 56, 50))
        charfixed = fixeddata_t(63, stats_t(39, 34, 46, 38), 26, badges, 65)
        charhalf = halfstate_t(charfixed, 63, 0, NOMODS, applyBadgeBoosts(
            badges, stats_t(39, 34, 46, 38)))
        initial_state = charhalf, starhalf, 0
        initial_statep = 0, initial_state

        dmin, dmax, frozen = self.min, self.max, self.frozen
        stateps = topoSort([initial_statep], self.getSuccessorsList)

        itercount = 0
        while dmax[initial_statep] - dmin[initial_statep] > tolerance:
            itercount += 1

            for sp in stateps:
                if sp in frozen:
                    continue

                if sp[0] == 0:
                    # choice node
                    dmin[sp] = max(dmin[sp2] for sp2 in self.getSuccessors(sp))
                    dmax[sp] = max(dmax[sp2] for sp2 in self.getSuccessors(sp))
                else:
                    dmin[sp] = sum(dmin[sp2] * p for sp2,
                                   p in self.getSuccessors(sp))
                    dmax[sp] = sum(dmax[sp2] * p for sp2,
                                   p in self.getSuccessors(sp))

                if dmin[sp] >= dmax[sp]:
                    dmax[sp] = dmin[sp] = (dmin[sp] + dmax[sp]) / 2
                    frozen.add(sp)
        return (dmax[initial_statep] + dmin[initial_statep]) / 2

def bench_mdp(loops):
    expected = 0.89873589887
    max_diff = 1e-6
    range_it = xrange(loops)

    t0 = perf.perf_counter()
    for _ in range_it:
        result = Battle().evaluate(0.192)
    dt = perf.perf_counter() - t0

    if abs(result - expected) > max_diff:
        raise Exception("invalid result: got %s, expected %s "
                        "(diff: %s, max diff: %s)"
                        % (result, expected, result - expected, max_diff))
    return dt

if __name__ == "__main__":
    runner = perf.Runner()
    runner.metadata['description'] = "MDP benchmark"
    runner.bench_time_func('mdp', bench_mdp)


"""Simple, brute-force N-Queens solver."""

import perf

__author__ = "collinwinter@google.com (Collin Winter)"

# Pure-Python implementation of itertools.permutations().
def permutations(iterable, r=None):
    """permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)"""
    pool = tuple(iterable)
    n = len(pool)
    if r is None:
        r = n
    indices = list(range(n))
    cycles = list(range(n - r + 1, n + 1))[::-1]
    yield tuple(pool[i] for i in indices[:r])
    while n:
        for i in reversed(range(r)):
            cycles[i] -= 1
            if cycles[i] == 0:
                indices[i:] = indices[i + 1:] + indices[i:i + 1]
                cycles[i] = n - i
            else:
                j = cycles[i]
                indices[i], indices[-j] = indices[-j], indices[i]
                yield tuple(pool[i] for i in indices[:r])
                break
        else:
            return

# From http://code.activestate.com/recipes/576647/
def n_queens(queen_count):
    """N-Queens solver.

    Args:
        queen_count: the number of queens to solve for. This is also the
            board size.

    Yields:
        Solutions to the problem. Each yielded value is looks like
        (3, 8, 2, 1, 4, ..., 6) where each number is the column position for the
        queen, and the index into the tuple indicates the row.
    """
    cols = range(queen_count)
    for vec in permutations(cols):
        if (queen_count == len(set(vec[i] + i for i in cols))
                        == len(set(vec[i] - i for i in cols))):
            yield vec

def bench_n_queens(queen_count):
    list(n_queens(queen_count))

if __name__ == "__main__":
    runner = perf.Runner()
    runner.metadata['description'] = "Simple, brute-force N-Queens solver"

    queen_count = 8
    runner.bench_func('nqueens', bench_n_queens, queen_count)


"""
Go board game
"""
import math
import random

import perf

SIZE = 9
GAMES = 200
KOMI = 7.5
EMPTY, WHITE, BLACK = 0, 1, 2
SHOW = {EMPTY: '.', WHITE: 'o', BLACK: 'x'}
PASS = -1
MAXMOVES = SIZE * SIZE * 3
TIMESTAMP = 0
MOVES = 0

def to_pos(x, y):
    return y * SIZE + x

def to_xy(pos):
    y, x = divmod(pos, SIZE)
    return x, y

class Square:

    def __init__(self, board, pos):
        self.board = board
        self.pos = pos
        self.timestamp = TIMESTAMP
        self.removestamp = TIMESTAMP
        self.zobrist_strings = [random.randrange(9223372036854775807)
                                for i in range(3)]

    def set_neighbours(self):
        x, y = self.pos % SIZE, self.pos // SIZE
        self.neighbours = []
        for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            newx, newy = x + dx, y + dy
            if 0 <= newx < SIZE and 0 <= newy < SIZE:
                self.neighbours.append(self.board.squares[to_pos(newx, newy)])

    def move(self, color):
        global TIMESTAMP, MOVES
        TIMESTAMP += 1
        MOVES += 1
        self.board.zobrist.update(self, color)
        self.color = color
        self.reference = self
        self.ledges = 0
        self.used = True
        for neighbour in self.neighbours:
            neighcolor = neighbour.color
            if neighcolor == EMPTY:
                self.ledges += 1
            else:
                neighbour_ref = neighbour.find(update=True)
                if neighcolor == color:
                    if neighbour_ref.reference.pos != self.pos:
                        self.ledges += neighbour_ref.ledges
                        neighbour_ref.reference = self
                    self.ledges -= 1
                else:
                    neighbour_ref.ledges -= 1
                    if neighbour_ref.ledges == 0:
                        neighbour.remove(neighbour_ref)
        self.board.zobrist.add()

    def remove(self, reference, update=True):
        self.board.zobrist.update(self, EMPTY)
        self.removestamp = TIMESTAMP
        if update:
            self.color = EMPTY
            self.board.emptyset.add(self.pos)
#            if color == BLACK:
#                self.board.black_dead += 1
#            else:
#                self.board.white_dead += 1
        for neighbour in self.neighbours:
            if neighbour.color != EMPTY and neighbour.removestamp != TIMESTAMP:
                neighbour_ref = neighbour.find(update)
                if neighbour_ref.pos == reference.pos:
                    neighbour.remove(reference, update)
                else:
                    if update:
                        neighbour_ref.ledges += 1

    def find(self, update=False):
        reference = self.reference
        if reference.pos != self.pos:
            reference = reference.find(update)
            if update:
                self.reference = reference
        return reference

    def __repr__(self):
        return repr(to_xy(self.pos))

class EmptySet:

    def __init__(self, board):
        self.board = board
        self.empties = list(range(SIZE * SIZE))
        self.empty_pos = list(range(SIZE * SIZE))

    def random_choice(self):
        choices = len(self.empties)
        while choices:
            i = int(random.random() * choices)
            pos = self.empties[i]
            if self.board.useful(pos):
                return pos
            choices -= 1
            self.set(i, self.empties[choices])
            self.set(choices, pos)
        return PASS

    def add(self, pos):
        self.empty_pos[pos] = len(self.empties)
        self.empties.append(pos)

    def remove(self, pos):
        self.set(self.empty_pos[pos], self.empties[len(self.empties) - 1])
        self.empties.pop()

    def set(self, i, pos):
        self.empties[i] = pos
        self.empty_pos[pos] = i

class ZobristHash:

    def __init__(self, board):
        self.board = board
        self.hash_set = set()
        self.hash = 0
        for square in self.board.squares:
            self.hash ^= square.zobrist_strings[EMPTY]
        self.hash_set.clear()
        self.hash_set.add(self.hash)

    def update(self, square, color):
        self.hash ^= square.zobrist_strings[square.color]
        self.hash ^= square.zobrist_strings[color]

    def add(self):
        self.hash_set.add(self.hash)

    def dupe(self):
        return self.hash in self.hash_set

class Board:

    def __init__(self):
        self.squares = [Square(self, pos) for pos in range(SIZE * SIZE)]
        for square in self.squares:
            square.set_neighbours()
        self.reset()

    def reset(self):
        for square in self.squares:
            square.color = EMPTY
            square.used = False
        self.emptyset = EmptySet(self)
        self.zobrist = ZobristHash(self)
        self.color = BLACK
        self.finished = False
        self.lastmove = -2
        self.history = []
        self.white_dead = 0
        self.black_dead = 0

    def move(self, pos):
        square = self.squares[pos]
        if pos != PASS:
            square.move(self.color)
            self.emptyset.remove(square.pos)
        elif self.lastmove == PASS:
            self.finished = True
        if self.color == BLACK:
            self.color = WHITE
        else:
            self.color = BLACK
        self.lastmove = pos
        self.history.append(pos)

    def random_move(self):
        return self.emptyset.random_choice()

    def useful_fast(self, square):
        if not square.used:
            for neighbour in square.neighbours:
                if neighbour.color == EMPTY:
                    return True
        return False

    def useful(self, pos):
        global TIMESTAMP
        TIMESTAMP += 1
        square = self.squares[pos]
        if self.useful_fast(square):
            return True
        old_hash = self.zobrist.hash
        self.zobrist.update(square, self.color)
        empties = opps = weak_opps = neighs = weak_neighs = 0
        for neighbour in square.neighbours:
            neighcolor = neighbour.color
            if neighcolor == EMPTY:
                empties += 1
                continue
            neighbour_ref = neighbour.find()
            if neighbour_ref.timestamp != TIMESTAMP:
                if neighcolor == self.color:
                    neighs += 1
                else:
                    opps += 1
                neighbour_ref.timestamp = TIMESTAMP
                neighbour_ref.temp_ledges = neighbour_ref.ledges
            neighbour_ref.temp_ledges -= 1
            if neighbour_ref.temp_ledges == 0:
                if neighcolor == self.color:
                    weak_neighs += 1
                else:
                    weak_opps += 1
                    neighbour_ref.remove(neighbour_ref, update=False)
        dupe = self.zobrist.dupe()
        self.zobrist.hash = old_hash
        strong_neighs = neighs - weak_neighs
        strong_opps = opps - weak_opps
        return not dupe and \
            (empties or weak_opps or (strong_neighs and (strong_opps or weak_neighs)))

    def useful_moves(self):
        return [pos for pos in self.emptyset.empties if self.useful(pos)]

    def replay(self, history):
        for pos in history:
            self.move(pos)

    def score(self, color):
        if color == WHITE:
            count = KOMI + self.black_dead
        else:
            count = self.white_dead
        for square in self.squares:
            squarecolor = square.color
            if squarecolor == color:
                count += 1
            elif squarecolor == EMPTY:
                surround = 0
                for neighbour in square.neighbours:
                    if neighbour.color == color:
                        surround += 1
                if surround == len(square.neighbours):
                    count += 1
        return count

    def check(self):
        for square in self.squares:
            if square.color == EMPTY:
                continue

            members1 = set([square])
            changed = True
            while changed:
                changed = False
                for member in members1.copy():
                    for neighbour in member.neighbours:
                        if neighbour.color == square.color and neighbour not in members1:
                            changed = True
                            members1.add(neighbour)
            ledges1 = 0
            for member in members1:
                for neighbour in member.neighbours:
                    if neighbour.color == EMPTY:
                        ledges1 += 1

            root = square.find()

            # print 'members1', square, root, members1
            # print 'ledges1', square, ledges1

            members2 = set()
            for square2 in self.squares:
                if square2.color != EMPTY and square2.find() == root:
                    members2.add(square2)

            ledges2 = root.ledges
            # print 'members2', square, root, members1
            # print 'ledges2', square, ledges2

            assert members1 == members2
            assert ledges1 == ledges2, ('ledges differ at %r: %d %d' % (
                square, ledges1, ledges2))

            set(self.emptyset.empties)

            empties2 = set()
            for square in self.squares:
                if square.color == EMPTY:
                    empties2.add(square.pos)

    def __repr__(self):
        result = []
        for y in range(SIZE):
            start = to_pos(0, y)
            result.append(''.join(
                [SHOW[square.color] + ' ' for square in self.squares[start:start + SIZE]]))
        return '\n'.join(result)

class UCTNode:

    def __init__(self):
        self.bestchild = None
        self.pos = -1
        self.wins = 0
        self.losses = 0
        self.pos_child = [None for x in range(SIZE * SIZE)]
        self.parent = None

    def play(self, board):
        """ uct tree search """
        color = board.color
        node = self
        path = [node]
        while True:
            pos = node.select(board)
            if pos == PASS:
                break
            board.move(pos)
            child = node.pos_child[pos]
            if not child:
                child = node.pos_child[pos] = UCTNode()
                child.unexplored = board.useful_moves()
                child.pos = pos
                child.parent = node
                path.append(child)
                break
            path.append(child)
            node = child
        self.random_playout(board)
        self.update_path(board, color, path)

    def select(self, board):
        """ select move; unexplored children first, then according to uct value """
        if self.unexplored:
            i = random.randrange(len(self.unexplored))
            pos = self.unexplored[i]
            self.unexplored[i] = self.unexplored[len(self.unexplored) - 1]
            self.unexplored.pop()
            return pos
        elif self.bestchild:
            return self.bestchild.pos
        else:
            return PASS

    def random_playout(self, board):
        """ random play until both players pass """
        for x in range(MAXMOVES):  # XXX while not self.finished?
            if board.finished:
                break
            board.move(board.random_move())

    def update_path(self, board, color, path):
        """ update win/loss count along path """
        wins = board.score(BLACK) >= board.score(WHITE)
        for node in path:
            if color == BLACK:
                color = WHITE
            else:
                color = BLACK
            if wins == (color == BLACK):
                node.wins += 1
            else:
                node.losses += 1
            if node.parent:
                node.parent.bestchild = node.parent.best_child()

    def score(self):
        winrate = self.wins / float(self.wins + self.losses)
        parentvisits = self.parent.wins + self.parent.losses
        if not parentvisits:
            return winrate
        nodevisits = self.wins + self.losses
        return winrate + math.sqrt((math.log(parentvisits)) / (5 * nodevisits))

    def best_child(self):
        maxscore = -1
        maxchild = None
        for child in self.pos_child:
            if child and child.score() > maxscore:
                maxchild = child
                maxscore = child.score()
        return maxchild

    def best_visited(self):
        maxvisits = -1
        maxchild = None
        for child in self.pos_child:
            #            if child:
            # print to_xy(child.pos), child.wins, child.losses, child.score()
            if child and (child.wins + child.losses) > maxvisits:
                maxvisits, maxchild = (child.wins + child.losses), child
        return maxchild

# def user_move(board):
#     while True:
#         text = six.moves.input('?').strip()
#         if text == 'p':
#             return PASS
#         if text == 'q':
#             raise EOFError
#         try:
#             x, y = [int(i) for i in text.split()]
#         except ValueError:
#             continue
#         if not (0 <= x < SIZE and 0 <= y < SIZE):
#             continue
#         pos = to_pos(x, y)
#         if board.useful(pos):
#             return pos

def computer_move(board):
    pos = board.random_move()
    if pos == PASS:
        return PASS
    tree = UCTNode()
    tree.unexplored = board.useful_moves()
    nboard = Board()
    for game in range(GAMES):
        node = tree
        nboard.reset()
        nboard.replay(board.history)
        node.play(nboard)
    return tree.best_visited().pos

def versus_cpu():
    random.seed(1)
    board = Board()
    return computer_move(board)

if __name__ == "__main__":
    kw = {}
    if perf.python_has_jit():
        # PyPy needs to compute more warmup values to warmup its JIT
        kw['warmups'] = 50
    runner = perf.Runner(**kw)
    runner.metadata['description'] = "Test the performance of the Go benchmark"
    runner.bench_func('go', versus_cpu)


from __future__ import division, with_statement, print_function, absolute_import

import logging

from performance.run import run_perf_script

# Benchmark groups. The "default" group is what's run if no -b option is
# specified.
DEFAULT_GROUP = [
    '2to3',
    # Issue #47: Chameleon doesn't support Python 3.8 yet
    # 'chameleon',
    'chaos',
    'crypto_pyaes',
    'deltablue',
    'django_template',
    'dulwich_log',
    'fannkuch',
    'float',
    # Issue #46: Genshi doesn't support Python 3.8 yet
    #'genshi',
    'go',
    'hexiom',
    'hg_startup',
    'html5lib',
    'json_dumps',
    'json_loads',
    'logging',
    'mako',
    'meteor_contest',
    'nbody',
    'nqueens',
    'pathlib',
    'pickle',
    'pickle_dict',
    'pickle_list',
    'pickle_pure_python',
    'pidigits',
    'pyflate',
    'python_startup',
    'python_startup_no_site',
    'raytrace',
    'regex_compile',
    'regex_dna',
    'regex_effbot',
    'regex_v8',
    'richards',
    'scimark',
    'spambayes',
    'spectral_norm',
    'sqlalchemy_declarative',
    'sqlalchemy_imperative',
    'sqlite_synth',
    'sympy',
    'telco',
    'tornado_http',
    'unpack_sequence',
    'unpickle',
    'unpickle_list',
    'unpickle_pure_python',
    'xml_etree',
]

BENCH_GROUPS = {
    # get_benchmarks() creates an "all" group which includes every benchmark
    # performance knows about.
    "default": DEFAULT_GROUP,
    "startup": ["normal_startup", "startup_nosite",
                "hg_startup"],
    "regex": ["regex_v8", "regex_effbot", "regex_compile",
              "regex_dna"],
    "serialize": ["pickle_pure_python", "unpickle_pure_python",  # Not for Python 3
                  "pickle", "unpickle",
                  "xml_etree",
                  "json_dumps", "json_loads"],
    "apps": ["2to3", "chameleon", "html5lib",
             "spambayes", "tornado_http"],
    "math": ["float", "nbody", "pidigits"],
    "template": ["django_template", "mako"],
}

def python2_only(func):
    func._python2_only = True
    return func

def BM_2to3(python, options):
    return run_perf_script(python, options, "2to3")

@python2_only
def BM_hg_startup(python, options):
    return run_perf_script(python, options, "hg_startup")

def BM_Chameleon(python, options):
    return run_perf_script(python, options, "chameleon")

def BM_Tornado_Http(python, options):
    return run_perf_script(python, options, "tornado_http")

def BM_Django_Template(python, options):
    return run_perf_script(python, options, "django_template")

def BM_Float(python, options):
    return run_perf_script(python, options, "float")

def BM_mako(python, options):
    return run_perf_script(python, options, "mako")

def BM_pathlib(python, options):
    return run_perf_script(python, options, "pathlib")

def pickle_benchmark(python, options, *extra_args):
    return run_perf_script(python, options, "pickle",
                           extra_args=list(extra_args))

def BM_pickle(python, options):
    return pickle_benchmark(python, options, "pickle")

def BM_unpickle(python, options):
    return pickle_benchmark(python, options, "unpickle")

def BM_pickle_list(python, options):
    return pickle_benchmark(python, options, "pickle_list")

def BM_pickle_dict(python, options):
    return pickle_benchmark(python, options, "pickle_dict")

def BM_unpickle_list(python, options):
    return pickle_benchmark(python, options, "unpickle_list")

def BM_pickle_pure_python(python, options):
    return pickle_benchmark(python, options, "--pure-python", "pickle")

def BM_unpickle_pure_python(python, options):
    return pickle_benchmark(python, options, "--pure-python", "unpickle")

def BM_xml_etree(python, options):
    return run_perf_script(python, options, "xml_etree")

def BM_json_loads(python, options):
    return run_perf_script(python, options, "json_loads")

def BM_json_dumps(python, options):
    return run_perf_script(python, options, "json_dumps")

def BM_NQueens(python, options):
    return run_perf_script(python, options, "nqueens")

def BM_Chaos(python, options):
    return run_perf_script(python, options, "chaos")

def BM_Fannkuch(python, options):
    return run_perf_script(python, options, "fannkuch")

def BM_Go(python, options):
    return run_perf_script(python, options, "go")

def BM_Meteor_Contest(python, options):
    return run_perf_script(python, options, "meteor_contest")

def BM_Spectral_Norm(python, options):
    return run_perf_script(python, options, "spectral_norm")

def BM_Telco(python, options):
    return run_perf_script(python, options, "telco")

def BM_hexiom(python, options):
    return run_perf_script(python, options, "hexiom")

def BM_raytrace(python, options):
    return run_perf_script(python, options, "raytrace")

def BM_logging(python, options):
    return run_perf_script(python, options, "logging")

def BM_python_startup(python, options):
    return run_perf_script(python, options, "python_startup")

def BM_python_startup_no_site(python, options):
    return run_perf_script(python, options, "python_startup",
                           extra_args=["--no-site"])

def BM_regex_v8(python, options):
    return run_perf_script(python, options, "regex_v8")

def BM_regex_effbot(python, options):
    return run_perf_script(python, options, "regex_effbot")

def BM_regex_compile(python, options):
    return run_perf_script(python, options, "regex_compile")

def BM_regex_dna(python, options):
    return run_perf_script(python, options, "regex_dna")

def BM_unpack_sequence(python, options):
    return run_perf_script(python, options, "unpack_sequence")

def BM_nbody(python, options):
    return run_perf_script(python, options, "nbody")

@python2_only
def BM_spambayes(python, options):
    return run_perf_script(python, options, "spambayes")

def BM_html5lib(python, options):
    return run_perf_script(python, options, "html5lib")

def BM_richards(python, options):
    return run_perf_script(python, options, "richards")

def BM_pidigits(python, options):
    return run_perf_script(python, options, "pidigits")

def BM_crypto_pyaes(python, options):
    return run_perf_script(python, options, "crypto_pyaes")

def BM_sympy(python, options):
    return run_perf_script(python, options, "sympy")

def BM_deltablue(python, options):
    return run_perf_script(python, options, "deltablue")

def BM_scimark(python, options):
    return run_perf_script(python, options, "scimark")

def BM_dulwich_log(python, options):
    return run_perf_script(python, options, "dulwich_log")

@python2_only
def BM_pyflate(python, options):
    return run_perf_script(python, options, "pyflate")

def BM_sqlite_synth(python, options):
    return run_perf_script(python, options, "sqlite_synth")

def BM_genshi(python, options):
    return run_perf_script(python, options, "genshi")

def BM_sqlalchemy_declarative(python, options):
    return run_perf_script(python, options, "sqlalchemy_declarative")

def BM_sqlalchemy_imperative(python, options):
    return run_perf_script(python, options, "sqlalchemy_imperative")

def BM_mdp(python, options):
    return run_perf_script(python, options, "mdp")

# End benchmarks, begin main entry point support.

def get_benchmarks():
    bench_funcs = dict((name[3:].lower(), func)
                       for name, func in globals().items()
                       if name.startswith("BM_"))

    bench_groups = BENCH_GROUPS.copy()

    # Calculate set of 2-and-3 compatible benchmarks.
    group2n3 = bench_groups["2n3"] = []
    for bm, func in bench_funcs.items():
        if not getattr(func, '_python2_only', False):
            group2n3.append(bm)

    # create the 'all' group
    bench_groups["all"] = sorted(bench_funcs)

    return (bench_funcs, bench_groups)

def filter_benchmarks(benchmarks, bench_funcs, base_ver):
    """Filters out benchmarks not supported by both Pythons.

    Args:
        benchmarks: a set() of benchmark names
        bench_funcs: dict mapping benchmark names to functions
        python: the interpereter commands (as lists)

    Returns:
        The filtered set of benchmark names
    """
    for bm in list(benchmarks):
        func = bench_funcs[bm]
        if getattr(func, '_python2_only', False) and (3, 0) <= base_ver:
            benchmarks.discard(bm)
            logging.info("Skipping Python2-only benchmark %s; "
                         "not compatible with Python %s" % (bm, base_ver))
            continue
    return benchmarks

def expand_benchmark_name(bm_name, bench_groups):
    """Recursively expand name benchmark names.

    Args:
        bm_name: string naming a benchmark or benchmark group.

    Yields:
        Names of actual benchmarks, with all group names fully expanded.
    """
    expansion = bench_groups.get(bm_name)
    if expansion:
        for name in expansion:
            for name in expand_benchmark_name(name, bench_groups):
                yield name
    else:
        yield bm_name

def select_benchmarks(benchmarks, bench_groups):
    legal_benchmarks = bench_groups["all"]
    benchmarks = benchmarks.split(",")
    positive_benchmarks = set(bm.lower()
                              for bm in benchmarks
                              if bm and not bm.startswith("-"))
    negative_benchmarks = set(bm[1:].lower()
                              for bm in benchmarks
                              if bm and bm.startswith("-"))

    should_run = set()
    if not positive_benchmarks:
        should_run = set(expand_benchmark_name("default", bench_groups))

    for name in positive_benchmarks:
        for bm in expand_benchmark_name(name, bench_groups):
            if bm not in legal_benchmarks:
                logging.warning("No benchmark named %s", bm)
            else:
                should_run.add(bm)

    for bm in negative_benchmarks:
        if bm in bench_groups:
            raise ValueError("Negative groups not supported: -%s" % bm)
        elif bm not in legal_benchmarks:
            logging.warning("No benchmark named %s", bm)
        else:
            should_run.remove(bm)
    return should_run


"""Benchmarks for Python's regex engine.

These are some of the original benchmarks used to tune Python's regex engine
in 2000 written by Fredrik Lundh. Retreived from
http://mail.python.org/pipermail/python-dev/2000-August/007797.html and
integrated into Unladen Swallow's perf.py in 2009 by David Laing.

These benchmarks are of interest since they helped to guide the original
optimization of the sre engine, and we shouldn't necessarily ignore them just
because they're "old".
"""

# Python imports
import re

# Local imports
import perf
from six.moves import xrange

USE_BYTES_IN_PY3K = False

def re_compile(s):
    if USE_BYTES_IN_PY3K:
        return re.compile(s.encode('latin1'))
    else:
        return re.compile(s)

# These are the regular expressions to be tested. These sync up,
# index-for-index with the list of strings generated by gen_string_table()
# below.

def gen_regex_table():
    return [
        re_compile('Python|Perl'),
        re_compile('Python|Perl'),
        re_compile('(Python|Perl)'),
        re_compile('(?:Python|Perl)'),
        re_compile('Python'),
        re_compile('Python'),
        re_compile('.*Python'),
        re_compile('.*Python.*'),
        re_compile('.*(Python)'),
        re_compile('.*(?:Python)'),
        re_compile('Python|Perl|Tcl'),
        re_compile('Python|Perl|Tcl'),
        re_compile('(Python|Perl|Tcl)'),
        re_compile('(?:Python|Perl|Tcl)'),
        re_compile('(Python)\\1'),
        re_compile('(Python)\\1'),
        re_compile('([0a-z][a-z0-9]*,)+'),
        re_compile('(?:[0a-z][a-z0-9]*,)+'),
        re_compile('([a-z][a-z0-9]*,)+'),
        re_compile('(?:[a-z][a-z0-9]*,)+'),
        re_compile('.*P.*y.*t.*h.*o.*n.*')]

def gen_string_table(n):
    """Generates the list of strings that will be used in the benchmarks.

    All strings have repeated prefixes and suffices, and n specifies the
    number of repetitions.
    """
    strings = []

    def append(s):
        if USE_BYTES_IN_PY3K:
            strings.append(s.encode('latin1'))
        else:
            strings.append(s)
    append('-' * n + 'Perl' + '-' * n)
    append('P' * n + 'Perl' + 'P' * n)
    append('-' * n + 'Perl' + '-' * n)
    append('-' * n + 'Perl' + '-' * n)
    append('-' * n + 'Python' + '-' * n)
    append('P' * n + 'Python' + 'P' * n)
    append('-' * n + 'Python' + '-' * n)
    append('-' * n + 'Python' + '-' * n)
    append('-' * n + 'Python' + '-' * n)
    append('-' * n + 'Python' + '-' * n)
    append('-' * n + 'Perl' + '-' * n)
    append('P' * n + 'Perl' + 'P' * n)
    append('-' * n + 'Perl' + '-' * n)
    append('-' * n + 'Perl' + '-' * n)
    append('-' * n + 'PythonPython' + '-' * n)
    append('P' * n + 'PythonPython' + 'P' * n)
    append('-' * n + 'a5,b7,c9,' + '-' * n)
    append('-' * n + 'a5,b7,c9,' + '-' * n)
    append('-' * n + 'a5,b7,c9,' + '-' * n)
    append('-' * n + 'a5,b7,c9,' + '-' * n)
    append('-' * n + 'Python' + '-' * n)
    return strings

def init_benchmarks(n_values=None):
    """Initialize the strings we'll run the regexes against.

    The strings used in the benchmark are prefixed and suffixed by
    strings that are repeated n times.

    The sequence n_values contains the values for n.
    If n_values is None the values of n from the original benchmark
    are used.

    The generated list of strings is cached in the string_tables
    variable, which is indexed by n.

    Returns:
    A list of string prefix/suffix lengths.
    """

    if n_values is None:
        n_values = (0, 5, 50, 250, 1000, 5000, 10000)

    string_tables = {n: gen_string_table(n) for n in n_values}
    regexs = gen_regex_table()

    data = []
    for n in n_values:
        for id in xrange(len(regexs)):
            regex = regexs[id]
            string = string_tables[n][id]
            data.append((regex, string))
    return data

def bench_regex_effbot(loops):
    if bench_regex_effbot.data is None:
        bench_regex_effbot.data = init_benchmarks()
    data = bench_regex_effbot.data

    range_it = xrange(loops)
    search = re.search
    t0 = perf.perf_counter()

    for _ in range_it:
        # Runs all of the benchmarks for a given value of n.
        for regex, string in data:
            # search 10 times
            search(regex, string)
            search(regex, string)
            search(regex, string)
            search(regex, string)
            search(regex, string)
            search(regex, string)
            search(regex, string)
            search(regex, string)
            search(regex, string)
            search(regex, string)

    return perf.perf_counter() - t0

# cached data, generated at the first call
bench_regex_effbot.data = None

def add_cmdline_args(cmd, args):
    if args.force_bytes:
        cmd.append("--force_bytes")

if __name__ == '__main__':
    runner = perf.Runner(add_cmdline_args=add_cmdline_args)
    runner.metadata['description'] = ("Test the performance of regexps "
                                      "using Fredik Lundh's benchmarks.")
    runner.argparser.add_argument("-B", "--force_bytes", action="store_true",
                                  help="Force testing bytes regexps "
                                       "under 3.x.")
    options = runner.parse_args()
    if options.force_bytes:
        USE_BYTES_IN_PY3K = True

    runner.bench_time_func('regex_effbot', bench_regex_effbot,
                           inner_loops=10)