import functools
import itertools
import json
import math
import tempfile
import time
from collections.abc import Iterator
from copy import deepcopy
from operator import contains, itemgetter, methodcaller, ne
from typing import Dict, Generator, List, TYPE_CHECKING, Tuple

import gevent

if TYPE_CHECKING:
    from locust.runners import WorkerNode

class UsersDispatcher(Iterator):
    """
    Iterator that dispatches the users to the workers.
    The users already running on the workers are also taken into
    account.

    The dispatcher waits an appropriate amount of time between each iteration
    in order for the spawn rate to be respected whether running in
    local or distributed mode.

    The spawn rate is only applied when additional users are needed.
    Hence, if the desired user count contains less users than what is currently running,
    the dispatcher won't wait and will only run for
    one iteration. The rationale for not stopping users at the spawn rate
    is that stopping them is a blocking operation, especially when
    a stop timeout is specified. When a stop timeout is specified combined with users having long-running tasks,
    attempting to stop the users at a spawn rate will lead to weird behaviours (users being killed even though the
    stop timeout is not reached yet).

    The terminology used in the users dispatcher is:
      - Dispatch cycle
            A dispatch cycle corresponds to a ramp-up from start to finish. So,
            going from 10 to 100 users with a spawn rate of 1/s corresponds to one
            dispatch cycle. An instance of the `UsersDispatcher` class "lives" for
            one dispatch cycle only.
      - Dispatch iteration
            A dispatch cycle contains one or more dispatch iterations. In the previous example
            of going from 10 to 100 users with a spawn rate of 1/s, there are 100 dispatch iterations.
            That is, from 10 to 11 users is a dispatch iteration, from 12 to 13 is another, and so on.
            If the spawn rate were to be 2/s, then there would be 50 dispatch iterations for this dispatch cycle.
            For a more extreme case with a spawn rate of 120/s, there would be only a single dispatch iteration
            from 10 to 100.
    """

    def __init__(self, worker_nodes: "List[WorkerNode]", user_classes_count: Dict[str, int], spawn_rate: float):
        """
        :param worker_nodes: List of worker nodes
        :param user_classes_count: Desired number of users for each class
        :param spawn_rate: The spawn rate
        """
        # NOTE: We use "sorted" in some places in this module. It is done to ensure repeatable behaviour.
        #       This is especially important when iterating over a dictionary which, prior to py3.7, was
        #       completely unordered. For >=Py3.7, a dictionary keeps the insertion order. Even then,
        #       it is safer to sort the keys when repeatable behaviour is required.
        self._worker_nodes = sorted(worker_nodes, key=lambda w: w.id)

        self._user_classes_count = user_classes_count

        self._user_classes = sorted(user_classes_count.keys())

        # Represents the already running users among the workers at the start of this dispatch cycle
        self._initial_dispatched_users = {
            worker_node.id: {
                user_class: worker_node.user_classes_count.get(user_class, 0)
                for user_class in self._user_classes_count.keys()
            }
            for worker_node in self._worker_nodes
        }

        self._desired_user_count = sum(self._user_classes_count.values())

        self._spawn_rate = spawn_rate

        self._desired_users_assigned_to_workers = _WorkersUsersAssignor(
            user_classes_count, worker_nodes
        ).desired_users_assigned_to_workers

        # This represents the desired users distribution minus the already running users among the workers.
        # The values inside this dictionary are updated during the current dispatch cycle. For example,
        # if we dispatch 1 user of UserClass1 to worker 1, then we will decrement by 1 the user count
        # for UserClass1 of worker 1. Naturally, the current dispatch cycle is done once all the values
        # reach zero.
        self._users_left_to_assigned = {
            worker_node.id: {
                user_class: max(
                    0,
                    self._desired_users_assigned_to_workers[worker_node.id][user_class]
                    - self._initial_dispatched_users[worker_node.id][user_class],
                )
                for user_class in self._user_classes_count.keys()
            }
            for worker_node in self._worker_nodes
        }

        self._user_count_per_dispatch = max(1, math.floor(self._spawn_rate))

        self._wait_between_dispatch = self._user_count_per_dispatch / self._spawn_rate

        # We use deepcopy because we will update the values inside `dispatched_users`
        # to keep track of the number of dispatched users for the current dispatch cycle.
        # It is essentially the same thing as for the `effective_assigned_users` dictionary,
        # but in reverse.
        self._dispatched_users = deepcopy(self._initial_dispatched_users)

        # Initialize the generator that is used in `__next__`
        self._dispatcher_generator = self._dispatcher()

        self._iteration = 0

        self._workers_desired_user_count = {
            worker_node_id: sum(desired_users_on_worker.values())
            for worker_node_id, desired_users_on_worker in self._desired_users_assigned_to_workers.items()
        }

    def __next__(self) -> Dict[str, Dict[str, int]]:
        self._iteration += 1
        return next(self._dispatcher_generator)

    def _dispatcher(self) -> Generator[Dict[str, Dict[str, int]], None, None]:
        """Main iterator logic for dispatching users during this dispatch cycle"""
        if self._desired_users_assignment_can_be_obtained_in_a_single_dispatch_iteration:
            yield self._desired_users_assigned_to_workers

        else:
            while not self._all_users_have_been_dispatched:
                ts = time.perf_counter()
                yield self._users_to_dispatch_for_current_iteration()
                if not self._all_users_have_been_dispatched:
                    delta = time.perf_counter() - ts
                    sleep_duration = max(0.0, self._wait_between_dispatch - delta)
                    assert sleep_duration <= 10, sleep_duration
                    gevent.sleep(sleep_duration)

    @property
    def _desired_users_assignment_can_be_obtained_in_a_single_dispatch_iteration(self) -> bool:
        if self._dispatched_user_count >= self._desired_user_count:
            # There is already more users running in total than the total desired user count
            return True

        if self._dispatched_user_count + self._user_count_per_dispatch > self._desired_user_count:
            # The spawn rate greater than the remaining users to dispatch
            return True

        # Workers having already more users than desired will show up zero
        # users left to dispatch in the following dictionary. And this, even
        # if these workers are missing users in one or more user classes.
        user_classes_count_left_to_dispatch_excluding_excess_users = {
            user_class: max(
                0,
                sum(map(itemgetter(user_class), self._desired_users_assigned_to_workers.values()))
                - self._dispatched_user_class_count(user_class),
            )
            for user_class in self._user_classes_count.keys()
        }

        # This condition is to cover a corner case for which there exists no dispatch solution that won't
        # violate the following constraints:
        #   - No worker run excess users at any point (except for the possible excess users already running)
        #   - No worker run excess users of any user class at any point (except for the
        #     possible excess users already running)
        #   - The total user count is never exceeded (except for the possible excess users already running)
        # In a situation like this, we have no choice but to immediately dispatch the final users immediately.
        workers_user_count = self._workers_user_count
        if (
            sum(
                1
                for user_class, user_class_count_left_to_dispatch in user_classes_count_left_to_dispatch_excluding_excess_users.items()
                if user_class_count_left_to_dispatch != 0
                for worker_node in self._worker_nodes
                if (
                    workers_user_count[worker_node.id]
                    < sum(self._desired_users_assigned_to_workers[worker_node.id].values())
                )
                if (
                    (
                        self._desired_users_assigned_to_workers[worker_node.id][user_class]
                        - self._dispatched_users[worker_node.id][user_class]
                    )
                    > 0
                )
            )
            == 0
        ):
            return True

        user_count_left_to_dispatch_excluding_excess_users = sum(
            user_classes_count_left_to_dispatch_excluding_excess_users.values()
        )
        return self._user_count_per_dispatch >= user_count_left_to_dispatch_excluding_excess_users

    def _users_to_dispatch_for_current_iteration(self) -> Dict[str, Dict[str, int]]:
        """
        Compute the users to dispatch for the current dispatch iteration.
        """
        user_count_in_current_dispatch = 0

        for i, user_class_to_add in enumerate(itertools.cycle(self._user_classes)):
            # For large number of user classes and large number of workers, this assertion might fail.
            # If this happens, you can remove it or increase the threshold. Right now, the assertion
            # is there as a safeguard for situations that can't be easily tested (i.e. large scale distributed tests).
            assert i < 100 * len(
                self._user_classes
            ), "Looks like dispatch is stuck in an infinite loop (iteration {}). Crash dump:\n{}\nAlso written to {}".format(
                i, *self._crash_dump()
            )

            if all(self._user_class_cannot_be_assigned_to_any_worker(user_class) for user_class in self._user_classes):
                # This means that we're at the last iteration of this dispatch cycle. If some user
                # classes are in excess, this last iteration will stop those excess users.
                self._dispatched_users.update(self._desired_users_assigned_to_workers)
                self._users_left_to_assigned.update(
                    {
                        worker_node_id: {user_class: 0 for user_class in user_classes_count.keys()}
                        for worker_node_id, user_classes_count in self._dispatched_users.items()
                    }
                )
                break

            if self._dispatched_user_class_count(user_class_to_add) >= self._user_classes_count[user_class_to_add]:
                continue

            if self._user_class_cannot_be_assigned_to_any_worker(user_class_to_add):
                continue

            if self._try_next_user_class_in_order_to_stay_balanced_during_ramp_up(user_class_to_add):
                continue

            for j, worker_node_id in enumerate(itertools.cycle(sorted(self._users_left_to_assigned.keys()))):
                assert j < int(
                    2 * self._number_of_workers
                ), "Looks like dispatch is stuck in an infinite loop (iteration {}). Crash dump:\n{}\nAlso written to {}".format(
                    j, *self._crash_dump()
                )
                if self._worker_is_full(worker_node_id):
                    continue
                if self._dispatched_user_count == self._desired_user_count or (
                    self._user_count_per_dispatch - user_count_in_current_dispatch >= self._user_count_left_to_assigned
                ):
                    # This means that we're at the last iteration of this dispatch cycle. If some user
                    # classes are in excess, this last iteration will stop those excess users.
                    self._dispatched_users.update(self._desired_users_assigned_to_workers)
                    self._users_left_to_assigned.update(
                        {
                            worker_node_id: {user_class: 0 for user_class in user_classes_count.keys()}
                            for worker_node_id, user_classes_count in self._dispatched_users.items()
                        }
                    )
                    break
                if self._users_left_to_assigned[worker_node_id][user_class_to_add] == 0:
                    continue
                if self._try_next_worker_in_order_to_stay_balanced_during_ramp_up(worker_node_id, user_class_to_add):
                    continue
                self._dispatched_users[worker_node_id][user_class_to_add] += 1
                self._users_left_to_assigned[worker_node_id][user_class_to_add] -= 1
                user_count_in_current_dispatch += 1
                break

            if user_count_in_current_dispatch == self._user_count_per_dispatch:
                break

            if self._dispatched_user_count == self._desired_user_count:
                break

        return {
            worker_node_id: dict(sorted(user_classes_count.items(), key=itemgetter(0)))
            for worker_node_id, user_classes_count in sorted(self._dispatched_users.items(), key=itemgetter(0))
        }

    def _crash_dump(self) -> Tuple[str, str]:
        """Parameters necessary to debug infinite loop issues.

        Users encountering an infinite loop issue should provide these informations.
        """
        crash_dump = json.dumps(
            {
                "spawn_rate": self._spawn_rate,
                "initial_dispatched_users": self._initial_dispatched_users,
                "desired_users_assigned_to_workers": self._desired_users_assigned_to_workers,
                "user_classes_count": self._user_classes_count,
                "initial_user_count": sum(map(sum, map(dict.values, self._initial_dispatched_users.values()))),
                "desired_user_count": sum(self._user_classes_count.values()),
                "number_of_workers": self._number_of_workers,
            },
            indent="  ",
        )
        fp = tempfile.NamedTemporaryFile(
            prefix="locust-dispatcher-crash-dump-", suffix=".json", mode="wt", delete=False
        )
        try:
            fp.write(crash_dump)
        finally:
            fp.close()
        return crash_dump, fp.name

    @property
    def _number_of_workers(self) -> int:
        return len(self._users_left_to_assigned)

    @property
    def _all_users_have_been_dispatched(self) -> bool:
        return self._user_count_left_to_assigned == 0

    @property
    def _user_count_left_to_assigned(self) -> int:
        return sum(map(sum, map(dict.values, self._users_left_to_assigned.values())))

    def _try_next_user_class_in_order_to_stay_balanced_during_ramp_up(self, user_class_to_add: str) -> bool:
        """
        Whether to skip to the next user class or not. This is done so that
        the distribution of user class stays approximately balanced from one dispatch
        iteration to another.
        """
        # For performance reasons, we use `functools.lru_cache()` on the `self._dispatched_user_classes_count`
        # method because its value does not change within the scope of the current method. However, the next time
        # `self._try_next_user_class_in_order_to_stay_balanced_during_ramp_up` is invoked, we need
        # `self._dispatched_user_classes_count` to be recomputed.
        self._dispatched_user_classes_count.cache_clear()

        if all(user_count > 0 for user_count in self._dispatched_user_classes_count().values()):
            # We're here because each user class have at least one user running. Thus,
            # we need to ensure that the distribution of users corresponds to the weights.
            if not self._adding_this_user_class_respects_distribution_better_than_adding_any_other_user_class(
                user_class_to_add
            ):
                return True
            else:
                return False

        else:
            # Because each user class doesn't have at least one running user, we use a simpler strategy
            # that makes sure each user class appears once.
            #
            # The following code checks if another user class that would better preserves the distribution exists.
            # If no such user class exists, this code will evaluate to `False` and the current user class
            # will be considered as the next user class to be added to the pool of running users.
            return any(
                True
                for next_user_class in filter(functools.partial(ne, user_class_to_add), self._user_classes)
                if sum(map(itemgetter(next_user_class), self._users_left_to_assigned.values())) > 0
                if self._dispatched_user_class_count(next_user_class) < self._user_classes_count[next_user_class]
                if not self._user_class_cannot_be_assigned_to_any_worker(next_user_class)
                if (
                    self._dispatched_user_classes_count()[user_class_to_add]
                    - self._dispatched_user_classes_count()[next_user_class]
                    >= 1
                )
            )

    def _adding_this_user_class_respects_distribution_better_than_adding_any_other_user_class(
        self, user_class_to_add: str
    ) -> bool:
        distance = self._distance_from_ideal_distribution()
        distance_after_adding_user_class = self._distance_from_ideal_distribution_after_adding_this_user_class(
            user_class_to_add
        )
        if distance_after_adding_user_class > distance and all(
            not self._adding_this_user_class_respects_distribution(user_class)
            for user_class in self._user_classes_count.keys()
            if user_class != user_class_to_add
            if not self._user_class_cannot_be_assigned_to_any_worker(user_class)
        ):
            # If we are here, it means that if one user of `user_class_to_add` is added
            # then the distribution will be the best we can get. In other words, adding
            # one user of any other user class won't yield a better distribution.
            return True
        return distance_after_adding_user_class <= distance

    def _adding_this_user_class_respects_distribution(self, user_class_to_add: str) -> bool:
        if (
            self._distance_from_ideal_distribution_after_adding_this_user_class(user_class_to_add)
            <= self._distance_from_ideal_distribution()
        ):
            return True
        return False

    def _distance_from_ideal_distribution(self) -> float:
        """How far are we from the ideal distribution given the current set of running users?"""
        weights = [
            self._dispatched_user_classes_count()[user_class] / sum(self._dispatched_user_classes_count().values())
            for user_class in self._user_classes
        ]

        return math.sqrt(sum(map(lambda x: (x[1] - x[0]) ** 2, zip(weights, self._desired_relative_weights()))))

    def _distance_from_ideal_distribution_after_adding_this_user_class(self, user_class_to_add: str) -> float:
        """
        How far are we from the ideal distribution if we were to add `user_class_to_add` to the pool of running users?
        """
        relative_weights_with_added_user_class = [
            (
                self._dispatched_user_classes_count()[user_class] + 1
                if user_class == user_class_to_add
                else self._dispatched_user_classes_count()[user_class]
            )
            / (sum(self._dispatched_user_classes_count().values()) + 1)
            for user_class in self._user_classes
        ]

        return math.sqrt(
            sum(
                map(
                    lambda x: (x[1] - x[0]) ** 2,
                    zip(relative_weights_with_added_user_class, self._desired_relative_weights()),
                )
            )
        )

    @functools.lru_cache()
    def _desired_relative_weights(self) -> List[float]:
        """The relative weight of each user class we desire"""
        return [self._user_classes_count[user_class] / self._desired_user_count for user_class in self._user_classes]

    @functools.lru_cache()
    def _dispatched_user_classes_count(self) -> Dict[str, int]:
        """The user count for each user class that are dispatched at this time"""
        return {
            user_class: self._dispatched_user_class_count(user_class) for user_class in self._user_classes_count.keys()
        }

    def _try_next_worker_in_order_to_stay_balanced_during_ramp_up(
        self, worker_node_id_to_add_user_on: str, user_class: str
    ) -> bool:
        """
        Whether to skip to the next worker or not. This is done so that
        each worker runs approximately the same amount of users during a ramp-up.
        """
        if self._dispatched_user_count == 0:
            return False

        workers_user_count = self._workers_user_count

        # Represents the ideal workers on which we'd want to add the user class
        # because these workers contain less users than all the other workers
        ideal_worker_node_ids = [
            ideal_worker_node_id
            for ideal_worker_node_id in workers_user_count.keys()
            if workers_user_count[ideal_worker_node_id] + 1 - min(workers_user_count.values()) < 2
        ]

        if worker_node_id_to_add_user_on in ideal_worker_node_ids:
            return False

        # Only keep the workers having less users than the target value as
        # we can't add users to those workers anyway.
        workers_user_count_without_excess_users = {
            worker_node_id: user_count
            for worker_node_id, user_count in workers_user_count.items()
            if user_count < sum(self._desired_users_assigned_to_workers[worker_node_id].values())
        }

        ideal_worker_on_which_to_add_user_exists = any(
            self._users_left_to_assigned[ideal_worker_node_id][user_class] > 0
            for ideal_worker_node_id in ideal_worker_node_ids
        )

        if worker_node_id_to_add_user_on not in workers_user_count_without_excess_users:
            return ideal_worker_on_which_to_add_user_exists

        if (
            workers_user_count_without_excess_users[worker_node_id_to_add_user_on]
            + 1
            - min(workers_user_count.values())
            >= 2
            and ideal_worker_on_which_to_add_user_exists
        ):
            # Adding the user to the current worker will result in this worker having more than 1
            # extra users compared to the other workers (condition on the left of the `and` above).
            # Moreover, we know there exists at least one other worker that would better host
            # the new user (condition on the right of the `and` above). Thus, we skip to the next worker node.
            return True

        return False

    @property
    def _workers_user_count(self) -> Dict[str, int]:
        """User count currently running on each of the workers"""
        return {
            worker_node_id: sum(dispatched_users_on_worker.values())
            for worker_node_id, dispatched_users_on_worker in self._dispatched_users.items()
        }

    def _dispatched_user_class_count(self, user_class: str) -> int:
        """Number of dispatched users at this time for the given user class"""
        return sum(map(itemgetter(user_class), self._dispatched_users.values()))

    @property
    def _dispatched_user_count(self) -> int:
        """Number of dispatched users at this time"""
        return sum(map(sum, map(dict.values, self._dispatched_users.values())))

    def _user_class_cannot_be_assigned_to_any_worker(self, user_class_to_add: str) -> bool:
        """No worker has enough place to accept this user class"""
        effective_user_count_of_that_user_class_that_can_be_added = sum(
            self._desired_users_assigned_to_workers[worker_node.id][user_class_to_add]
            - self._dispatched_users[worker_node.id][user_class_to_add]
            for worker_node in self._worker_nodes
        )
        if effective_user_count_of_that_user_class_that_can_be_added <= 0:
            return True

        if any(
            True
            for worker_node in self._worker_nodes
            if not self._worker_is_full(worker_node.id)
            if self._users_left_to_assigned[worker_node.id][user_class_to_add] > 0
        ):
            return False

        return True

    def _worker_is_full(self, worker_node_id: str) -> bool:
        """The worker cannot accept more users without exceeding the maximum user count it can run"""
        return self._workers_user_count[worker_node_id] >= self._workers_desired_user_count[worker_node_id]

class _WorkersUsersAssignor:
    """Helper to compute the users assigned to the workers"""

    def __init__(self, user_classes_count: Dict[str, int], worker_nodes: "List[WorkerNode]"):
        self._user_classes_count = user_classes_count
        self._user_classes = sorted(user_classes_count.keys())
        self._worker_nodes = sorted(worker_nodes, key=lambda w: w.id)

    @property
    def desired_users_assigned_to_workers(self) -> Dict[str, Dict[str, int]]:
        """The users assigned to the workers.

        The assignment is done in a way that each worker gets around the same number of users of each user class.
        If some user classes are more represented than others, then it is not possible to equally distribute
        the users from each user class to all workers. It is done in a best-effort.

        The assignment also ensures that each worker runs the same amount of users (irrespective of the user class
        of those users). If the total user count does not yield an integer when divided by the number of workers,
        then some workers will have one more user than the others.
        """
        assigned_users = {
            worker_node.id: {user_class: 0 for user_class in self._user_classes} for worker_node in self._worker_nodes
        }

        # We need to copy to prevent modifying `user_classes_count`.
        user_classes_count = self._user_classes_count.copy()

        user_count = sum(user_classes_count.values())

        # If `remainder > 0`, it means that some workers will have `users_per_worker + 1` users.
        users_per_worker, remainder = divmod(user_count, len(self._worker_nodes))

        for user_class in self._user_classes:
            if sum(user_classes_count.values()) == 0:
                # No more users of any user class to assign to workers, so we can exit this loop.
                break

            # Assign users of `user_class` to the workers in a round-robin fashion.
            for worker_node in itertools.cycle(self._worker_nodes):
                if user_classes_count[user_class] == 0:
                    break

                number_of_users_left_to_assign = user_count - self._number_of_assigned_users_across_workers(
                    assigned_users
                )

                if (
                    self._number_of_assigned_users_for_worker(assigned_users, worker_node) == users_per_worker
                    and number_of_users_left_to_assign > remainder
                ):
                    continue

                elif (
                    self._number_of_assigned_users_for_worker(assigned_users, worker_node) == users_per_worker + 1
                    and number_of_users_left_to_assign < remainder
                ):
                    continue

                assigned_users[worker_node.id][user_class] += 1
                user_classes_count[user_class] -= 1

        return assigned_users

    @staticmethod
    def _number_of_assigned_users_for_worker(
        assigned_users: Dict[str, Dict[str, int]], worker_node: "WorkerNode"
    ) -> int:
        """User count running on the given worker"""
        return sum(assigned_users[worker_node.id].values())

    @staticmethod
    def _number_of_assigned_users_across_workers(assigned_users: Dict[str, Dict[str, int]]) -> int:
        """Total user count running on the workers"""
        return sum(map(sum, map(methodcaller("values"), assigned_users.values())))


# coding:utf-8

from scrapy.selector import HtmlXPathSelector
from scrapy.contrib.spiders import CrawlSpider
from scrapy.http.request import Request
from youtube.items import youtubeItem

class youtubeK18Spider(CrawlSpider):
    name = 'youtube_k18'
    allowed_domains = ['www.youtube.com']
    start_urls = ['http://www.youtube.com/education?hl=en']

    def parse(self, response):
        filter_cate = ['Primary & Secondary Education']
        hxs = HtmlXPathSelector(response)
        sites = hxs.select('//div[@class=\'edu-mixed-collection\']')
        for site in sites:
            item = youtubeItem()
            category = site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/text()').extract()[0]
            category = category.replace(u'»',' ')
            category = category.strip();
            is_target = category in filter_cate
            if is_target == True:
                item['category1'] = category
                url = 'http://www.youtube.com/education' + site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/@href').extract()[0]
                request = Request(url, meta={'item': item}, callback=self.parse_category2)
                yield request

    def parse_category2(self, response):
        hxs = HtmlXPathSelector(response)
        sites = hxs.select('//div[@class=\'yt-uix-slider browse-collection yt-rounded\']')
        for site in sites:
            item = response.meta['item']
            has_data = len(site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/text()'))
            if has_data == 0:
                continue
            category = site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/text()').extract()[0]
            category = category.replace(u'»',' ')
            category = category.strip();
            url = site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/@href').extract()[0]
            if url != '':
                item['category2'] = category
                url = 'http://www.youtube.com/education' + url
                request = Request(url, meta={'item': item}, callback=self.parse_category3)
                yield request

    def parse_category3(self, response):
        hxs = HtmlXPathSelector(response)
        sites = hxs.select('//div[@class=\'yt-uix-slider browse-collection yt-rounded\']')
        for site in sites:
            item = response.meta['item']
            has_data = len(site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/text()'))
            if has_data == 0:
                continue
            category = site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/text()').extract()[0]
            category = category.replace(u'»',' ')
            category = category.strip();
            url = site.select('.//div[@class=\'yt-uix-slider-title\']/h2/a/@href').extract()[0]
            if url != '':
                item['category3'] = category
                url = 'http://www.youtube.com/education' + url
                request = Request(url, meta={'item': item}, callback=self.parse_category4)
                yield request

    def parse_category4(self, response):
        hxs = HtmlXPathSelector(response)
        site_first =  hxs.select('//button[@class=\'yt-uix-button-toggled yt-uix-button\']')
        sites = hxs.select('//button[@class=\' yt-uix-button\']')
        sites.append(site_first)
        filter_urls = []
        for site in sites:
            item = response.meta['item']
            has_data = len(site.select('./@href').extract())
            if has_data == 0:
                continue
            url = site.select('./@href').extract()[0]
            url = 'http://www.youtube.com/education' + url
            is_filter_target = url in filter_urls
            if is_filter_target == True:
                continue
            filter_urls.append(url)
            request = Request(url, meta={'item': item}, callback=self.parse_detail1)
            yield request

    def parse_detail1(self, response):
        hxs = HtmlXPathSelector(response)
        sites = hxs.select('/html/body/div/div[2]/div/div[2]/div[3]/div/div/ul/li')
        for site in sites:
            item = response.meta['item']
            url = site.select('div/div/div/h3/a/@href').extract()[0]
            url = 'http://www.youtube.com' + url
            request = Request(url, meta={'item': item}, callback=self.parse_detail2)
            yield request

    def parse_detail2(self, response):
        hxs = HtmlXPathSelector(response)
        sites = hxs.select('/html/body/div/div[2]/div/div/div[3]/div/ol/li')
        for site in sites:
            item = response.meta['item']
            item['img'] = site.select('div/a/span/span/img/@src').extract()
            item['detail_title1'] = site.select('h3/a/text()').extract()
            item['detail_link'] = 'http://www.youtube.com' + site.select('h3/a/@href').extract()[0]
            yield item

        sites = hxs.select('/html/body/div/div[2]/div/div/div[3]/div/div[2]/div[2]/div/ul/li')
        for site in sites:
            item = response.meta['item']
            item['img'] = site.select('a/span/span/span/img/@src').extract()
            item['detail_title1'] = site.select('a/span[2]/text()').extract()
            has_data = len(site.select('a/@href').extract())
            #TODO load more
            if has_data == 0:
                continue
            item['detail_link'] = 'http://www.youtube.com' + site.select('a/@href').extract()[0]
            yield item


import os
import sys
import numpy as np
import tensorflow as tf
from tensorflow.keras.losses import categorical_crossentropy as logloss
from tensorflow.keras.losses import kullback_leibler_divergence
from tensorflow.keras.metrics import categorical_accuracy
from sklearn.metrics import f1_score, precision_score, recall_score, confusion_matrix
import argparse
import h5py

from load_data import *
from model import *
from config_cnn import *

def acc(y_true, y_pred):
	y_true_hard = y_true[:, :2]
	y_pred_hard = y_pred[:, :2]
	return categorical_accuracy(y_true_hard, y_pred_hard)

def categorical_crossentropy(y_true, y_pred):
	y_true_hard = y_true[:, :2]
	y_pred_hard = y_pred[:, :2]
	return logloss(y_true_hard, y_pred_hard)

def kld_loss(y_true, y_pred):
	y_true_softs = y_true[:, 2:]
	y_pred_softs = y_pred[:, 2:]
	return kullback_leibler_divergence(y_true_softs, y_pred_softs)

def kd_loss(alpha, temperature):

	def custom_loss(y_true, y_pred):

		y_true, y_true_softs = y_true[:, :2], y_true[:, 2:]
		y_pred, y_pred_softs = y_pred[:, :2], y_pred[:, 2:]

		loss = (1-alpha)*logloss(y_true, y_pred) + alpha*(temperature**2)*kullback_leibler_divergence(y_true_softs, y_pred_softs)

		return loss

	return custom_loss

def sample_scores(loaded_model, model_type, song):

	x_test, y_test = load_xy_data(song, MEL_JAMENDO_DIR, JAMENDO_LABEL_DIR)

	y_pred = loaded_model.predict(x_test, verbose=1)

	if model_type=='kd': y_pred = y_pred[:, :2]

	y_pred = np.argmax(y_pred, axis=1)

	y_pred = y_pred.reshape(-1).astype(int)
	y_test = y_test.reshape(-1).astype(int)

	accuracy_single = (len(y_test) - np.sum(np.abs(y_pred - y_test)))/ len(y_test)
	f1 = f1_score(y_test, y_pred, average='binary')
	pr = precision_score(y_test, y_pred, average='binary')
	re = recall_score(y_test, y_pred, average='binary')

	print('\nSample Scores...\n')
	print('Accuracy: ' + str(accuracy_single))
	print('Precision: ', pr)
	print('Recall: ', re)
	print('F1-score: ', f1)

	del x_test

	return y_pred, y_test

def test(model_name, model_type, wts_dir, args):

	if model_type=='kd':
		model = RNN_small(timesteps=RNN_INPUT_SIZE)
		model.pop()
		model_logits = model.layers[-1].output
		model_logits_T = Lambda(lambda x: x/args.temperature)(model_logits)
		probs_T = Softmax(axis=1)(model_logits_T)
		probs_1 = Softmax(axis=1)(model_logits)
		output = Concatenate()([probs_1, probs_T])
		model = Model(inputs=model.input, outputs=output)

		model.load_weights(wts_dir+model_name)
	else:
		model = tf.keras.models.load_model(wts_dir+model_name)

	print(model.summary())

	list_of_songs = os.listdir(MEL_JAMENDO_DIR + 'test')

	y_preds = []
	y_tests = []

	for song in list_of_songs:

		y_pred, y_test = sample_scores(model, model_type, song)

		for i in range(len(y_pred)):
			y_preds.append(y_pred[i])
			y_tests.append(y_test[i])

	# convert list to np array
	y_preds = np.array(y_preds)
	y_tests = np.array(y_tests)

	# calculate scores
	acc = (len(y_tests) - np.sum(np.abs(y_preds - y_tests))) / float(len(y_tests))

	f1 = f1_score(y_tests, y_preds, average='binary')
	pr = precision_score(y_tests, y_preds, average='binary')
	re = recall_score(y_tests, y_preds, average='binary')

	tn, fp, fn, tp = confusion_matrix(y_tests, y_preds).ravel()
	fp_rate = fp / (fp + tn)
	fn_rate = fn / (fn + tp)

	print("TEST SCORES OVERALL\n")
	print('Acc %.4f' % acc)
	print('Precision %.4f' % pr)
	print('Recall %.4f' % re)
	print('F1-score %.4f' % f1)
	print('fp rate', fp_rate, 'fn_rate', fn_rate)

	scores = [acc, pr, re, f1, fp_rate, fn_rate]

	return scores


from flask_sqlalchemy import SQLAlchemy
from flask import *
from flask_cors import CORS

SCC = Flask(__name__)
SCC.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///SCCData.db'
SCC.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
db = SQLAlchemy(SCC)
CORS(SCC)

class Users(db.Model):
    __tablename__ = 'Users'

    Name = db.Column(db.VARCHAR(40), nullable=False)
    SID = db.Column(db.String(8), nullable=False, primary_key=True)
    Mobile = db.Column(db.String(10), nullable=False)
    MobileVerified = db.Column(db.BOOLEAN, nullable=False)

    def __repr__(self):
        return '(Name: ' + self.Name + ' SID: ' + str(self.SID) + ' Mobile: ' + str(self.Mobile) + ' Mobile Verified: ' + str(self.MobileVerified) + ' )'

@SCC.route('/')
def hello_world():
    return render_template("Landing_page.html")

@SCC.route('/otp', methods=['POST', 'GET'])
def send_otp():
    result = False
    if request.method == 'POST':
        data = request.form
        print(data)
        if data:
            Name = data['Name']
            SID = str(data['SID'])
            print(SID)
            Mobile = str(data['Mobile'])
            MobileVerified = False
            user = Users(Name=Name, SID=SID, Mobile=Mobile, MobileVerified=MobileVerified)
            curr_session = db.session
            try:
                curr_session.add(user)
                curr_session.commit()
                result = True
            except:
                curr_session.rollback()
                curr_session.flush()
                result = False
        elif data is None:
            return render_template('Landing_page.html', message="Data invalid")
    else:
        return render_template('Landing_page.html', message="Invalid Method")

    if result:
        return render_template('Questions.html')
    else:
        return render_template('Landing_page.html', message="Some unknown error occured.")

if __name__ == '__main__':
    SCC.run(debug=True)


from flask import Flask, render_template, redirect, url_for, request
from FUNKss import SaveMessage

app = Flask(__name__)
logged_in = True

@app.route('/')
def main():
	global logged_in
	print(logged_in)
	return render_template('index.html')

@app.route('/next')
def other():
	global logged_in
	if logged_in == False:
		print(logged_in)
		return redirect(url_for('login'))
	else:
		print(logged_in)
		return render_template('other.html')

@app.route('/login', methods=['GET','POST'])
def login():
	global logged_in
	error_msg = None
	if request.method == 'POST':
		if request.form['username'] != 'admin' or request.form['password'] != 'password':
			error_msg = "Invalid password or username"
			logged_in = False
			print(logged_in)
		else:
			logged_in = True
			print(logged_in)
			return redirect(url_for('main'))
	return render_template('login.html', error=error_msg)

@app.route('/logout')
def logout():
	global logged_in
	logged_in = False
	print(logged_in)
	return "You have succesfully logged out"

@app.route('/entries', methods=['GET','POST'])
def write():
	global logged_in
	if logged_in == False:
		return redirect(url_for('login'))
	else:
		if request.method == 'POST':
			a = request.form['entry']
			b = request.form['named']
			SaveMessage(str(b),str(a))
		return render_template('in.html')
	return redirect(url_for('main'))

# @app.route('/blog')
# def blog():
#  	return render_template('blog.html',title=title,data=data)

if __name__ == "__main__":
	app.run(debug=True)


"""Overlay for third-party chex.dataclass decorator.

See https://github.com/deepmind/chex#dataclass-dataclasspy. Typing-wise, the
differences between @dataclasses.dataclass and @chex.dataclass are:
* The latter has a mappable_dataclass parameter, defaulting to True, which makes
  the dataclass inherit from Mapping.
* Chex dataclasses have replace, from_tuple, and to_tuple methods.
"""

from pytype import overlay
from pytype import overlay_utils
from pytype.abstract import abstract
from pytype.overlays import classgen
from pytype.overlays import dataclass_overlay
from pytype.pytd import pytd

class ChexOverlay(overlay.Overlay):

  def __init__(self, ctx):
    member_map = {
        "dataclass": Dataclass.make,
    }
    ast = ctx.loader.import_name("chex")
    super().__init__(ctx, "chex", member_map, ast)

class Dataclass(dataclass_overlay.Dataclass):
  """Implements the @dataclass decorator."""

  _DEFAULT_ARGS = {**dataclass_overlay.Dataclass._DEFAULT_ARGS,
                   "mappable_dataclass": True}

  @classmethod
  def make(cls, ctx):
    return super().make(ctx, "chex")

  def _add_replace_method(self, node, cls):
    cls.members["replace"] = classgen.make_replace_method(
        self.ctx, node, cls, kwargs_name="changes")

  def _add_from_tuple_method(self, node, cls):
    # from_tuple is discouraged anyway, so we provide only bare-bones types.
    cls.members["from_tuple"] = overlay_utils.make_method(
        ctx=self.ctx,
        node=node,
        name="from_tuple",
        params=[overlay_utils.Param("args")],
        return_type=cls,
        kind=pytd.MethodTypes.STATICMETHOD,
    )

  def _add_to_tuple_method(self, node, cls):
    # to_tuple is discouraged anyway, so we provide only bare-bones types.
    cls.members["to_tuple"] = overlay_utils.make_method(
        ctx=self.ctx,
        node=node,
        name="to_tuple",
        return_type=self.ctx.convert.tuple_type,
    )

  def _add_mapping_base(self, node, cls):
    mapping = self.ctx.convert.name_to_value("typing.Mapping")
    # The class's MRO is constructed from its bases at the moment the class is
    # created, so both need to be updated.
    bases = cls.bases()
    # If any class in Mapping's MRO already exists in the list of bases, Mapping
    # needs to be inserted before it, otherwise we put it at the end.
    mapping_mro = {x.full_name for x in mapping.mro}
    cls_bases = [x.data[0].full_name for x in bases]
    cls_mro = [x.full_name for x in cls.mro]
    bpos = [i for i, x in enumerate(cls_bases) if x in mapping_mro]
    mpos = [i for i, x in enumerate(cls_mro) if x in mapping_mro]
    if bpos:
      bpos, mpos = bpos[0], mpos[0]
      bases.insert(bpos, mapping.to_variable(node))
      cls.mro = cls.mro[:mpos] + (mapping,) + cls.mro[mpos:]
    else:
      bases.append(mapping.to_variable(node))
      cls.mro = cls.mro + (mapping,)

  def decorate(self, node, cls):
    super().decorate(node, cls)
    if not isinstance(cls, abstract.InterpreterClass):
      return
    self._add_replace_method(node, cls)
    self._add_from_tuple_method(node, cls)
    self._add_to_tuple_method(node, cls)
    if not self.args[cls]["mappable_dataclass"]:
      return
    self._add_mapping_base(node, cls)


class Solution(object):
    def fourSumCount(self, A, B, C, D):
        """
        :type A: List[int]
        :type B: List[int]
        :type C: List[int]
        :type D: List[int]
        :rtype: int
        """
        count = 0
        sum_1, sum_2 = dict(), dict()
        length = len(A)
        for i in range(length):
            for j in range(length):
                cur_sum = A[i] + B[j]
                if sum_1.has_key(cur_sum):
                    sum_1[cur_sum] += 1
                else:
                    sum_1[cur_sum] = 1
        for i in range(length):
            for j in range(length):
                cur_sum = C[i] + D[j]
                if sum_2.has_key(cur_sum):
                    sum_2[cur_sum] += 1
                else:
                    sum_2[cur_sum] = 1
        # loop all the hash maps
        for key in sum_1:
            to_find = - key
            if sum_2.has_key(to_find):
                count += sum_1[key] * sum_2[to_find]
        return count


import pandas as pd
import matplotlib.pyplot as plt

def plot_term_matrix(matrix, precision):
    plt.subplots(figsize=(20, 25))
    plt.spy(matrix, precision=precision, markersize=1)

def convert_file_to_dataframe(file):
    with open(file, 'r', encoding='utf8') as reader:
        lines = reader.readlines()
        sentence = []
        label = []
        for l in lines:
            split = l.split('\t')
            sentence.append(split[0])
            label.append(split[1].rstrip())

        df = pd.DataFrame({'sentence': sentence, 'label': label})

    return df


from django.conf.urls import url, include
from rest_framework import routers
from turbine_data import views

router = routers.DefaultRouter()
router.register(r'turbine-data', views.TurbineDataViewSet)

urlpatterns = [
    url(r'^', include(router.urls), name='turbine-data'),
    url(
        r'^turbine-data/last',
        views.getLastTurbineData,
        name='turbine-data-last'
    ),
    url(
        r'^start_year:(?P<start_year>\d{4})&&finish_year:(?P<finish_year>\d{4})/$',
        views.getTurbineDataByYear,
        name='turbine-data-year'
    ),
    url(
        r'^year:(?P<year>\d{4})&&start_month:(?P<start_month>\d{2})&&finish_month:(?P<finish_month>\d{2})/$', views.getTurbineDataByMonth,
        name='turbine-data-month'
    ),
    url(
        r'^start:(?P<start_year>\d{4})-(?P<start_month>\d{2})-(?P<start_day>\d{2})&&finish:(?P<finish_year>\d{4})-(?P<finish_month>\d{2})-(?P<finish_day>\d{2})/$',
        views.getTurbineDataByCompleteDate,
        name='turbine-data-complete'
    ),
    url(
        r'^start:(?P<start_year>\d{4})-(?P<start_month>\d{2})-(?P<start_day>\d{2})/$',
        views.getTurbineDataByDay,
        name='turbine-data-day'
    ),
]


import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt
from tqdm import tqdm

def probit_vector(u_vec):
    return np.array([stats.norm.cdf(u,0,1) for u in u_vec])

def make_data_and_gibbs(w, mu, sigma, N=100, B_burn = 1000, B_save = 500):
    ######### Synthesize data
    x = np.random.randn(N,D)
    y = np.squeeze(np.array([np.random.binomial(1, p) for p in probit_vector(np.matmul(w.reshape((1,2)),x.T))]))
    ######### Done synthesizing data

    ######### Gibbs Sampling
    w_gibbs_list = []

    # 0. Initialize w
    w_gibbs = np.random.randn(D)

    for b in tqdm(range(B_burn+B_save)):
        # 1. Sample z_sampled[n] ~ p(z_n|x[n],y[n],w_gibbs) for n=1,...,N
        z_sampled = np.empty(N)
        for n in range(N):
            condition_satisfied = False
            while not condition_satisfied:
                z_sampled[n] = np.random.normal(np.dot(x[n],w_gibbs), 1)
                if y[n] == 1:
                    condition_satisfied = z_sampled[n] >= 0
                else:
                    condition_satisfied = z_sampled[n] < 0
        # z_sampled

        # 2. Sample w_gibbs ~ p(w|{x[n],y[n],z_sampled[n]})
        J = np.linalg.inv(sigma) + np.matmul(x.T,x)
        sigma_cond_gibbs = np.linalg.inv(J)
        h = np.matmul(np.linalg.inv(sigma), mu)
        for n in range(N):
            h += z_sampled[n]*x[n]
        w_gibbs = np.random.multivariate_normal(np.matmul(sigma_cond_gibbs,h), sigma_cond_gibbs)

        # 3. Store w_gibbs if bn >= B_burn
        if b >= B_burn:
            w_gibbs_list.append(w_gibbs)

    return w_gibbs_list

def try_Ns_and_plot(w, mu, sigma, Ns=[100, 500], B_burn = 1000, B_save = 500):
    f,ax = plt.subplots(1,len(Ns), figsize=(9,3))

    for i in range(len(Ns)):
        N = Ns[i]
        w_gibbs_list = make_data_and_gibbs(w, mu, sigma, N, B_burn, B_save)
        w_gibbs_xs = [w_g[0] for w_g in w_gibbs_list]
        w_gibbs_ys = [w_g[1] for w_g in w_gibbs_list]
        ax[i].scatter(w_gibbs_xs, w_gibbs_ys, c='blue')
        ax[i].scatter(*w, c='red')
        ax[i].set_aspect('equal')
        ax[i].set_xlim([w[0]-.5, w[0]+.5])
        ax[i].set_ylim([w[1]-.5, w[1]+.5])
        ax[i].set_title("N = " + str(N))
    f.show()
    return f

######### Synthesize params
D = 2
mu = np.zeros(D)
temp = np.random.randn(D,D)
sigma = temp.T * temp
w = np.random.multivariate_normal(mu, sigma)

######## Run exps and plot
f = try_Ns_and_plot(w, mu, sigma, [100,300,500])


import re
import os, sys

reps_regex = re.compile("\d*")

def find_workload_dirs(inputdir, regex_string):
    # Gather VCs tested from directory names
    workload_regex = re.compile(regex_string)
    all_directories = filter(lambda x: os.path.isdir(os.path.join(inputdir, x)), os.listdir(inputdir))
    matching_directories = [os.path.join(inputdir, dirname) for dirname in all_directories if workload_regex.match(dirname)]
    return matching_directories

def find_num_repetitions(single_configuration_dir):
    return len([1 for rep_dirname in os.listdir(single_configuration_dir) if reps_regex.match(rep_dirname)])

def get_and_validate_num_repetition(directories_to_check):
    num_repetitions = -1
    for directory in directories_to_check:
        rep_list_len = find_num_repetitions(directory)
        if (num_repetitions == -1):
            num_repetitions = rep_list_len
        else:
            if (rep_list_len != num_repetitions):
                print("Directory {} had {} reps, but expected {}".format(directory, rep_list_len, num_repetitions))
                exit(2)
    if (num_repetitions == -1):
        print("No directies to check have been provided. Cannot validate num of repetitions")
        exit(2)

    return num_repetitions

def find_num_vc(workload_dir_path):
    basedir = os.path.basename(workload_dir_path)
    start_idx = basedir.find("_") + 1 # num of vc occures right after the underscore
    end_idx = basedir.find("vc", start_idx)
    num_vc = int(basedir[start_idx:end_idx])
    return num_vc

def find_num_multigets(workload_dir_path):
    basedir = os.path.basename(workload_dir_path)
    start_idx = basedir.find("_") + 1 # num of multigets occures right after the underscore
    end_idx = basedir.find("multiget", start_idx)
    num_vc = int(basedir[start_idx:end_idx])
    return num_vc

def find_num_workers(workload_dir_path):
    basedir = os.path.basename(workload_dir_path)
    start_idx = basedir.find("vc") + 2 # num of workers occures right after vc
    end_idx = basedir.find("workers", start_idx)
    num_workers = int(basedir[start_idx:end_idx])
    return num_workers

def find_worker_configurations(workload_directories):
    all_worker_configurations = []
    for directory in workload_directories:
        num_workers = find_num_workers(directory)
        if num_workers not in all_worker_configurations:
            all_worker_configurations.append(find_num_workers(directory))

    return all_worker_configurations

def find_matching_worker_config_dirs(all_directories, worker_config):
    return [rep_dirname for rep_dirname in all_directories if str(worker_config) + "workers" in rep_dirname]

def get_only_subdir(dirpath):
    filter_dirs = filter(lambda x: os.path.isdir(os.path.join(dirpath, x)), os.listdir(dirpath))
    subdirs = [os.path.join(dirpath, dirname) for dirname in filter_dirs]
    if (len(subdirs) != 1):
        print ("Found more or less than one subdirectories in {}".format(dirpath))
        sys.exit(2)
    return os.path.join(dirpath, subdirs[0])


import random
import sys
import numpy as np
import tvm
from tvm import comm_reducer
from tvm.arith import Domain
from tvm.testing import estimate_performance
from tvm.ir_pass import Simplify, Equal, LiftNonzeronessCondition, IsSumCombiner, \
    CanFactorZeroFromCombiner, InlineTailCall, InlineTensors, SolveSystemOfInequalities, \
    SimplifyDomain, SimplifyReductionDomain, ExtractAsTensorMaybe, ExtractReductions, \
    ExtractNonTopReductions, OptimizeAndLiftNonzeronessConditions, SolveSystemOfEquations

def get_shape(tensor):
    return [s.value for s in tensor.shape]

def check_eq(t1, t2, args):
    s1 = tvm.create_schedule(t1.op)
    m1 = tvm.build(s1, [t1] + args)

    s2 = tvm.create_schedule(t2.op)
    m2 = tvm.build(s2, [t2] + args)

    for _ in range(5):
        arg_vals = [tvm.ndarray.array(np.random.uniform(-10, 10, size=get_shape(a))
                                      .astype(a.dtype))
                    for a in [t1] + args]
        m1(*arg_vals)
        res1 = arg_vals[0].asnumpy()
        m2(*arg_vals)
        res2 = arg_vals[0].asnumpy()

        np.testing.assert_allclose(res1, res2, atol=1e-3, rtol=1e-2)

def check_symeq(expr1, expr2):
    expr1 = tvm.ir_pass.Simplify(tvm.ir_pass.CanonicalSimplify(expr1))
    expr2 = tvm.ir_pass.Simplify(tvm.ir_pass.CanonicalSimplify(expr2))

    if tvm.ir_pass.Equal(expr1, expr2):
        return

    diff = tvm.ir_pass.Simplify(tvm.ir_pass.CanonicalSimplify(expr1 - expr2))
    if not Equal(diff, tvm.const(0, expr1.dtype)):
        raise AssertionError("Expressions {} and {} are not equal, their diff is {}"
                             .format(expr1, expr2, diff))

def compute(shape, fcompute):
    """Like tvm.compute but automatically extracts reductions."""
    return tvm.compute(shape,
                       lambda *vs: ExtractNonTopReductions(
                           fcompute(*vs), vs, {v: tvm.Range(0, s) for v, s in zip(vs, shape)}))

def check_tensor_symeq(A, B):
    if not isinstance(B, tvm.tensor.Tensor):
        B = compute(A.shape, B)
    vmap = {a.var: b.var for a, b in zip(A.op.axis, B.op.axis)}
    expr_a = tvm.ir_pass.Substitute(A.op.body[A.value_index], vmap)
    expr_b = B.op.body[B.value_index]
    expr_a = tvm.ir_pass.CanonicalSimplify(InlineTensors(expr_a, [], True))
    expr_b = tvm.ir_pass.CanonicalSimplify(InlineTensors(expr_b, [], True))
    if not Equal(expr_a, expr_b):
        print(expr_a)
        print(expr_b)
        raise AssertionError("The expressions are not equal")

def run_expr(expr, vranges):
    def _compute_body(*us):
        vmap = {v: u + r.min for (v, r), u in zip(vranges.items(), us)}
        return tvm.ir_pass.Substitute(expr, vmap)

    A = compute([r.extent.value for v, r in vranges.items()], _compute_body)
    args = [tvm.ndarray.empty(A.shape, A.dtype)]
    sch = tvm.create_schedule(A.op)
    mod = tvm.build(sch, [A])
    mod(*args)
    lowered = tvm.lower(sch, [A], simple_mode=True)
    return args[0].asnumpy()

def check_bruteforce(bool_expr, vranges, cond=None):
    if cond is not None:
        bool_expr = tvm.any(tvm.expr.Not(cond), bool_expr)

    res = run_expr(bool_expr, vranges)
    if not np.all(res):
        indices = list(np.argwhere(res == 0)[0])
        counterex = [(str(v), i + r.min) for (v, r), i in zip(vranges.items(), indices)]
        counterex = sorted(counterex, key=lambda x: x[0])
        counterex = ", ".join([v + " = " + str(i) for v, i in counterex])
        raise AssertionError("Expression {}\nis not true on {}\n"
                             "Counterexample: {}"
                             .format(tvm.ir_pass.CanonicalSimplify(bool_expr), vranges, counterex))

def all(*conds):
    return tvm.all(tvm.const(1, 'bool'), *conds)

def reduce_over_domain(domain, combiner, expr):
    axes = [tvm.reduce_axis(domain.ranges[v], v.name) for v in domain.variables]
    vars_to_axes = {v: a.var for v, a in zip(domain.variables, axes)}
    cond = all(*[tvm.ir_pass.Substitute(c, vars_to_axes) for c in domain.conditions])
    expr = tvm.ir_pass.Substitute(expr, vars_to_axes)
    reduction = combiner(expr, axes, cond)
    return reduction

def check_domain_transformation(domain_transform, vranges={}):
    def _check_forward(domain1, domain2, varmap, backvarmap, vranges=vranges):
        all_vranges = vranges.copy()
        all_vranges.update({v: r for v, r in domain1.ranges.items()})

        # Check that the transformation is injective
        cond_on_vars = tvm.const(1, 'bool')
        for v in domain1.variables:
            v_back = tvm.ir_pass.Simplify(tvm.ir_pass.Substitute(varmap[v], backvarmap))
            cond_on_vars = all(cond_on_vars, v == v_back)
        # Also we have to check that the new conds are true when old conds are true
        cond_subst = tvm.ir_pass.Substitute(all(*domain2.conditions), backvarmap)
        # We have to include conditions from vranges too
        for v in domain2.variables:
            if v in domain2.ranges:
                r = domain2.ranges[v]
                range_cond = all(v >= r.min, v < r.min + r.extent)
                range_cond = tvm.ir_pass.Substitute(range_cond, backvarmap)
                cond_subst = all(cond_subst, range_cond)
        cond_subst = tvm.ir_pass.Simplify(cond_subst)
        check_bruteforce(all(cond_subst, cond_on_vars), all_vranges, cond=all(*domain1.conditions))

        # Additionally check that some kind of a checksum is the same
        if domain1.variables:
            expr = sum([v*(i + 1) for i, v in enumerate(domain1.variables)])
        else:
            expr = tvm.const(1, 'int32')
        expr1 = reduce_over_domain(domain1, xor_combiner, expr)
        expr2 = reduce_over_domain(domain2, xor_combiner, tvm.ir_pass.Substitute(expr, varmap))
        #  expr1 = reduce_over_domain(domain1, sum_combiner, tvm.const(1, 'int64'))
        #  expr2 = reduce_over_domain(domain2, sum_combiner, tvm.const(1, 'int64'))
        #  print(run_expr(expr1, vranges))
        #  print(run_expr(expr2, vranges))
        check_bruteforce(expr1 == expr2, vranges)

    _check_forward(domain_transform.old_domain, domain_transform.new_domain,
                   domain_transform.old_to_new, domain_transform.new_to_old)
    _check_forward(domain_transform.new_domain, domain_transform.old_domain,
                   domain_transform.new_to_old, domain_transform.old_to_new)

prod_combiner = comm_reducer(lambda x, y: x*y, lambda t0: tvm.const(1, t0))
sum_combiner = comm_reducer(lambda x, y: x + y, lambda t0: tvm.const(0, t0))
sum2_combiner = comm_reducer(lambda x, y: y + x, lambda t0: tvm.const(0, t0))
sum_derivative_combiner = comm_reducer(lambda x, y: (x[0] + y[0], y[1] + x[1]),
                                       lambda t0, t1: (tvm.const(0, t0), tvm.const(0, t1)))
prod_derivative_combiner = comm_reducer(lambda x, y: (x[0]*y[0], x[0]*y[1] + x[1]*y[0]),
                                        lambda t0, t1: (tvm.const(1, t0), tvm.const(0, t1)))
sum_both_combiner = comm_reducer(lambda x, y: (x[0] + y[0], x[0] + y[0] + x[1] + y[1]),
                                        lambda t0, t1: (tvm.const(0, t0), tvm.const(0, t1)))
xor_combiner = comm_reducer(lambda x, y: x ^ y, lambda t0: tvm.const(0, t0))

m_param = tvm.var('m_param')
sum_or_prod_combiner = comm_reducer(lambda x, y: tvm.expr.Select(m_param < 0, x + y, x*y),
                                    lambda t0: tvm.expr.Select(m_param < 0,
                                                               tvm.const(0, t0), tvm.const(1, t0)))
shifted_sum_combiner = comm_reducer(lambda x, y: x + y - m_param,
                                    lambda t0: m_param)

def test_is_sum_combiner():
    k = tvm.reduce_axis((0, 10), name="k")
    i = tvm.const(0, "int32")
    f = tvm.const(0.0, "float32")
    assert IsSumCombiner(sum_combiner(i, k).combiner)
    assert IsSumCombiner(sum_combiner(f, k).combiner)
    assert IsSumCombiner(sum2_combiner(i, k).combiner)
    assert IsSumCombiner(sum2_combiner(f, k).combiner)
    assert not IsSumCombiner(sum_derivative_combiner((f, f), k)[0].combiner)
    assert not IsSumCombiner(prod_combiner(f, k).combiner)
    assert not IsSumCombiner(prod_derivative_combiner((f, f), k)[1].combiner)
    assert not IsSumCombiner(sum_or_prod_combiner(f, k).combiner)
    assert not IsSumCombiner(sum_or_prod_combiner(f, k).combiner, {m_param: tvm.Range(-5, 1)})
    assert IsSumCombiner(sum_or_prod_combiner(f, k).combiner, {m_param: tvm.Range(-5, -1)})
    assert not IsSumCombiner(shifted_sum_combiner(i, k).combiner)
    assert IsSumCombiner(shifted_sum_combiner(i, k).combiner, {m_param: tvm.Range(0, 1)})

def test_can_factor_zero_from_combiner():
    k = tvm.reduce_axis((0, 10), name="k")
    i = tvm.const(0, "int32")
    f = tvm.const(0.0, "float32")
    assert CanFactorZeroFromCombiner(sum_combiner(i, k).combiner, 0)
    assert CanFactorZeroFromCombiner(sum2_combiner(f, k).combiner, 0)
    assert CanFactorZeroFromCombiner(sum_derivative_combiner((f, f), k)[0].combiner, 0)
    assert CanFactorZeroFromCombiner(sum_derivative_combiner((f, f), k)[0].combiner, 1)
    assert not CanFactorZeroFromCombiner(prod_derivative_combiner((f, f), k)[0].combiner, 0)
    assert CanFactorZeroFromCombiner(prod_derivative_combiner((f, f), k)[0].combiner, 1)
    assert CanFactorZeroFromCombiner(sum_both_combiner((f, f), k)[0].combiner, 0)
    assert not CanFactorZeroFromCombiner(sum_both_combiner((f, f), k)[0].combiner, 1)
    assert not CanFactorZeroFromCombiner(sum_or_prod_combiner(f, k).combiner, 0,
                                         {m_param: tvm.Range(-5, 1)})
    assert CanFactorZeroFromCombiner(sum_or_prod_combiner(f, k).combiner, 0,
                                     {m_param: tvm.Range(-5, -1)})
    assert not CanFactorZeroFromCombiner(shifted_sum_combiner(i, k).combiner, 0)
    assert CanFactorZeroFromCombiner(shifted_sum_combiner(i, k).combiner, 0,
                                     {m_param: tvm.Range(0, 1)})

def test_lift_nonzeroness_condition():
    k = tvm.reduce_axis((0, 5), name="k")
    l = tvm.reduce_axis((0, 5), name="l")
    n = tvm.reduce_axis((0, 5), name="n")
    A = tvm.placeholder((10,), name='A')

    def _check(shape, fun, A=A):
        T1 = tvm.compute(shape, fun)
        T2 = tvm.compute(shape, lambda *args: LiftNonzeronessCondition(fun(*args)))
        check_eq(T1, T2, [A])
        assert isinstance(T2.op.body[0], tvm.expr.Select)

    _check((10,), lambda i: A[i])
    _check((10,), lambda i: A[i] + (i % 2 == 0))
    _check((10,), lambda i: A[i]*(i % 2 == 0) + (i % 2 == 0))
    _check((10,), lambda i: tvm.expr.Select((i % 2 == 0), A[i], 0.0))
    _check((10,), lambda i: tvm.expr.Select((i % 2 == 0), A[i], 0.0) + (i % 2 == 0))
    _check((10,), lambda i: tvm.expr.Select((i % 2 == 0), 0.0, A[i]) + (i % 2 == 0))
    def e1(i): return tvm.expr.Select((i % 2 == 1), 0.0, A[i])
    def e2(i): return tvm.expr.Select((i % 2 == 0), A[(i + 1) % 10], 0.0)
    def e3(i): return tvm.expr.Select((i % 2 == 1), A[i], 0.0)
    _check((10,), lambda i: e1(i) + e2(i) + e3(i) + e1(i)*e2(i))
    _check((10,), lambda i: e1(i)*e3(i))
    _check((10,), lambda i: e1(i)*e2(i))
    _check((10,10), lambda i, j: A[i]*(i == j) + A[j]*(i == 2*j) + A[j]*(j == i))
    _check((10,10), lambda i, j: tvm.min(A[i]*(i == j), A[j]*(i == 2*j)))
    _check((10,10), lambda i, j: tvm.max(A[i]*(i == j), A[j]*(i == 2*j)))
    _check((10,10), lambda i, j: A[i]*(i == j) - A[j]*(i == 2*j))
    _check((10,10), lambda i, j: A[i]*(i == j) / (1 + tvm.abs(A[j]*(i == 2*j))))
    _check((10,10), lambda i, j: i*(i < j) + j*(i > j))
    _check((10,10), lambda i, j: i*(i < j) % (1 + j*(i > j)))

    def _check_symeq(expr1, expr2):
        expr1 = LiftNonzeronessCondition(expr1)
        expr2 = LiftNonzeronessCondition(expr2)
        print(expr1)
        print(expr2)
        print()
        check_symeq(expr1, expr2)

    _check_symeq(tvm.expr.Select(tvm.expr.EQ(k, l), 0.0, tvm.expr.Cast('float32', (k < n))),
                 tvm.expr.Select(tvm.expr.And((k < n), tvm.expr.NE(k, l)), 1.0, 0.0))
    _check_symeq(tvm.min(tvm.expr.Cast('int32', k < n)*l, tvm.expr.Select(k >= n, 0, 1)),
                 tvm.expr.Select(k < n, tvm.min(l, 1), 0))

    expr1 = tvm.if_then_else(k < n,
                             tvm.expr.Select(tvm.expr.EQ(k, l), A[k],  0.0),
                             tvm.expr.Select(l < n, A[l], 0.0))
    expr2 = tvm.expr.Select(tvm.any(all(k < n, tvm.expr.EQ(k, l)),
                                    all(k >= n, l < n)),
                            tvm.if_then_else(k < n, A[k], A[l]),
                            0.0)
    check_symeq(LiftNonzeronessCondition(expr1), expr2)

def test_inline_tail_call():
    A = tvm.compute((10, 10), lambda i, j: i + j*j)
    B = tvm.compute((5, 6), lambda k, l: A[k + l, k + 1])
    C = InlineTailCall(B)
    resbody = lambda k, l: k + l + (k + 1)*(k + 1)
    check_symeq(C.op.body[0], resbody(*[iv.var for iv in C.op.axis]))

def test_inline_tensors():
    A = tvm.compute((10, 10), lambda i, j: i + j)
    B = tvm.compute((10, 10), lambda i, j: i * j)
    C = tvm.compute((10, 10), lambda i, j: A[i, j] + B[i, j])
    k = tvm.reduce_axis((0, 5), name="k")
    D = tvm.compute((10, 10), lambda i, j: tvm.sum(A[i, k], k))
    E = tvm.compute((10, 10), lambda i, j: A[2, j] + C[i, 2] + D[i, j])
    F = tvm.compute((10, 10), lambda i, j: tvm.exp(A[i, j]) + B[i, A[i, j]])

    R = InlineTensors(E)
    resbody = lambda i, j: 2 + j + i + 2 + i*2 + D[i, j]
    check_symeq(R.op.body[0], resbody(*[iv.var for iv in R.op.axis]))

    R = InlineTensors(E, [A])
    resbody = lambda i, j: 2 + j + C[i, 2] + D[i, j]
    check_symeq(R.op.body[0], resbody(*[iv.var for iv in R.op.axis]))

    R = InlineTensors(E, [A, C])
    resbody = lambda i, j: 2 + j + ((i + 2) + B[i, 2]) + D[i, j]
    check_symeq(R.op.body[0], resbody(*[iv.var for iv in R.op.axis]))

    R = InlineTensors(E, [B, C])
    resbody = lambda i, j: A[2, j] + (A[i, 2] + i*2) + D[i, j]
    check_symeq(R.op.body[0], resbody(*[iv.var for iv in R.op.axis]))

    R = InlineTensors(F)
    resbody = lambda i, j: tvm.exp(i + j) + i * (i + j)
    check_symeq(R.op.body[0], resbody(*[iv.var for iv in R.op.axis]))

    R = InlineTensors(F, [A])
    resbody = lambda i, j: tvm.exp(i + j) + B[i, (i + j)]
    check_symeq(R.op.body[0], resbody(*[iv.var for iv in R.op.axis]))

    R = InlineTensors(F, [B])
    resbody = lambda i, j: tvm.exp(A[i, j]) + i * A[i, j]
    check_symeq(R.op.body[0], resbody(*[iv.var for iv in R.op.axis]))

def test_solve_system_of_equations():
    seed = random.randrange(sys.maxsize)
    print("\nseed: {}\n".format(seed))
    random.seed(seed)

    def _check(variables, formulas, coef=(-5, 5), bounds=(-20, 20)):
        vs = [tvm.var("x" + str(i)) for i in range(variables)]

        fs = []
        for i in range(formulas):
            s1 = sum([v*random.randint(coef[0], coef[1]) for v in vs])
            s1 += random.randint(coef[0], coef[1])
            s2 = sum([v*random.randint(coef[0], coef[1]) for v in vs])
            s2 += random.randint(coef[0], coef[1])
            if random.random() < 0.7:
                op = tvm.expr.EQ
            else:
                op = random.choice([tvm.expr.LE, tvm.expr.LT, tvm.expr.GE, tvm.expr.GT])
            fs.append(op(s1, s2))

        vranges = {v: tvm.Range(bounds[0], bounds[1] + 1) for v in vs}

        domain = tvm.arith.Domain(vs, fs, vranges)
        domain_transform = SolveSystemOfEquations(domain)
        print("old", domain.conditions)
        print("new", domain_transform.new_domain.conditions)
        print(domain_transform.old_to_new)
        print(domain_transform.new_to_old)
        print()

        check_domain_transformation(domain_transform)

        # leaving some variables as parameters should also be ok
        for k in [1, 2]:
            if len(vs) > k:
                domain = tvm.arith.Domain(vs[:-k], fs, vranges)
                domain_transform = SolveSystemOfEquations(domain)
                print("old", domain.conditions)
                print("new", domain_transform.new_domain.conditions)
                print(domain_transform.old_to_new)
                print(domain_transform.new_to_old)
                print()

                param_ranges = {v: vranges[v] for v in vs[-k:]}
                check_domain_transformation(domain_transform, param_ranges)

    for i in range(3):
        _check(1, 1)
    for i in range(3):
        _check(1, 2)

    for i in range(3):
        _check(2, 1)
    for i in range(3):
        _check(2, 2)
    for i in range(3):
        _check(2, 3)

    for i in range(5):
        _check(3, 3, coef=(-2,2))
    for i in range(5):
        _check(3, 4, coef=(-2,2))

    for i in range(5):
        _check(4, 3, coef=(-1,1))

    for i in range(5):
        _check(10, 2, coef=(-1,1), bounds=(0, 4))
    for i in range(5):
        _check(10, 3, coef=(0,1), bounds=(0, 4))

def test_solve_system_of_inequalities():
    seed = random.randrange(sys.maxsize)
    print("\nseed: {}\n".format(seed))
    random.seed(seed)

    def _check(variables, formulas, coef=(-5, 5), bounds=(-20, 20)):
        vs = [tvm.var("x" + str(i)) for i in range(variables)]

        fs = []
        for i in range(formulas):
            s1 = sum([v*random.randint(coef[0], coef[1]) for v in vs])
            s1 += random.randint(coef[0], coef[1])
            s2 = sum([v*random.randint(coef[0], coef[1]) for v in vs])
            s2 += random.randint(coef[0], coef[1])
            op = random.choice([tvm.expr.EQ, tvm.expr.LE, tvm.expr.LT, tvm.expr.GE, tvm.expr.GT])
            fs.append(op(s1, s2))

        vranges = {v: tvm.Range(bounds[0], bounds[1] + 1) for v in vs}

        before = all(*fs)
        print(before)
        after = all(*SolveSystemOfInequalities(fs, vs, vranges))
        print(after)
        print()

        check_bruteforce(before == after, vranges)

    for i in range(3):
        _check(1, 1)
    for i in range(3):
        _check(1, 2)

    for i in range(3):
        _check(2, 1)
    for i in range(3):
        _check(2, 2)
    for i in range(3):
        _check(2, 3)

    # Somewhere here coefficients in the results become too large, leading to overflow,
    # so we use smaller initial coefficients

    for i in range(5):
        _check(3, 3, coef=(-2,2))
    for i in range(5):
        _check(3, 4, coef=(-2,2))

    for i in range(5):
        _check(4, 3, coef=(-1,1))

    for i in range(5):
        _check(10, 2, coef=(-1,1), bounds=(0, 4))
    for i in range(5):
        _check(10, 3, coef=(0,1), bounds=(0, 4))

def test_simplify_domain():
    # Note that here we test both SimplifyDomain and SimplifyReductionDomain.
    def _check(cond, axis, volume, vranges={}):
        vranges_with_axis = dict(vranges)
        vranges_with_axis.update({iv.var: iv.dom for iv in axis})
        variables = [iv.var for iv in axis]
        domain = Domain(variables, cond, vranges_with_axis)
        domain_tr = SimplifyDomain(domain)

        print("old", domain)
        print("new", domain_tr.new_domain)
        print("old_to_new", domain_tr.old_to_new)
        print("new_to_old", domain_tr.new_to_old)
        print()

        check_domain_transformation(domain_tr, vranges)

        # Also check SimplifyReductionDomain
        reduction = xor_combiner(sum([v*(i + 1) for i, v in enumerate(axis)]), axis)
        new_reduction = SimplifyReductionDomain(reduction, vranges)
        check_bruteforce(reduction == new_reduction, vranges)

        if volume is not None:
            vol = np.prod([domain_tr.new_domain.ranges[v].extent.value
                           for v in domain_tr.new_domain.variables])
            if vol != volume:
                raise AssertionError("New volume is {} != {}\n"
                                     "Old domain {}\nNew domain {}"
                                     .format(vol, volume, domain, domain_tr.new_domain))

    k = tvm.reduce_axis((0, 5), name="k")
    l = tvm.reduce_axis((0, 5), name="l")
    n = tvm.reduce_axis((0, 5), name="n")

    _check((k <= l), [k, l, n], 125)
    _check((k < l), [k, l, n], 80)
    _check(tvm.expr.EQ(k, l), [k, l, n], 25)
    _check(all(tvm.expr.EQ(k, l), (l < n)), [k, l, n], 16)
    _check(tvm.expr.EQ(2*l, k), [k, l, n], 15)
    _check(tvm.expr.EQ(2*l, k), [n, l, k], 15)
    _check(all(l - k < 2, 2*n == k), [k, l, n], 15)
    _check(all(l - k < 2, l >= k), [k, l, n], 50)

    some_var = tvm.var('some_var')
    _check(all(l - k < some_var, l >= k), [k, l, n], 50, {some_var: tvm.Range(0, 3)})
    _check(all(l - k < some_var, l >= k), [k, l, n], 25, {some_var: tvm.Range(0, 2)})

    k = tvm.reduce_axis((-3, 2), name="k")
    l = tvm.reduce_axis((-3, 2), name="l")
    n = tvm.reduce_axis((-3, 2), name="n")

    _check((k < l), [k, l, n], 80)
    _check(tvm.expr.EQ(k, l), [k, l, n], 25)
    _check(all(tvm.expr.EQ(k, l), (l < n)), [k, l, n], 16)
    # Now there are only two possible values for l: {l = -1, k = -2} and {l = 0, k = 0}
    _check(tvm.expr.EQ(2*l, k), [k, l, n], 10)
    _check(tvm.expr.EQ(2*l, k), [n, l, k], 10)
    _check(all(l - k < 2, 2*n == k), [k, l, n], 10)
    _check(all(l - k < 2, l >= k), [k, l, n], 50)

    some_var = tvm.var('some_var')
    _check(all(l - k < some_var, l >= k), [k, l, n], 50, {some_var: tvm.Range(0, 3)})
    _check(all(l - k < some_var, l >= k), [k, l, n], 25, {some_var: tvm.Range(0, 2)})

    k = tvm.reduce_axis((0, 3), name="k")
    m = tvm.reduce_axis((0, 2), name="m")
    x = tvm.reduce_axis((0, 4), name="x")
    _check(all(k + m*3 == x), [k, m, x], 4)

    k = tvm.reduce_axis((0, 6), name="k")
    l = tvm.reduce_axis((0, 5), name="l")
    n = tvm.reduce_axis((0, 30), name="n")

    _check(all(k + l*6 == n), [k, l, n], 30)
    _check(all(k + l*6 == n), [n, k, l], 30)
    _check(all(k + l*6 == n), [n, l, k], 30)

    _check(all(tvm.truncdiv(n, 5) == k, tvm.truncmod(n, 5) == l), [l, k, n], 30)
    _check(all(tvm.truncdiv(n, 5) == k, tvm.truncmod(n, 5) == l), [n, l, k], 30)
    _check(all(tvm.truncdiv(n, 5) == k, tvm.truncmod(n, 5) == l), [n, k, l], 30)

    _check(all(n // 5 == k, n % 5 == l), [l, k, n], 30)
    _check(all(n // 5 == k, n % 5 == l), [n, l, k], 30)
    _check(all(n // 5 == k, n % 5 == l), [n, k, l], 30)

    k = tvm.reduce_axis((0, 10), name="k")
    l = tvm.reduce_axis((0, 10), name="l")
    _check(all(tvm.truncmod(l + k, 3) <= 1, tvm.truncdiv(l + k, 3) <= 2), [l, k], 48)
    _check(all((l + k)%3 <= 1, (l + k)//3 <= 2), [l, k], 48)

    # Some real-life examples (check only correctness)
    # 6400 -> 2916
    jac_i3 = tvm.reduce_axis((0, 10), name="jac_i3")
    jac_i0 = tvm.reduce_axis((0, 1), name="jac_i0")
    xx = tvm.reduce_axis((0, 4), name="xx")
    yy = tvm.reduce_axis((0, 4), name="yy")
    jac_i2 = tvm.reduce_axis((0, 10), name="jac_i2")
    ff = tvm.reduce_axis((0, 2), name="ff")
    jac_i1 = tvm.reduce_axis((0, 2), name="jac_i1")
    nn = tvm.reduce_axis((0, 1), name="nn")
    _check(all(jac_i3 <= xx*2 + 2, jac_i2 <= yy*2 + 2, yy*2 <= jac_i2, xx*2 <= jac_i3),
           [nn, ff, yy, xx, jac_i0, jac_i1, jac_i2, jac_i3],
           None)

    # 32 -> 8
    n1_k1 = tvm.reduce_axis((0, 2), name="n1_k1")
    n0 = tvm.reduce_axis((0, 1), name="n0")
    n2_k2 = tvm.reduce_axis((0, 4), name="n2_k2")
    n3_k3 = tvm.reduce_axis((0, 4), name="n3_k3")
    ax3 = tvm.var("ax3")
    ax2 = tvm.var("ax2")
    ax1 = tvm.var("ax1")
    ax0 = tvm.var("ax0")
    _check(all(ax3 <= n3_k3*2 + 2, ax2 <= n2_k2*2 + 2, n3_k3*2 <= ax3, n2_k2*2 <= ax2),
           [n0, n1_k1, n2_k2, n3_k3],
           None,
           {ax3: tvm.Range(0, 10),
            ax2: tvm.Range(0, 10),
            ax1: tvm.Range(0, 2),
            ax0: tvm.Range(0, 1)})

def test_extract_as_tensor_maybe():
    def _check(shape, fcompute, volume=None, vranges={}):
        def fcompute_extracted(*variables):
            vranges_updated = dict(vranges)
            vranges_updated.update({v: tvm.Range(0, s) for v, s in zip(variables, shape)})
            expr = fcompute(*variables)
            if isinstance(expr, tvm.expr.Select):
                new_true_value = ExtractAsTensorMaybe(expr.true_value,
                                                      expr.condition,
                                                      variables,
                                                      vranges_updated)
                expr = tvm.expr.Select(expr.condition,
                                       new_true_value,
                                       expr.false_value)
                if volume is not None:
                    assert isinstance(new_true_value, tvm.expr.Call)
                    vol = np.prod([iv.dom.extent.value for iv in new_true_value.func.axis])
                    if vol != volume:
                        raise AssertionError("New volume is {} != {}"
                                             .format(vol, volume))
            return expr

        A = tvm.compute(shape, fcompute)
        B = tvm.compute(shape, fcompute_extracted)
        check_eq(A, B, [])

    _check((10, 10), lambda i, j: tvm.expr.Select(i < 3, i + j, 0), volume=30)
    _check((10, 10), lambda i, j: tvm.expr.Select(i < 3, j, 0), volume=10)
    _check((10, 10), lambda i, j: tvm.expr.Select(i < 3, i, 0), volume=3)
    _check((10, 10), lambda i, j: tvm.expr.Select(all(i < j, j < 5), i + j, 0), volume=16)
    # This one doesn't get extracted
    _check((10, 10), lambda i, j: tvm.expr.Select(i <= j, i + j, 0))

def test_extract_reductions():
    k = tvm.reduce_axis((0, 10), name="k")
    l = tvm.reduce_axis((0, 10), name="l")
    n = tvm.reduce_axis((0, 10), name="n")

    A = tvm.compute((10, 10),
                    lambda i, j:
                        ExtractReductions(sum_combiner(i + k + xor_combiner(j*k + l, l), k),
                                          [i, j],
                                          {i: tvm.Range(0, 10), j: tvm.Range(0, 10)}))
    B = tvm.compute((10, 10), lambda j, k: xor_combiner(j*k + l, l))
    C = tvm.compute((10, 10), lambda i, j: sum_combiner(i + k + B[j, k], k))
    check_eq(C, A, [])

    fcompute = lambda i, j: \
        ExtractReductions(sum_both_combiner((prod_derivative_combiner((i*n + 2*k, j + k), k)[1],
                                             xor_combiner(j*n + l, l)), n)[1],
                          [i, j],
                          {i: tvm.Range(0, 10), j: tvm.Range(0, 10)})
    A = tvm.compute((10, 10), fcompute)
    _, B = tvm.compute((10, 10, 10),
                       lambda i, j, n: prod_derivative_combiner((i*n + 2*k, j + k), k))
    C = tvm.compute((10, 10), lambda j, n: xor_combiner(j*n + l, l))
    _, D = tvm.compute((10, 10), lambda i, j: sum_both_combiner((B[i, j, n], C[j, n]), n))
    check_eq(A, D, [])

def test_optimize_and_lift_nonzeroness():
    k = tvm.reduce_axis((0, 10), name="k")
    l = tvm.reduce_axis((0, 10), name="l")
    n = tvm.reduce_axis((0, 10), name="n")
    A = tvm.placeholder((10, 10), name="A")

    zero = tvm.const(0, 'float32')

    B = compute((10, 10), lambda i, j: tvm.sum((i == j)*A[i, k] + A[k, j]*(i == j), k))
    B = OptimizeAndLiftNonzeronessConditions(B)
    R = lambda i, j: tvm.expr.Select(i == j,
                                     tvm.sum(A[j, k] + A[k, j], k),
                                     zero)
    check_tensor_symeq(B, R)

    # TODO: This test is unstable: sometimes the resulting condition looks like
    #       (i == j)*(j == i) instead of (i == j)
    #  B = compute((10, 10), lambda i, j: tvm.sum((i == j)*(i == k)*A[i, k] +
    #                                             (i == j)*A[k, j]*(i == k), k))
    #  B = OptimizeAndLiftNonzeronessConditions(B)
    #  R = lambda i, j: tvm.expr.Select(i == j, A[j, j]*2.0, zero)
    #  check_tensor_symeq(B, R)

    B = compute((10, 10), lambda i, j: tvm.sum((i < j)*(j < k)*A[j, k], k))
    B = OptimizeAndLiftNonzeronessConditions(B)
    k1 = tvm.reduce_axis((2, 10), name="k1")
    R = compute((10, 10), lambda i, j:
                tvm.expr.Select(all(i < j, j < 10),
                                tvm.sum(tvm.expr.Select(j < k1, A[j, k1], zero), k1),
                                zero))
    check_eq(B, R, [A])
    assert estimate_performance(B) <= estimate_performance(R)

    # TODO: This one requires better propagation of equalities involving outer variables
    #  B = compute((10, 10), lambda i, j: tvm.sum((i <= j)*(j <= k)*A[j, k], k, where=(i >= k)))
    #  B = OptimizeAndLiftNonzeronessConditions(B)
    #  R = compute((10, 10), lambda i, j: tvm.expr.Select((i == j), A[i, i], zero))
    #  check_eq(B, R, [A])
    #  assert estimate_performance(B) <= estimate_performance(R)

    B = compute((10, 10),
                lambda i, j: prod_derivative_combiner((A[j, k], (i <= j)*(j < k)*A[i, k]), k)[1])
    B = OptimizeAndLiftNonzeronessConditions(B)
    R = compute((10, 10), lambda i, j:
                tvm.expr.Select(all(i <= j, j < 10),
                                prod_derivative_combiner((A[j, k], (j < k)*A[i, k]), k)[1],
                                zero))
    check_eq(B, R, [A])
    assert estimate_performance(B) <= estimate_performance(R)

    B = compute((10,), lambda i:
                tvm.sum(A[i, k]*tvm.any(all(i < 5, k < 6), all(i > 5, k > 4)), k))
    B = OptimizeAndLiftNonzeronessConditions(B)
    R = compute((10,), lambda i:
                tvm.expr.Select(tvm.any(i < 5, i > 5),
                                tvm.sum(A[i, k], k, where=all(tvm.any(i < 5, k > 4),
                                                                  tvm.any(i > 5, k < 6))),
                                zero))
    check_eq(B, R, [A])
    assert estimate_performance(B) <= estimate_performance(R)

    # Specifying ranges of parameters
    B = compute((10, 10), lambda i, j: sum_or_prod_combiner((i == j)*A[i, k] + A[k, j]*(i == j), k))
    B = OptimizeAndLiftNonzeronessConditions(B, {m_param: tvm.Range(-5, -3)})
    R = lambda i, j: tvm.expr.Select(i == j,
                                     tvm.sum(A[j, k] + A[k, j], k),
                                     zero)
    check_tensor_symeq(B, R)

    B = compute((10, 10), lambda i, j: tvm.sum(((i - k) <= m_param) * A[i, k], k))
    B = OptimizeAndLiftNonzeronessConditions(B, {m_param: tvm.Range(11, 20)})
    R = lambda i, j: tvm.sum(A[i, k], k)
    check_tensor_symeq(B, R)

if __name__ == "__main__":
    test_is_sum_combiner()
    test_can_factor_zero_from_combiner()
    test_lift_nonzeroness_condition()
    test_inline_tail_call()
    test_inline_tensors()
    test_solve_system_of_equations()
    test_solve_system_of_inequalities()
    test_simplify_domain()
    test_extract_as_tensor_maybe()
    test_extract_reductions()
    test_optimize_and_lift_nonzeroness()


#!/usr/bin/env python
#tensorflow 2 integration of MOCO

import argparse
import builtins
import math
import os
import random
import shutil
import time
import warnings

#tensorflow imports
import tensorflow as tf

""" import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim
import torch.multiprocessing as mp
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models """

# a = tf.constant([-1.0, 0.0, 1.0], dtype = tf.float32)
# b = tf.keras.activations.softsign(a)

# print(b.numpy())

model_names = sorted(name for name in models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(models.__dict__[name]))

parser = argparse.ArgumentParser(description="TensorFlow V2 ImageNet Training")
parser.add_argument('data', metavar="DIR",
                    help="path to dataset")

parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet50',
                    choices=model_names,
                    help='model_architecture:' + '|'.join(model_names) + '(default: resnet50)')

parser.add_argument('-j', '--workers', default=32, type=int, metavar='N',
                    help='number of data loading workers (default: 32)')
parser.add_argument('--epochs', default=200, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=256, type=int,
                    metavar='N',
                    help='mini-batch size (default: 256), this is the total '
                         'batch size of all GPUs on the current node when '
                         'using Data Parallel or Distributed Data Parallel')
parser.add_argument('--lr', '--learning-rate', default=0.03, type=float,
                    metavar='LR', help='initial learning rate', dest='lr')
parser.add_argument('--schedule', default=[120, 160], nargs='*', type=int,
                    help='learning rate schedule (when to drop lr by 10x)')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum of SGD solver')
parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)',
                    dest='weight_decay')
parser.add_argument('-p', '--print-freq', default=10, type=int,
                    metavar='N', help='print frequency (default: 10)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('--world-size', default=-1, type=int,
                    help='number of nodes for distributed training')
parser.add_argument('--rank', default=-1, type=int,
                    help='node rank for distributed training')
parser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str,
                    help='url used to set up distributed training')
parser.add_argument('--dist-backend', default='nccl', type=str,
                    help='distributed backend')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--gpu', default=None, type=int,
                    help='GPU id to use.')
parser.add_argument('--multiprocessing-distributed', action='store_true',
                    help='Use multi-processing distributed training to launch '
                         'N processes per node, which has N GPUs. This is the '
                         'fastest way to use PyTorch for either single node or '
                         'multi node data parallel training')

# moco specific configs:
parser.add_argument('--moco-dim', default=128, type=int,
                    help='feature dimension (default: 128)')
parser.add_argument('--moco-k', default=65536, type=int,
                    help='queue size; number of negative keys (default: 65536)')
parser.add_argument('--moco-m', default=0.999, type=float,
                    help='moco momentum of updating key encoder (default: 0.999)')
parser.add_argument('--moco-t', default=0.07, type=float,
                    help='softmax temperature (default: 0.07)')

# options for moco v2
parser.add_argument('--mlp', action='store_true',
                    help='use mlp head')
parser.add_argument('--aug-plus', action='store_true',
                    help='use moco v2 data augmentation')
parser.add_argument('--cos', action='store_true',
                    help='use cosine lr schedule')

def main():
    args  = parser.parse_args()

    if args.seed is not None:
        random.seed(args.seed)
        tf.random.set_seed(args.seed)
        warnings.warn('You have chosen to seed training. '
                      'This will turn on the CUDNN deterministic setting, '
                      'which can slow down your training considerably! '
                      'You may see unexpected behavior when restarting '
                      'from checkpoints.')

    if args.gpu is not None:
        warnings.warn('You have chosen a specific GPU. This will completely '
                      'disable data parallelism.')

    if args.dist_url == "env://" and args.world_size == - 1:
            args.world_size = int(os.environ["WORLD_SIZE"])

    args.distributed = args.world_size > 1 or args.multiprocessing_distributed

    ngpus_per_node = len(tf.config.list_physical_devices('GPU'))

    if args.multiprocessing_distributed:
                # Since we have ngpus_per_node processes per node, the total world_size
        # needs to be adjusted accordingly
        args.world_size = ngpus_per_node * args.world_size


from bs4 import BeautifulSoup as bs
import requests as rq

errors = []

def get_next_page_link(web_page_soup):
	next_page_tag = web_page_soup.find('a',{'rel':'next'})
	if next_page_tag is None:
		return False
	else:
		return next_page_tag.get('href')

def get_download_page_link(web_page_soup):
	return web_page_soup.find('a',{'id':'download'}).get('href')

def get_download_link(download_page_link):
	download_page_soup = None
	i=0
	while i<30:
		try:
			download_page_soup = bs(rq.get(download_page_link).content,'lxml')
		except Exception as e:
			i+=1
			print("connection problem. retry ",i)
		if download_page_soup is not None:
			break
	for x in download_page_soup.find_all('a'):
		if x.text == '' and '.mp4' in x.get('href'):
			return x.get('href')

def get_anime_download_links(first_episode_page_link):
	page_link = first_episode_page_link
	download_links = []
	while True:
		print("Scraping link:",page_link)
		web_page_soup = None
		i=0
		while i<30:
			try:
				web_page_soup = bs(rq.get(page_link).content,'lxml')
			except Exception as e:
				i+=1
				print("connection problem. retry ",i)

			if web_page_soup is not None:
				break;
		if web_page_soup is not None:
			individual_download_link = get_download_link(get_download_page_link(web_page_soup))
			if individual_download_link is not None:
				download_links.append(individual_download_link)
			else:
				download_links.append("/"*30+"\nError obtaining download link\n"+"\\"*30)
			next_page = get_next_page_link(web_page_soup)
		if not next_page:
			break
		else:
			page_link = next_page
	return download_links

if __name__ == '__main__':
	statring_page = input('Enter first episode page link:')
	file_name = input('\nEnter distination file name:')
	download_links = get_anime_download_links(statring_page)
	no_of_errors = 0
	with open(file_name,'a') as f:
		for link in download_links:
			if link is not None:
				f.write(link)
			else:
				no_of_errors+=1
			f.write("\n")
	print("No. of errors =",no_of_errors)


import os
import sys

ERROR = False

def main(function):
    try:
        arguments = sys.argv[1:]
        assert arguments
        for path in arguments:
            assert os.path.isdir(path)
        for path in arguments:
            engine(path, function)
    except:
        sys.stdout.write('Usage: %s <directory>' % os.path.basename(sys.argv[0]))

def engine(path, function):
    global ERROR
    for root, dirs, files in os.walk(path):
        for name in files:
            path = os.path.join(root, name)
            try:
                function(path)
            except:
                sys.stderr.write('%sError: %s' % (ERROR and '\n' or '', path))
                ERROR = True

def zero(path):
    size = os.path.getsize(path)
    if size:
        data = open(path, 'wb')
        todo = size
        if todo >= 2 ** 20:
            buff = '\x00' * 2 ** 20
            while todo >= 2 ** 20:
                data.write(buff)
                todo = size - data.tell()
        data.write('\x00' * todo)
        data.close()

if __name__ == '__main__':
    main(zero)

Python 2.X/ZERO/Command Line/Zero Programs/zero.py