<commit_before>import logging
import requests
import json
import hashlib

from .base import Provider as BaseProvider

LOGGER = logging.getLogger(__name__)

def ProviderParser(subparser):
    subparser.add_argument('--auth-key', help='specify the key to access the API')
    subparser.add_argument('--auth-secret', help='specify the secret to access the API')

# Implements the DNS GoDaddy provider.
# Some general remarks about this provider, because it uses a weirdly designed API.
# Indeed, there is no direct way to insert, update or delete a specific record. Furthermore, there is no unique identifier for a record.
# Instead GoDaddy use a replace approach: for a given set of records (everything, or matching particular type and/or name),
#   one can replace this set with the new set sent throught API.
# With this approach:
#   - adding a record consists in appending a record to the obtained set and call replace with the updated set,
#   - updating a record consists in modifying a record in the obtained set and call replace with the updated set,
#   - deleting a record consists in removing a record in the obtained set and call replace with the updated set.
# In parallel, as said before, there is no unique identifier. This provider then implement a pseudo-identifier,
#   to allow an easy update or delete using the '--identifier' lexicon parameter.
# But you need to call the 'list' command just before executing and update/delete action, because identifier value
#   is tied to the content of the record, and will change anytime something is changed in the record.
class Provider(BaseProvider):

    def __init__(self, options, engine_overrides=None):
        super(Provider, self).__init__(options, engine_overrides)
        self.domain_id = None
        self.api_endpoint = self.engine_overrides.get('api_endpoint', 'https://api.godaddy.com/v1')

    def authenticate(self):
        domain = self.options.get('domain')

        result = self._get('/domains/{0}'.format(domain))
        self.domain_id = result['domainId']

    def list_records(self, type=None, name=None, content=None):
        domain = self.options.get('domain')

        url = '/domains/{0}/records'.format(domain)
        if type:
            url += '/{0}'.format(type)
        if name:
            url += '/{0}'.format(self._relative_name(name))

        raws = self._get(url)

        records = []
        for raw in raws:
            records.append({
                'identifier': Provider._identifier(raw),
                'type': raw['type'],
                'name': self._full_name(raw['name']),
                'ttl': raw['ttl'],
                'content': raw['data']
            })

        if content:
            records = [record for record in records if record['data'] == content]

        LOGGER.debug('list_records: %s', records)

        return records

    def create_record(self, type, name, content):
        domain = self.options.get('domain')
        relative_name = self._relative_name(name)
        ttl = self.options.get('ttl')

        # Retrieve existing data for given type and name, and append a new record.
        records = self._get('/domains/{0}/records/{1}/{2}'.format(domain, type, relative_name))

        data = {'data': content}
        if ttl:
            data['ttl'] = ttl

        records.append(data)

        # Synchronize data with inserted record into DNS zone for given type and name.
        self._put('/domains/{0}/records/{1}/{2}'.format(domain, type, relative_name), records)

        LOGGER.debug('create_record: %s %s %s', type, name, content)

        return True

    def update_record(self, identifier, type=None, name=None, content=None):
        # No identifier is used with GoDaddy.
        # We can rely either:
        #   - only on type/name to get the relevant records, both of them are required or we will could update to much records ...,
        #   - or by the pseudo-identifier provided
        # Furthermore for type/name approach, we cannot update all matching records, as it would lead
        #   to an error (two entries of same type + name cannot have the same content).
        # So for type/name approach, we search first matching record for type/name on which content is different,
        #   and we update it before synchronizing the DNS zone.
        if not identifier and not type:
            raise Exception('ERROR: type is required')
        if not identifier and not name:
            raise Exception('ERROR: name is required')

        domain = self.options.get('domain')
        relative_name = self._relative_name(name)

        # Retrieve existing data for given type and name
        records = self._get('/domains/{0}/records/{1}/{2}'.format(domain, type, relative_name))

        # Get the record to update:
        #   - either explicitly by its identifier,
        #   - or the first matching by its type+name where content does not match (first match, see first method comment for explanation).
        for record in records:
            if ((identifier and Provider._identifier(record) == identifier) or
                record['type'] == type and self._relative_name(record['name']) == relative_name and record['data'] != content):
                record['data'] = content
                break

        # Synchronize data with updated records into DNS zone for given type and name.
        self._put('/domains/{0}/records/{1}/{2}'.format(domain, type, relative_name), records)

        LOGGER.debug('update_record: %s %s %s', type, name, content)

        return True

    def delete_record(self, identifier=None, type=None, name=None, content=None):
        # For the LOL. GoDaddy does not accept an empty array when updating a particular set of records.
        # It means that you cannot request to remove all records matching a particular type and/or name.
        # Instead, we get ALL records in the DNS zone, update the set, and replace EVERYTHING in the DNS zone.
        # You will always have at minimal NS/SRV entries in the array, otherwise your DNS zone is broken,
        #   and updating the zone is the least of your problem ...
        domain = self.options.get('domain')

        # Retrieve all records in the DNS zone
        records = self._get('/domains/{0}/records'.format(domain))

        relative_name = None
        if name:
            relative_name = self._relative_name(name)

        # Filter out all records which match the pattern (either identifier, or some combination of type/name/content).
        filtered_records = []
        if identifier:
            filtered_records = [record for record in records if Provider._identifier(record) != identifier]
        else:
            for record in records:
                if ((not type and not relative_name and not content) or
                    (type and not relative_name and not content and record['type'] != type) or
                    (not type and relative_name and not content and self._relative_name(record['name']) != relative_name) or
                    (not type and not relative_name and content and record['data'] != content) or
                    (type and relative_name and not content and (record['type'] != type or self._relative_name(record['name']) != relative_name)) or
                    (type and not relative_name and content and (record['type'] != type or record['data'] != content)) or
                    (not type and relative_name and content and (self._relative_name(record['name']) != relative_name or record['data'] != content)) or
                    (type and relative_name and content and (record['type'] != type or self._relative_name(record['name']) != relative_name or record['data'] != content))):
                    filtered_records.append(record)

        # Synchronize data with expurged entries into DNS zone.
        self._put('/domains/{0}/records'.format(domain, type, relative_name), filtered_records)

        LOGGER.debug('delete_records: %s %s %s', type, name, content)

        return True

    # GoDaddy provides no identifier for a record, which is a problem where identifiers can be used (delete and update).
    # To circumvent this, we implement a pseudo-identifier, which is basically a hash of type+name+content of a given record.
    # It is far from perfect, as the identifier will change each time we change something in the record ...
    # But at least, one can use 'lexicon godaddy list ...' then 'lexicon godaddy update --identifier ...' to modify specific record.
    # However, 'lexicon godaddy list ...' should be called each time DNS zone had been changed to calculate new identifiers.
    @staticmethod
    def _identifier(record):
        sha256 = hashlib.sha256()
        sha256.update(record.get('type', ''))
        sha256.update(record.get('name', ''))
        sha256.update(record.get('data', ''))
        return sha256.hexdigest()[0:7]

    def _request(self, action='GET', url='/', data=None, query_params=None):
        if not data:
            data = {}
        if not query_params:
            query_params = {}

        result = requests.request(action, self.api_endpoint + url,
                                  params=query_params,
                                  data=json.dumps(data),
                                  headers={
                                      'Content-Type': 'application/json',
                                      'Accept': 'application/json',
                                      # GoDaddy use a key/secret pair to authenticate
                                      'Authorization': 'sso-key {0}:{1}'.format(
                                          self.options.get('auth_key'),
                                          self.options.get('auth_secret'))
                                  })

        result.raise_for_status()

        try:
            # Return the JSON body response if exists.
            return result.json()
        except ValueError:
            # For some requests command (eg. PUT), GoDaddy will not return any JSON, just an HTTP status without body.
            return None
<commit_msg>Ready for integration tests.<commit_after>import logging
import requests
import json
import hashlib

from .base import Provider as BaseProvider

LOGGER = logging.getLogger(__name__)

def ProviderParser(subparser):
    subparser.add_argument('--auth-key', help='specify the key to access the API')
    subparser.add_argument('--auth-secret', help='specify the secret to access the API')

# Implements the DNS GoDaddy provider.
# Some general remarks about this provider, because it uses a weirdly designed API.
# Indeed, there is no direct way to insert, update or delete a specific record. Furthermore, there is no unique identifier for a record.
# Instead GoDaddy use a replace approach: for a given set of records one can replace this set with a new set sent throught API.
# For the sake of simplicity and consistency accoss the provider edit methods, the set will be always all records in the DNS zone.
# With this approach:
#   - adding a record consists in appending a record to the obtained set and call replace with the updated set,
#   - updating a record consists in modifying a record in the obtained set and call replace with the updated set,
#   - deleting a record consists in removing a record in the obtained set and call replace with the updated set.
# In parallel, as said before, there is no unique identifier. This provider then implement a pseudo-identifier,
#   to allow an easy update or delete using the '--identifier' lexicon parameter.
# But you need to call the 'list' command just before executing and update/delete action, because identifier value
#   is tied to the content of the record, and will change anytime something is changed in the record.
class Provider(BaseProvider):

    def __init__(self, options, engine_overrides=None):
        super(Provider, self).__init__(options, engine_overrides)
        self.domain_id = None
        self.api_endpoint = self.engine_overrides.get('api_endpoint', 'https://api.godaddy.com/v1')

    def authenticate(self):
        domain = self.options.get('domain')

        result = self._get('/domains/{0}'.format(domain))
        self.domain_id = result['domainId']

    def list_records(self, type=None, name=None, content=None):
        domain = self.options.get('domain')

        url = '/domains/{0}/records'.format(domain)
        if type:
            url += '/{0}'.format(type)
        if name:
            url += '/{0}'.format(self._relative_name(name))

        raws = self._get(url)

        records = []
        for raw in raws:
            records.append({
                'identifier': Provider._identifier(raw),
                'type': raw['type'],
                'name': self._full_name(raw['name']),
                'ttl': raw['ttl'],
                'content': raw['data']
            })

        if content:
            records = [record for record in records if record['data'] == content]

        LOGGER.debug('list_records: %s', records)

        return records

    def create_record(self, type, name, content):
        domain = self.options.get('domain')
        relative_name = self._relative_name(name)
        ttl = self.options.get('ttl')

        # Retrieve existing data in DNS zone.
        records = self._get('/domains/{0}/records'.format(domain))

        # Append a new entry corresponding to given parameters.
        data = {'type': type, 'name': relative_name, 'data': content}
        if ttl:
            data['ttl'] = ttl

        records.append(data)

        # Synchronize data with inserted record into DNS zone.
        self._put('/domains/{0}/records'.format(domain), records)

        LOGGER.debug('create_record: %s %s %s', type, name, content)

        return True

    def update_record(self, identifier, type=None, name=None, content=None):
        # No identifier is used with GoDaddy.
        # We can rely either:
        #   - only on type/name to get the relevant records, both of them are required or we will could update to much records ...,
        #   - or by the pseudo-identifier provided
        # Furthermore for type/name approach, we cannot update all matching records, as it would lead
        #   to an error (two entries of same type + name cannot have the same content).
        # So for type/name approach, we search first matching record for type/name on which content is different,
        #   and we update it before synchronizing the DNS zone.
        if not identifier and not type:
            raise Exception('ERROR: type is required')
        if not identifier and not name:
            raise Exception('ERROR: name is required')

        domain = self.options.get('domain')
        relative_name = None
        if name:
            relative_name = self._relative_name(name)

        # Retrieve existing data in DNS zone.
        records = self._get('/domains/{0}/records'.format(domain))

        # Get the record to update:
        #   - either explicitly by its identifier,
        #   - or the first matching by its type+name where content does not match (first match, see first method comment for explanation).
        for record in records:
            if ((identifier and Provider._identifier(record) == identifier) or
                (not identifier and record['type'] == type and self._relative_name(record['name']) == relative_name and record['data'] != content)):
                record['data'] = content
                break

        # Synchronize data with updated records into DNS zone.
        self._put('/domains/{0}/records'.format(domain), records)

        LOGGER.debug('update_record: %s %s %s', type, name, content)

        return True

    def delete_record(self, identifier=None, type=None, name=None, content=None):
        # For the LOL. GoDaddy does not accept an empty array when updating a particular set of records.
        # It means that you cannot request to remove all records matching a particular type and/or name.
        # Instead, we get ALL records in the DNS zone, update the set, and replace EVERYTHING in the DNS zone.
        # You will always have at minimal NS/SRV entries in the array, otherwise your DNS zone is broken,
        #   and updating the zone is the least of your problem ...
        domain = self.options.get('domain')

        # Retrieve all records in the DNS zone
        records = self._get('/domains/{0}/records'.format(domain))

        relative_name = None
        if name:
            relative_name = self._relative_name(name)

        # Filter out all records which match the pattern (either identifier, or some combination of type/name/content).
        filtered_records = []
        if identifier:
            filtered_records = [record for record in records if Provider._identifier(record) != identifier]
        else:
            for record in records:
                if ((not type and not relative_name and not content) or
                    (type and not relative_name and not content and record['type'] != type) or
                    (not type and relative_name and not content and self._relative_name(record['name']) != relative_name) or
                    (not type and not relative_name and content and record['data'] != content) or
                    (type and relative_name and not content and (record['type'] != type or self._relative_name(record['name']) != relative_name)) or
                    (type and not relative_name and content and (record['type'] != type or record['data'] != content)) or
                    (not type and relative_name and content and (self._relative_name(record['name']) != relative_name or record['data'] != content)) or
                    (type and relative_name and content and (record['type'] != type or self._relative_name(record['name']) != relative_name or record['data'] != content))):
                    filtered_records.append(record)

        # Synchronize data with expurged entries into DNS zone.
        self._put('/domains/{0}/records'.format(domain), filtered_records)

        LOGGER.debug('delete_records: %s %s %s', type, name, content)

        return True

    # GoDaddy provides no identifier for a record, which is a problem where identifiers can be used (delete and update).
    # To circumvent this, we implement a pseudo-identifier, which is basically a hash of type+name+content of a given record.
    # It is far from perfect, as the identifier will change each time we change something in the record ...
    # But at least, one can use 'lexicon godaddy list ...' then 'lexicon godaddy update --identifier ...' to modify specific record.
    # However, 'lexicon godaddy list ...' should be called each time DNS zone had been changed to calculate new identifiers.
    @staticmethod
    def _identifier(record):
        sha256 = hashlib.sha256()
        sha256.update(record.get('type', '').encode('utf-8'))
        sha256.update(record.get('name', '').encode('utf-8'))
        sha256.update(record.get('data', '').encode('utf-8'))
        return sha256.hexdigest()[0:7]

    def _request(self, action='GET', url='/', data=None, query_params=None):
        if not data:
            data = {}
        if not query_params:
            query_params = {}

        result = requests.request(action, self.api_endpoint + url,
                                  params=query_params,
                                  data=json.dumps(data),
                                  headers={
                                      'Content-Type': 'application/json',
                                      'Accept': 'application/json',
                                      # GoDaddy use a key/secret pair to authenticate
                                      'Authorization': 'sso-key {0}:{1}'.format(
                                          self.options.get('auth_key'),
                                          self.options.get('auth_secret'))
                                  })

        result.raise_for_status()

        try:
            # Return the JSON body response if exists.
            return result.json()
        except ValueError:
            # For some requests command (eg. PUT), GoDaddy will not return any JSON, just an HTTP status without body.
            return None
<|endoftext|>


<commit_before># ====== Host all the text definition =======
import re

# the word is a sequence (more than one) of non-space character
from typing import List

# this is private because we want regex composition to be as few as possible
# use if you need to use this regex,
# please consider using the function provided below instead
_WORD_REGEX_STR = r'\S+'

# the left boundary of a word is either the start of the passage or a space
# this is private because we want regex composition to be as few as possible
# use if you need to use this regex,
# please consider using the function provided below instead
_SINGLE_LEFT_WORD_BOUNDARY_REGEX_STR = r'^|\s'

# the right boundary of a word is either the end of the passage or a space
# this is private because we want regex composition to be as few as possible
# use if you need to use this regex,
# please consider using the function provided below instead
_SINGLE_RIGHT_WORD_BOUNDARY_REGEX_STR = r'\s|$'

# ============== compiled version ==================
# the compiled version of word regex
WORD_REGEX = re.compile(_WORD_REGEX_STR, re.UNICODE)

def get_all_words_in_text(text: str) -> List[str]:
    """Get all the words in a given text in the order that they appear.

    :param text: the text to get all the word from
    :return: a list of words in the order that they appear in text
    """
    # the `split` and `strip` method handles all kinds of white spaces
    # including the
    return text.strip().split()

def get_single_word_count_in_text(text: str, word: str) -> int:
    """Get the number of single word appear in text.

    :param text: the whole text to find the word
    :param word: the word to find in the whole text
    :return: the number of times the `word` appear in `text`
    """
    return get_all_words_in_text(text).count(word)
<commit_msg>add additional definition<commit_after># ====== Host all the text definition =======
import re

# the word is a sequence (more than one) of non-space character
from typing import List

# this is private because we want regex composition to be as few as possible
# use if you need to use this regex,
# please consider using the function provided below instead
_WORD_REGEX_STR = r'\S+'

# the left boundary of a word is either the start of the passage or a space
# this is private because we want regex composition to be as few as possible
# use if you need to use this regex,
# please consider using the function provided below instead
_SINGLE_LEFT_WORD_BOUNDARY_REGEX_STR = r'^|\s'

# the right boundary of a word is either the end of the passage or a space
# this is private because we want regex composition to be as few as possible
# use if you need to use this regex,
# please consider using the function provided below instead
_SINGLE_RIGHT_WORD_BOUNDARY_REGEX_STR = r'\s|$'

# ============== compiled version ==================
# the compiled version of word regex
WORD_REGEX = re.compile(_WORD_REGEX_STR, re.UNICODE)

def get_all_words_in_text(text: str) -> List[str]:
    """Get all the words in a given text in the order that they appear.

    :param text: the text to get all the word from
    :return: a list of words in the order that they appear in text
    """
    # the `split` and `strip` method handles all kinds of white spaces
    # including the
    return text.strip().split()

def get_single_word_count_in_text(text: str, word: str) -> int:
    """Get the number of single word appear in text.

    :param text: the whole text to find the word
    :param word: the word to find in the whole text
    :return: the number of times the `word` appear in `text`
    """
    return get_all_words_in_text(text).count(word)

def get_words_with_right_boundary(text: str) -> List[str]:
    """Get the word with all its rightward spaces

    We will discard the spaces in front of the first word
    :param text: the input text
    :return: a list of string, each string start with a word,
        together with all of its space on the right
    """
    # this is a list where
    # the first element is a white space (or empty)
    # the second element is a word
    # the third element is a whitespace and so on.
    reg_split_list = re.split("(" + _WORD_REGEX_STR + ")", text)

    return [
        # joins the word with its right boundary
        # (all the space on the right side until another word)
        "".join(reg_split_list[start: start + 1])
        # we cannot use keyword parameter on range,
        # because stupid python (3.6.1) does not allow us to...
        # the range means: range(start=1, end=len(reg_split_list), step=2)
        for start in range(1, len(reg_split_list), 2)
    ]
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# This tries to fix bug 1357563 in a best effort way. It may find a proper page
# to redirect users to but if it fails it will raise a 404.
#
# Flow:
#
# 1. Check if the response is 404 and it's from a Lithium URL (/t5/*)
# 2. Checks cache, redirects if successful
# 3. Checks Lithium localized product_pages from LithiumRedirectionMiddleware.product_pages, redirects if successful
# 4. Fetches the page from Lithium and tries to extract title of the page.
#    Use og:title instead of HTML Title because Lithium strips the later down to 50 chars or something.
#    If it fails it returns 404
# 5a. Try to find a Document with the same title, redirect if successful
# 5b. Try to find a Topic with the same title under a product. Product is also extracted by title. For non-English pages we need to do a reverse translation lookup from TOPICS
# 5c. Try to find a Product. For non-English pages we need to do a reverse translation lookup from PRODUCTS
#
# In all cases we store a successful redirect in the cache.
# In all cases we append ?cache=no in the redirect to bust the cache
# All redirects are 302s
#
#

import hashlib
import re
import urllib

from django.core.cache import cache
from django.conf import settings
from django.http import HttpResponseRedirect
from django.utils.translation import pgettext
from django.utils import translation

import requests
from raven.contrib.django.raven_compat.models import client as sentry_client

from kitsune.products.models import Product, Topic
from kitsune.sumo.views import handle404
from kitsune.wiki.models import Document

# build a list of topics translated for reverse lookup
TOPICS = {}
for topic in set(Topic.objects.values_list('title', flat=True)):
    for lang in settings.SUMO_LANGUAGES:
        with translation.override(lang):
            trans_text = pgettext('DB: products.Topic.title', topic)
        TOPICS[trans_text] = topic

PRODUCTS = {}
for product in set(Product.objects.values_list('title', flat=True)):
    for lang in settings.SUMO_LANGUAGES:
        with translation.override(lang):
            trans_text = pgettext('DB: products.Product.title', product)
        PRODUCTS[trans_text] = product

class LithiumRedirectionMiddleware():

    meta_regex = re.compile('<meta content="(.*?)" property="(.*?)"/>')
    product_pages = {
        u'Mozilla-Hilfe-Deutsch': '/de/?cache=no',
        u'Soporte-de-Mozilla-Español': '/es/?cache=no',
        u'Magyar': '/hu/?cache=no',
        u'Pomoc-Mozilli-polski': '/pl/?cache=no',
        u'Podpora-Mozilly-slovenčina': '/sk/?cache=no',
        u'Поддержка-Mozilla-Русский': '/ru/?cache=no',
        u'Podpora-Mozilly-Čeština': '/cs/?cache=no',
        u'Mozilla-Support-Ελληνικά': '/el/?cache=no',
        u'Mozilla-tuki-suomi': '/fi/?cache=no',
        u'Dukungan-Mozilla-Bahasa': '/hi-IN/?cache=no',
        u'Suporte-Mozilla-Português-do': '/pt-br/?cache=no',
        u'Mozilla-Destek-Türkçe': '/tr/?cache=no',
        u'Підтримка-Mozilla-Українська': '/uk/?cache=no',
        u'Mozilla-Support-Dansk': '/da/?cache=no',
        u'Mozilla-Support-English': '/en/?cache=no',
        u'Assistance-de-Mozilla-Français': '/fr/?cache=no',
        u'Supporto-Mozilla-Italiano': '/it/?cache=no',
        u'Mozilla-Support-Nederlands': '/nl/?cache=no',
        u'Apoio-da-Mozilla-Português': '/pt/?cache=no',
        u'Mozilla-Support-Svenska': '/sv/?cache=no',
        'Mozilla-\xe3\x82\xb5\xe3\x83\x9d\xe3\x83\xbc\xe3\x83\x88-\xe6\x97\xa5\xe6\x9c\xac\xe8\xaa\x9e': '/ja/?cache=no',
        'Mozilla-\xe6\x8a\x80\xe8\xa1\x93\xe6\x94\xaf\xe6\x8f\xb4-\xe6\xad\xa3\xe9\xab\x94\xe4\xb8\xad\xe6\x96\x87-\xe7\xb9\x81\xe9\xab\x94': '/zh-TW/?cache=no',
        'Mozilla-\xeb\x8f\x84\xec\x9b\x80\xeb\xa7\x90-\xed\x95\x9c\xea\xb5\xad\xec\x96\xb4': '/ko/?cache=no',
        'Mozilla-\xe6\x8a\x80\xe6\x9c\xaf\xe6\x94\xaf\xe6\x8c\x81-\xe4\xb8\xad\xe6\x96\x87-\xe7\xae\x80\xe4\xbd\x93': '/zh-CN/?cache=no',
        '\xd8\xaf\xd8\xb9\xd9\x85-\xd9\x81\xd9\x8e\xd9\x8a\xd9\x8e\xd8\xb1\xd9\x81\xd9\x8f\xd9\x83\xd8\xb3-\xd8\xb9\xd8\xb1\xd8\xa8\xd9\x8a': '/ar/?cache=no',
    }

    extra_pages = {
        '/t5/How-To/Cookies/ta-p/16348': '/kb/cookies-information-websites-store-on-your-computer',
        '/t5/Fix-slowness-crashing-error/What-does-quot-Your-connection-is-not-secure-quot-mean/ta-p/30354': '/kb/what-does-your-connection-is-not-secure-mean',
        '/t5/Firefox-for-Android/ct-p/Firefox-Android': '/products/mobile',
    }

    def process_response(self, request, response):
        if (response.status_code != 404):
            return response

        urlparts = request.path_info.split('/')
        try:
            if not urlparts[1].startswith('t5'):
                return response
        except KeyError:
            return response

        # Oh boy, Lithium URL. Let's see if we can turn things around.

        # Cache first
        cache_key = hashlib.sha1(request.path_info.encode('utf-8')).hexdigest()
        cached_redirect = cache.get(cache_key)
        if cached_redirect:
            return HttpResponseRedirect(cached_redirect)

        if request.path_info in self.extra_pages:
            return HttpResponseRedirect(self.extra_pages[request.path_info])

        if urlparts[2] in self.product_pages:
            return HttpResponseRedirect(self.product_pages[urlparts[2]])

        path = request.path_info
        url = 'https://secure02.lithium.com.edgekey.net' + path
        try:
            response = requests.get(
                url,
                headers={'Host': 'support.mozilla.org'},
                timeout=3,
                verify=False,
            )
        except requests.exceptions.Timeout:
            sentry_client.captureException()
            return handle404(request)

        try:
            response.raise_for_status()
        except requests.exceptions.RequestException:
            sentry_client.captureException()
            return handle404(request)
        else:
            for content, prop in self.meta_regex.findall(response.content):
                try:
                    content = content.strip().decode('utf-8')
                    prop = prop.strip().decode('utf-8')
                except UnicodeError:
                    break

                if prop == 'og:title':
                    # order of lookup is important
                    try:
                        bar = Document.objects.get(title=content)
                    except (Document.DoesNotExist, Document.MultipleObjectsReturned):
                        pass
                    except:
                        sentry_client.captureException()
                    else:
                        url = bar.get_absolute_url() + '?cache=no'
                        cache.set(cache_key, url, None)
                        return HttpResponseRedirect(url)

                    if '-' in content:
                        product, title = content.split('-', 1)
                        product = product.strip()
                        title = title.strip()

                        title = TOPICS.get(title, title)
                        product = PRODUCTS.get(product, product)

                        try:
                            bar = Topic.objects.get(title=title, product__title=product)
                        except (Topic.DoesNotExist, Topic.MultipleObjectsReturned):
                            pass
                        except:
                            sentry_client.captureException()
                        else:
                            url = bar.get_absolute_url() + '?cache=no'
                            cache.set(cache_key, url, None)
                            return HttpResponseRedirect(url)

                    product = PRODUCTS.get(content, content)
                    try:
                        bar = Product.objects.get(title=product)
                    except (Product.DoesNotExist, Product.MultipleObjectsReturned):
                        pass
                    except:
                        sentry_client.captureException()
                    else:
                        url = bar.get_absolute_url() + '?cache=no'
                        cache.set(cache_key, url, None)
                        return HttpResponseRedirect(url)

        # everything failed, kittens die, it's a 404
        return handle404(request)
<commit_msg>[fix bug 1358573] Add redirect to disable-or-remove-add-ons<commit_after># -*- coding: utf-8 -*-
#
# This tries to fix bug 1357563 in a best effort way. It may find a proper page
# to redirect users to but if it fails it will raise a 404.
#
# Flow:
#
# 1. Check if the response is 404 and it's from a Lithium URL (/t5/*)
# 2. Checks cache, redirects if successful
# 3. Checks Lithium localized product_pages from LithiumRedirectionMiddleware.product_pages, redirects if successful
# 4. Fetches the page from Lithium and tries to extract title of the page.
#    Use og:title instead of HTML Title because Lithium strips the later down to 50 chars or something.
#    If it fails it returns 404
# 5a. Try to find a Document with the same title, redirect if successful
# 5b. Try to find a Topic with the same title under a product. Product is also extracted by title. For non-English pages we need to do a reverse translation lookup from TOPICS
# 5c. Try to find a Product. For non-English pages we need to do a reverse translation lookup from PRODUCTS
#
# In all cases we store a successful redirect in the cache.
# In all cases we append ?cache=no in the redirect to bust the cache
# All redirects are 302s
#
#

import hashlib
import re
import urllib

from django.core.cache import cache
from django.conf import settings
from django.http import HttpResponseRedirect
from django.utils.translation import pgettext
from django.utils import translation

import requests
from raven.contrib.django.raven_compat.models import client as sentry_client

from kitsune.products.models import Product, Topic
from kitsune.sumo.views import handle404
from kitsune.wiki.models import Document

# build a list of topics translated for reverse lookup
TOPICS = {}
for topic in set(Topic.objects.values_list('title', flat=True)):
    for lang in settings.SUMO_LANGUAGES:
        with translation.override(lang):
            trans_text = pgettext('DB: products.Topic.title', topic)
        TOPICS[trans_text] = topic

PRODUCTS = {}
for product in set(Product.objects.values_list('title', flat=True)):
    for lang in settings.SUMO_LANGUAGES:
        with translation.override(lang):
            trans_text = pgettext('DB: products.Product.title', product)
        PRODUCTS[trans_text] = product

class LithiumRedirectionMiddleware():

    meta_regex = re.compile('<meta content="(.*?)" property="(.*?)"/>')
    product_pages = {
        u'Mozilla-Hilfe-Deutsch': '/de/?cache=no',
        u'Soporte-de-Mozilla-Español': '/es/?cache=no',
        u'Magyar': '/hu/?cache=no',
        u'Pomoc-Mozilli-polski': '/pl/?cache=no',
        u'Podpora-Mozilly-slovenčina': '/sk/?cache=no',
        u'Поддержка-Mozilla-Русский': '/ru/?cache=no',
        u'Podpora-Mozilly-Čeština': '/cs/?cache=no',
        u'Mozilla-Support-Ελληνικά': '/el/?cache=no',
        u'Mozilla-tuki-suomi': '/fi/?cache=no',
        u'Dukungan-Mozilla-Bahasa': '/hi-IN/?cache=no',
        u'Suporte-Mozilla-Português-do': '/pt-br/?cache=no',
        u'Mozilla-Destek-Türkçe': '/tr/?cache=no',
        u'Підтримка-Mozilla-Українська': '/uk/?cache=no',
        u'Mozilla-Support-Dansk': '/da/?cache=no',
        u'Mozilla-Support-English': '/en/?cache=no',
        u'Assistance-de-Mozilla-Français': '/fr/?cache=no',
        u'Supporto-Mozilla-Italiano': '/it/?cache=no',
        u'Mozilla-Support-Nederlands': '/nl/?cache=no',
        u'Apoio-da-Mozilla-Português': '/pt/?cache=no',
        u'Mozilla-Support-Svenska': '/sv/?cache=no',
        'Mozilla-\xe3\x82\xb5\xe3\x83\x9d\xe3\x83\xbc\xe3\x83\x88-\xe6\x97\xa5\xe6\x9c\xac\xe8\xaa\x9e': '/ja/?cache=no',
        'Mozilla-\xe6\x8a\x80\xe8\xa1\x93\xe6\x94\xaf\xe6\x8f\xb4-\xe6\xad\xa3\xe9\xab\x94\xe4\xb8\xad\xe6\x96\x87-\xe7\xb9\x81\xe9\xab\x94': '/zh-TW/?cache=no',
        'Mozilla-\xeb\x8f\x84\xec\x9b\x80\xeb\xa7\x90-\xed\x95\x9c\xea\xb5\xad\xec\x96\xb4': '/ko/?cache=no',
        'Mozilla-\xe6\x8a\x80\xe6\x9c\xaf\xe6\x94\xaf\xe6\x8c\x81-\xe4\xb8\xad\xe6\x96\x87-\xe7\xae\x80\xe4\xbd\x93': '/zh-CN/?cache=no',
        '\xd8\xaf\xd8\xb9\xd9\x85-\xd9\x81\xd9\x8e\xd9\x8a\xd9\x8e\xd8\xb1\xd9\x81\xd9\x8f\xd9\x83\xd8\xb3-\xd8\xb9\xd8\xb1\xd8\xa8\xd9\x8a': '/ar/?cache=no',
    }

    extra_pages = {
        '/t5/How-To/Cookies/ta-p/16348': '/kb/cookies-information-websites-store-on-your-computer',
        '/t5/Fix-slowness-crashing-error/What-does-quot-Your-connection-is-not-secure-quot-mean/ta-p/30354': '/kb/what-does-your-connection-is-not-secure-mean',
        '/t5/Firefox-for-Android/ct-p/Firefox-Android': '/products/mobile',
        '/t5/Manage-preferences-and-add-ons/Disable-or-remove-Add-ons/ta-p/1000': '/kb/disable-or-remove-add-ons',
    }

    def process_response(self, request, response):
        if (response.status_code != 404):
            return response

        urlparts = request.path_info.split('/')
        try:
            if not urlparts[1].startswith('t5'):
                return response
        except KeyError:
            return response

        # Oh boy, Lithium URL. Let's see if we can turn things around.

        # Cache first
        cache_key = hashlib.sha1(request.path_info.encode('utf-8')).hexdigest()
        cached_redirect = cache.get(cache_key)
        if cached_redirect:
            return HttpResponseRedirect(cached_redirect)

        if request.path_info in self.extra_pages:
            return HttpResponseRedirect(self.extra_pages[request.path_info])

        if urlparts[2] in self.product_pages:
            return HttpResponseRedirect(self.product_pages[urlparts[2]])

        path = request.path_info
        url = 'https://secure02.lithium.com.edgekey.net' + path
        try:
            response = requests.get(
                url,
                headers={'Host': 'support.mozilla.org'},
                timeout=3,
                verify=False,
            )
        except requests.exceptions.Timeout:
            sentry_client.captureException()
            return handle404(request)

        try:
            response.raise_for_status()
        except requests.exceptions.RequestException:
            sentry_client.captureException()
            return handle404(request)
        else:
            for content, prop in self.meta_regex.findall(response.content):
                try:
                    content = content.strip().decode('utf-8')
                    prop = prop.strip().decode('utf-8')
                except UnicodeError:
                    break

                if prop == 'og:title':
                    # order of lookup is important
                    try:
                        bar = Document.objects.get(title=content)
                    except (Document.DoesNotExist, Document.MultipleObjectsReturned):
                        pass
                    except:
                        sentry_client.captureException()
                    else:
                        url = bar.get_absolute_url() + '?cache=no'
                        cache.set(cache_key, url, None)
                        return HttpResponseRedirect(url)

                    if '-' in content:
                        product, title = content.split('-', 1)
                        product = product.strip()
                        title = title.strip()

                        title = TOPICS.get(title, title)
                        product = PRODUCTS.get(product, product)

                        try:
                            bar = Topic.objects.get(title=title, product__title=product)
                        except (Topic.DoesNotExist, Topic.MultipleObjectsReturned):
                            pass
                        except:
                            sentry_client.captureException()
                        else:
                            url = bar.get_absolute_url() + '?cache=no'
                            cache.set(cache_key, url, None)
                            return HttpResponseRedirect(url)

                    product = PRODUCTS.get(content, content)
                    try:
                        bar = Product.objects.get(title=product)
                    except (Product.DoesNotExist, Product.MultipleObjectsReturned):
                        pass
                    except:
                        sentry_client.captureException()
                    else:
                        url = bar.get_absolute_url() + '?cache=no'
                        cache.set(cache_key, url, None)
                        return HttpResponseRedirect(url)

        # everything failed, kittens die, it's a 404
        return handle404(request)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from kivy.lang import Builder
from kivy.clock import Clock
from kivy.uix.modalview import ModalView
from kivy.properties import ListProperty
from kivy.metrics import sp, dp
from kivymd.theming import ThemableBehavior
from kivy.uix.floatlayout import FloatLayout
from kivymd.elevationbehavior import ElevationBehavior

Builder.load_string("""
<MDDatePicker>:
    size_hint: (None, None)
    size: dp(260), dp(30)+dp(150)+dp(290)
    pos_hint: {'center_x': .5, 'center_y': .5}
    canvas:
        Color:
            rgb: app.theme_cls.primary_dark
        Rectangle:
            size: dp(260), dp(30)
            pos: root.pos[0], root.pos[1] + root.height-dp(30)
        Color:
            rgb: app.theme_cls.primary_color
        Rectangle:
            size: dp(260), dp(150)
            pos: root.pos[0], root.pos[1] + root.height-(dp(30)+dp(150))
        Color:
            rgb: app.theme_cls.bg_normal
        Rectangle:
            size: dp(260), dp(290)
            pos: root.pos[0], root.pos[1] + root.height-(dp(30)+dp(150)+dp(290))

    MDFlatButton:
        pos: root.pos[0]+root.size[0]-dp(72)*2, root.pos[1] + dp(10)
        text: "Cancel"
        on_release: root.close_cancel()
    MDFlatButton:
        pos: root.pos[0]+root.size[0]-dp(72), root.pos[1] + dp(10)
        text: "OK"
        on_release: root.close_ok()

""")

class MDDatePicker(FloatLayout, ThemableBehavior, ElevationBehavior, ModalView):
    background_color = ListProperty([0, 0, 0, 0])

    def __init__(self, **kwargs):
        super(MDDatePicker, self).__init__(**kwargs)

    def close_cancel(self):
        self.dismiss()

    def close_ok(self):
        self.dismiss()

if __name__ == "__main__":
    from kivy.app import App
    from kivymd.theming import ThemeManager

    class DatePickerApp(App):
        theme_cls = ThemeManager()
        theme_cls.primary_palette = "Teal"
        # theme_cls.primary_hue = "200"

        def open_dialog(self):
            self.date_picker = MDDatePicker()
            # self.date_picker.bind(time=self.get_time)
            '''try:
                self.time_picker.set_time(self.last_time)
            except TypeError:
                pass'''
            self.date_picker.open()

        def build(self):
            main_widget = Builder.load_string("""
#:import MDRaisedButton kivymd.button.MDRaisedButton
FloatLayout:
    MDRaisedButton:
        size_hint: None, None
        size: 3 * dp(48), dp(48)
        # center_x: self.parent.center_x
        text: 'Switch theme color (debug)'
        on_release: app.theme_cls.primary_palette = 'Teal' if app.theme_cls.primary_palette != 'Teal' else 'DeepOrange'
        opposite_colors: True
    MDRaisedButton:
        size_hint: None, None
        size: 3 * dp(48), dp(48)
        center_x: self.parent.center_x
        text: 'Switch theme style'
        on_release: app.theme_cls.theme_style = 'Dark' if app.theme_cls.theme_style != 'Dark' else 'Light'
        opposite_colors: True
    MDRaisedButton:
        size_hint: None, None
        pos_hint: {'center_x': .5, 'center_y': .5}
        size: 3 * dp(48), dp(48)
        center_x: self.parent.center_x
        text: 'Open date picker'
        on_release: app.open_dialog()
        opposite_colors: True
""")
            return main_widget

    DatePickerApp().run()

<commit_msg>Added MDDatePicker<commit_after># -*- coding: utf-8 -*-

from kivy.lang import Builder
from kivy.clock import Clock
from kivy.uix.modalview import ModalView
from kivy.metrics import dp
from kivymd.label import MDLabel
from kivymd.theming import ThemableBehavior
from kivy.uix.floatlayout import FloatLayout
from kivy.uix.gridlayout import GridLayout
from kivymd.elevationbehavior import ElevationBehavior
import calendar
from datetime import date
import datetime
from kivy.properties import StringProperty, ListProperty, NumericProperty, OptionProperty, ObjectProperty
from kivy.uix.anchorlayout import AnchorLayout
from kivy.uix.behaviors import ButtonBehavior
from kivymd.ripplebehavior import RectangularRippleBehavior
from kivymd.backgroundcolorbehavior import BackgroundColorBehavior
from kivy.animation import Animation
from kivymd.color_definitions import colors
from kivy.utils import get_color_from_hex
from kivy.uix.widget import WidgetException

Builder.load_string("""
<MDDatePicker>:
    size_hint: (None, None)
    size: dp(300), dp(30)+dp(130)+dp(300)
    pos_hint: {'center_x': .5, 'center_y': .5}
    canvas:
        Color:
            rgb: app.theme_cls.primary_dark
        Rectangle:
            size: dp(300), dp(30)
            pos: root.pos[0], root.pos[1] + root.height-dp(30)
        Color:
            rgb: app.theme_cls.primary_color
        Rectangle:
            size: dp(300), dp(130)
            pos: root.pos[0], root.pos[1] + root.height-(dp(30)+dp(130))
        Color:
            rgb: app.theme_cls.bg_normal
        Rectangle:
            size: dp(300), dp(300)
            pos: root.pos[0], root.pos[1] + root.height-(dp(30)+dp(130)+dp(300))

    MDLabel:
        id: label_weekday
        font_style: "Headline"
        size_hint: (None, None)
        size: root.width, dp(30)
        pos: root.pos
        pos_hint: {'center_x': 0.5, 'center_y': 0.97}
        valign: "middle"
        halign: "center"

    MDLabel:
        id: label_date
        font_style: "Display3"
        size_hint: (None, None)
        size: root.width, dp(30)
        pos: root.pos
        pos_hint: {'center_x': 0.5, 'center_y': 0.85}
        valign: "middle"
        halign: "center"

    MDLabel:
        id: label_short_month
        font_style: "Headline"
        size_hint: (None, None)
        size: root.width, dp(30)
        pos: root.pos
        pos_hint: {'center_x': 0.5, 'center_y': 0.75}
        valign: "middle"
        halign: "center"

    MDLabel:
        id: label_year
        font_style: "Headline"
        size_hint: (None, None)
        size: root.width, dp(30)
        pos: root.pos
        pos_hint: {'center_x': 0.5, 'center_y': 0.7}
        valign: "middle"
        halign: "center"

    MDLabel:
        id: label_current_month
        font_style: "Body2"
        text: "September 2016"
        size_hint: (None, None)
        size: root.width, dp(30)
        pos: root.pos
        theme_text_color: 'Primary'
        pos_hint: {'center_x': 0.5, 'center_y': 0.6}
        valign: "middle"
        halign: "center"

    MDIconButton:
        icon: 'arrow-left'
        pos_hint: {'center_x': 0.275, 'center_y': 0.6}
        on_release: root.prev_month()

    MDIconButton:
        icon: 'arrow-right'
        pos_hint: {'center_x': 0.725, 'center_y': 0.6}
        on_release: root.next_month()

    MDFlatButton:
        pos: root.pos[0]+root.size[0]-dp(72)*2, root.pos[1] + dp(10)
        text: "Cancel"
        on_release: root.close_cancel()
    MDFlatButton:
        pos: root.pos[0]+root.size[0]-dp(72), root.pos[1] + dp(10)
        text: "OK"
        on_release: root.close_ok()
<CalendarButton>
    canvas:
        Color:
            #rgba: self.background_color if self.state == 'normal' else self._bg_color_down
            rgba: self._current_button_color
        Rectangle:
            size: self.size
            pos: self.pos
    size_hint: (None, None)
    height: dp(36)
    width: _label.texture_size[0] + dp(16)
    padding: (dp(8), 0)
    theme_text_color: 'Custom'
    text_color: root.theme_cls.primary_color
    MDLabel:
        id: _label
        text: root._text
        font_style: 'Button'
        size_hint_x: None
        text_size: (None, root.height)
        height: self.texture_size[1]
        theme_text_color: root.theme_text_color
        text_color: root.text_color
        valign: 'middle'
        halign: 'center'
        opposite_colors: root.opposite_colors

<CalendarSelector>
    canvas:
        Color:
            rgba: self.theme_color_with_alpha
        Ellipse:
            size: self.size
            pos: self.pos
""")

class CalendarButton(ThemableBehavior,
                     RectangularRippleBehavior,
                     ButtonBehavior,
                     BackgroundColorBehavior,
                     AnchorLayout):
    width = NumericProperty(dp(64),
                            min=dp(64),
                            max=None,
                            errorhandler=lambda x: dp(64))
    text_color = ListProperty()
    text = StringProperty('')
    theme_text_color = OptionProperty(None,
                                      allownone=True,
                                      options=['Primary', 'Secondary', 'Hint', 'Error', 'Custom'])
    text_color = ListProperty(None,
                              allownone=True)

    _text = StringProperty('')
    _bg_color_down = ListProperty([0, 0, 0, 0])
    _current_button_color = ListProperty([0, 0, 0, 0])

    def __init__(self, **kwargs):
        super(CalendarButton, self).__init__(**kwargs)
        self._current_button_color = self.background_color
        '''self._bg_color_down = get_color_from_hex(
            colors[self.theme_cls.theme_style]['FlatButtonDown'])'''

        Clock.schedule_once(lambda x: self.ids._label.bind(
            texture_size=self.update_width_on_label_texture))

    def update_width_on_label_texture(self, instance, value):
        self.ids._label.width = value[0]

    def on_text(self, instance, value):
        self._text = value.upper()

    def on_touch_down(self, touch):
        if touch.is_mouse_scrolling:
            return False
        elif not self.collide_point(touch.x, touch.y):
            return False
        elif self in touch.ud:
            return False
        elif self.disabled:
            return False
        else:
            self.fade_bg = Animation(duration=.2,
                                     _current_button_color=get_color_from_hex(
                                        colors[self.theme_cls.theme_style]['FlatButtonDown']))
            self.fade_bg.start(self)
            return super(CalendarButton, self).on_touch_down(touch)

    def on_touch_up(self, touch):
        if touch.grab_current is self:
            self.fade_bg.stop_property(self, '_current_button_color')
            Animation(duration=.05,
                      _current_button_color=self.background_color).start(self)
        return super(CalendarButton, self).on_touch_up(touch)

class DateButton(CalendarButton):
    def __init__(self, cls, **kwargs):
        super(CalendarButton, self).__init__(**kwargs)
        self.width = kwargs['size'][0]
        self.size = kwargs['size']
        self.text = kwargs['text']
        if self.text == str(date.today().day) \
                and cls.month == date.today().month \
                and cls.year == date.today().year:
            pass
        else:
            self.theme_text_color = 'Primary'

class CalendarSelector(CalendarButton):
    theme_color_with_alpha = ListProperty([0.0, 0.0, 0.0, 0.0])

    def __init__(self, cls):
        super(CalendarButton, self).__init__()
        self.selected_month = cls.month
        self.selected_year = cls.year
        self.selected_day = cls.day
        self.current_button = None
        self.disabled = True
        self.theme_color_with_alpha = self.theme_cls.primary_color
        self.theme_color_with_alpha[3] = 0.4

    def update(self, cls):
        if self.selected_month == cls.month and self.selected_year == cls.year:
            try:
                cls.add_widget(self)
                self.move(cls, self.current_button)
            except WidgetException:
                pass
        else:
            try:
                cls.remove_widget(self)
            except WidgetException:
                pass

    def move(self, cls, inst):
        if not inst:
            pass
        else:
            self.current_button = inst
            self.selected_month = cls.month
            self.selected_year = cls.year
            self.pos = inst.pos
            self.size = inst.size

    def receive_lookout(self, inst, cls):
        self.current_button = inst
        self.move(cls, self.current_button)
        self.update(cls)
        # cls.add_widget(self)

    def get_lookout(self):
        return self.selected_day

class MDDatePicker(FloatLayout,
                   ThemableBehavior,
                   ElevationBehavior,
                   ModalView):
    background_color = ListProperty([0, 0, 0, 0])
    date = ObjectProperty()

    def __init__(self, **kwargs):
        super(MDDatePicker, self).__init__(**kwargs)
        self.date = None
        self.layout = None
        self.selector = None

    def close_cancel(self):
        self.dismiss()

    def close_ok(self):
        self.date = datetime.datetime.strptime("".join([str(self.day),
                                                        str(self.month),
                                                        str(self.year)]), "%d%m%Y").date()
        self.dismiss()

    def set_date_str(self, day, month, year):
        self.day = day
        self.month = month
        self.year = year

    def set_date(self, the_date):
        try:
            self.day = the_date.day
            self.month = the_date.month
            self.year = the_date.year
        except AttributeError:
            raise TypeError("<DatePicker>.set_date requires a datetime.date, if you would prefer you can pass in <day>,"
                            " <month>, <year> separately as strings with <DatePicker>.set_date_str")

    def open(self, *args):
        super(MDDatePicker, self).open(*args)
        try:
            self.year
            self.month
            self.day
        except AttributeError:
            self.set_date_str(date.today().day, date.today().month, date.today().year)
        self.ids.label_weekday.text = str(datetime.date(self.year, self.month, self.day).strftime("%A"))
        self.ids.label_date.text = str(self.day)
        self.ids.label_short_month.text = calendar.month_name[self.month][:3].upper()
        self.ids.label_year.text = str(self.year)
        self.ids.label_current_month.text = calendar.month_name[self.month] + " " + str(self.year)
        self.selector = CalendarSelector(self)
        self.generate_calendar(year=self.year,
                               month=self.month,
                               lookout=self.selector.get_lookout())

    def get_touch(self, instance):
        self.selector.move(self, instance)
        self.selector.update(self)
        self.ids.label_weekday.text = str(datetime.date(self.year, self.month, int(instance.text)).strftime("%A"))
        self.ids.label_date.text = str(int(instance.text))
        self.ids.label_short_month.text = calendar.month_abbr[self.month].upper()
        self.ids.label_year.text = str(self.year)
        self.day = instance.text

    def next_month(self):
        self.month += 1
        if self.month == 13:
            self.month = 1
            self.year += 1
        self.selector.update(self)
        self.remove_widget(self.layout)
        del self.layout
        self.generate_calendar(year=self.year,
                               month=self.month)
        self.ids.label_current_month.text = calendar.month_name[self.month] + " " + str(self.year)

    def prev_month(self):
        self.month -= 1
        if self.month == 0:
            self.month = 12
            self.year -= 1
        self.selector.update(self)
        self.remove_widget(self.layout)
        del self.layout
        self.generate_calendar(year=self.year,
                               month=self.month)
        self.ids.label_current_month.text = calendar.month_name[self.month] + " " + str(self.year)

    def generate_calendar(self, year, month, lookout=None):
        actual_date = datetime.date(year, month, 1)
        cal = calendar.Calendar()
        self.layout = GridLayout(cols=7, size=(dp(250), dp(35*8)),
                                 size_hint=(None, None),
                                 pos=(dp(200), dp(50)),
                                 pos_hint={'center_x': .5, 'center_y': .275})
        for i in range(0, 2):
            self.layout.add_widget(MDLabel(size=(dp(35), dp(35)),
                                           size_hint=(None, None), text="S",
                                           halign="center",
                                           theme_text_color='Primary'))
        for i in cal.iterweekdays():
            if str(calendar.day_abbr[i][0]) == "S":
                pass
            else:
                self.layout.add_widget(MDLabel(size=(dp(35), dp(35)),
                                               size_hint=(None, None),
                                               text=calendar.day_abbr[i][0],
                                               halign="center",
                                               theme_text_color='Primary'))

        def first_day_of_month(d):
            return date(d.year, d.month, 1)

        month_start_col = first_day_of_month(actual_date).weekday()
        if month_start_col == 5:
            month_start_col = -2
        if month_start_col == 6:
            month_start_col = -1
        for i in range(-2, month_start_col):
            self.layout.add_widget(MDLabel(size=(dp(35), dp(35)),
                                           size_hint=(None, None)))
        for i in cal.itermonthdays2(year, month):
            if str(i[0]) == "0":
                pass
            else:
                date_button = DateButton(self, text=str(i[0]),
                                         size=(dp(35), dp(35)),
                                         size_hint=(None, None),
                                         pos=self.pos)
                date_button.bind(on_press=self.get_touch)
                self.layout.add_widget(date_button)
                if lookout:
                    if str(lookout) == str(i[0]):
                        self.selector.receive_lookout(date_button, self)
                        self.selector.move(cls=self, inst=date_button)
                        self.actual_start_lookout = date_button
                        Clock.schedule_once(lambda x: self.selector.move(cls=self,
                                                                         inst=self.actual_start_lookout),
                                            0.0000001)
        self.add_widget(self.layout)

if __name__ == "__main__":
    from kivy.app import App
    from kivymd.theming import ThemeManager

    class DatePickerApp(App):
        theme_cls = ThemeManager()

        def get_date(self, instance, the_date):
            print(the_date)
            self.last_date = the_date

        def open_dialog(self):
            self.date_picker = MDDatePicker()
            self.day = 19
            self.month = 9
            self.year = 2020
            # self.date_picker.set_date(datetime.datetime.strptime("".join([str(self.day),
            #                                                              str(self.month),
            #                                                              str(self.year)]), "%d%m%Y").date())'''
            # self.date_picker.set_date_str(self.day, self.month, self.year)
            self.date_picker.bind(date=self.get_date)
            try:
                self.date_picker.set_date(self.last_date)
            except AttributeError:
                pass
            self.date_picker.open()

        def build(self):
            main_widget = Builder.load_string("""
#:import MDRaisedButton kivymd.button.MDRaisedButton
#:import MDThemePicker kivymd.theme_picker.MDThemePicker
FloatLayout:
    MDRaisedButton:
        size_hint: None, None
        size: 3 * dp(48), dp(48)
        # center_x: self.parent.center_x
        text: 'Switch theme color (debug)'
        on_release: MDThemePicker().open()
        opposite_colors: True
    MDRaisedButton:
        size_hint: None, None
        size: 3 * dp(48), dp(48)
        center_x: self.parent.center_x
        text: 'Switch theme style'
        on_release: app.theme_cls.theme_style = 'Dark' if app.theme_cls.theme_style != 'Dark' else 'Light'
        opposite_colors: True
    MDRaisedButton:
        size_hint: None, None
        pos_hint: {'center_x': .5, 'center_y': .5}
        size: 3 * dp(48), dp(48)
        center_x: self.parent.center_x
        text: 'Open date picker'
        on_release: app.open_dialog()
        opposite_colors: True
""")
            return main_widget

    DatePickerApp().run()
<|endoftext|>


<commit_before>#!/usr/bin/env python3

from ..core import Module, Blob
from ..dataclasses import Table
from .hdf5 import HDF5Header

import km3io
import numpy as np
from collections import defaultdict

USR_MC_TRACKS_KEYS = [b'energy_lost_in_can', b'bx', b'by', b'ichan', b'cc']

class EventPump(Module):
    def configure(self):
        self._filename = self.require("filename")
        self._split_hits = self.get("split_hits", default=True)
        self.skip_hits = self.get("skip_hits", default=False)
        self.skip_mc_hits = self.get("skip_mc_hits", default=False)
        self.skip_mc_tracks = self.get("skip_mc_tracks", default=False)
        self.skip_header = self.get("skip_header", default=False)

        self._reader = km3io.OfflineReader(self._filename)

        self.header = None
        self.raw_header = None
        if not self.skip_header:
            if self._reader.header is not None:
                self.header = HDF5Header.from_km3io(self._reader.header)
                self.raw_header = self._generate_raw_header()

        self.blobs = self._blob_generator()

    def process(self, blob=None):
        return next(self.blobs)

    def __iter__(self):
        return self

    def __next__(self):
        return next(self.blobs)

    def __getitem__(self, item):
        if not isinstance(item, int):
            raise TypeError("Only integer indices are supported.")
        return self._parse_event(self._reader.events[item], Blob())

    def get_number_of_blobs(self):
        return len(self._reader.events)

    def _blob_generator(self):
        for event in self._reader.events:
            blob = Blob()

            self._parse_event(event, blob)
            yield blob

    def _parse_event(self, event, blob):
        blob["RawHeader"] = self.raw_header
        blob["Header"] = self.header

        blob['Event'] = event
        if not self.skip_hits and event.n_hits > 0:
            blob['Hits'] = self._parse_hits(event.hits)
        if not self.skip_mc_hits and event.n_mc_hits > 0:
            blob['McHits'] = self._parse_mc_hits(event.mc_hits)
        if not self.skip_mc_tracks and event.n_mc_tracks > 0:
            blob['McTracks'] = self._parse_mc_tracks(event.mc_tracks)
        return blob

    def _generate_raw_header(self):
        header = {}
        raw_header = {"field_names": [], "parameter": [], "field_values": []}
        for n, x in self._reader._fobj['Head'
                                       ]._map_3c_string_2c_string_3e_.items():
            header[n.decode("utf-8")] = x.decode("utf-8").strip()
        for attribute, fields in km3io.definitions.mc_header.items():
            values = header.get(attribute, '').split()
            if not values:
                continue
            raw_header["field_values"].append(" ".join(map(str, values)))
            raw_header["field_names"].append(" ".join(fields))
            raw_header["parameter"].append(attribute)
        return Table(
            raw_header,
            h5loc="/raw_header",
            name="RawHeader",
            h5singleton=True
        )

    def _parse_eventinfo(self, event):
        wgt1, wgt2, wgt3, wgt4 = self._parse_wgts(event.w)
        tab_data = {
            'event_id': event.id,
            'run_id': event.run_id,
            'weight_w1': wgt1,
            'weight_w2': wgt2,
            'weight_w3': wgt3,
            'weight_w4': wgt4,
            'timestamp': event.t_sec,
            'nanoseconds': event.t_ns,
            'mc_time': event.mc_t,
            'trigger_mask': event.trigger_mask,
            'trigger_counter': event.trigger_counter,
            'overlays': event.overlays,
            'det_id': event.det_id,
            'frame_index': event.frame_index,
            'mc_run_id': event.mc_run_id,
        }
        info = Table(tab_data, h5loc='/event_info', name='EventInfo')
        return info

    @staticmethod
    def _parse_wgts(wgt):
        if len(wgt) == 3:
            wgt1, wgt2, wgt3 = wgt
            wgt4 = np.nan
        elif len(wgt) == 4:
            # what the hell is w4?
            wgt1, wgt2, wgt3, wgt4 = wgt
        else:
            wgt1 = wgt2 = wgt3 = wgt4 = np.nan
        return wgt1, wgt2, wgt3, wgt4

    def _parse_usr_to_dct(self, mc_tracks):
        dct = defaultdict(list)
        for k in USR_MC_TRACKS_KEYS:
            dec_key = k.decode('utf_8')
            for i in range(mc_tracks.usr_names.shape[0]):
                value = np.nan
                if k in mc_tracks.usr_names[i]:
                    mask = mc_tracks.usr_names[i] == k
                    value = mc_tracks.usr[i][mask][0]
                dct[dec_key].append(value)
        return dct

    def _parse_mc_tracks(self, mc_tracks):
        dct = {
            'dir_x': mc_tracks.dir_x,
            'dir_y': mc_tracks.dir_y,
            'dir_z': mc_tracks.dir_z,
            'pos_x': mc_tracks.pos_x,
            'pos_y': mc_tracks.pos_y,
            'pos_z': mc_tracks.pos_z,
            'energy': mc_tracks.E,
            'time': mc_tracks.t,
            'type': mc_tracks.type,
            'id': mc_tracks.id,
            'length': mc_tracks.len
        }
        dct.update(self._parse_usr_to_dct(mc_tracks))
        return Table(dct, name='McTracks', h5loc='/mc_tracks', split_h5=True)

    def _parse_mc_hits(self, mc_hits):
        return Table({
            "a": mc_hits.a,
            "origin": mc_hits.origin,
            "pmt_id": mc_hits.pmt_id,
            "time": mc_hits.t
        },
                     h5loc="/mc_hits",
                     split_h5=self._split_hits,
                     name='McHits')

    def _parse_hits(self, hits):
        return Table({
            "channel_id": hits.channel_id,
            "dom_id": hits.dom_id,
            "time": hits.t,
            "tot": hits.tot,
            "triggered": hits.trig
        },
                     h5loc="/hits",
                     split_h5=self._split_hits,
                     name='Hits')
<commit_msg>Introduced a flag for reading usr data<commit_after>#!/usr/bin/env python3

from ..core import Module, Blob
from ..dataclasses import Table
from .hdf5 import HDF5Header

import km3io
import numpy as np
from collections import defaultdict

USR_MC_TRACKS_KEYS = [b'energy_lost_in_can', b'bx', b'by', b'ichan', b'cc']

class EventPump(Module):
    def configure(self):
        self._filename = self.require("filename")
        self._split_hits = self.get("split_hits", default=True)
        self.skip_hits = self.get("skip_hits", default=False)
        self.skip_mc_hits = self.get("skip_mc_hits", default=False)
        self.skip_mc_tracks = self.get("skip_mc_tracks", default=False)
        self.skip_header = self.get("skip_header", default=False)

        self._read_usr_data = self.get("read_usr_data", default=False)
        if self._read_usr_data:
            self.log.warning(
                "Reading usr-data will massively decrease the performance."
            )

        self._reader = km3io.OfflineReader(self._filename)

        self.header = None
        self.raw_header = None
        if not self.skip_header:
            if self._reader.header is not None:
                self.header = HDF5Header.from_km3io(self._reader.header)
                self.raw_header = self._generate_raw_header()

        self.blobs = self._blob_generator()

    def process(self, blob=None):
        return next(self.blobs)

    def __iter__(self):
        return self

    def __next__(self):
        return next(self.blobs)

    def __getitem__(self, item):
        if not isinstance(item, int):
            raise TypeError("Only integer indices are supported.")
        return self._parse_event(self._reader.events[item], Blob())

    def get_number_of_blobs(self):
        return len(self._reader.events)

    def _blob_generator(self):
        for event in self._reader.events:
            blob = Blob()

            self._parse_event(event, blob)
            yield blob

    def _parse_event(self, event, blob):
        blob["RawHeader"] = self.raw_header
        blob["Header"] = self.header

        blob['Event'] = event
        if not self.skip_hits and event.n_hits > 0:
            blob['Hits'] = self._parse_hits(event.hits)
        if not self.skip_mc_hits and event.n_mc_hits > 0:
            blob['McHits'] = self._parse_mc_hits(event.mc_hits)
        if not self.skip_mc_tracks and event.n_mc_tracks > 0:
            blob['McTracks'] = self._parse_mc_tracks(event.mc_tracks)
        return blob

    def _generate_raw_header(self):
        header = {}
        raw_header = {"field_names": [], "parameter": [], "field_values": []}
        for n, x in self._reader._fobj['Head'
                                       ]._map_3c_string_2c_string_3e_.items():
            header[n.decode("utf-8")] = x.decode("utf-8").strip()
        for attribute, fields in km3io.definitions.mc_header.items():
            values = header.get(attribute, '').split()
            if not values:
                continue
            raw_header["field_values"].append(" ".join(map(str, values)))
            raw_header["field_names"].append(" ".join(fields))
            raw_header["parameter"].append(attribute)
        return Table(
            raw_header,
            h5loc="/raw_header",
            name="RawHeader",
            h5singleton=True
        )

    def _parse_eventinfo(self, event):
        wgt1, wgt2, wgt3, wgt4 = self._parse_wgts(event.w)
        tab_data = {
            'event_id': event.id,
            'run_id': event.run_id,
            'weight_w1': wgt1,
            'weight_w2': wgt2,
            'weight_w3': wgt3,
            'weight_w4': wgt4,
            'timestamp': event.t_sec,
            'nanoseconds': event.t_ns,
            'mc_time': event.mc_t,
            'trigger_mask': event.trigger_mask,
            'trigger_counter': event.trigger_counter,
            'overlays': event.overlays,
            'det_id': event.det_id,
            'frame_index': event.frame_index,
            'mc_run_id': event.mc_run_id,
        }
        info = Table(tab_data, h5loc='/event_info', name='EventInfo')
        return info

    @staticmethod
    def _parse_wgts(wgt):
        if len(wgt) == 3:
            wgt1, wgt2, wgt3 = wgt
            wgt4 = np.nan
        elif len(wgt) == 4:
            # what the hell is w4?
            wgt1, wgt2, wgt3, wgt4 = wgt
        else:
            wgt1 = wgt2 = wgt3 = wgt4 = np.nan
        return wgt1, wgt2, wgt3, wgt4

    def _parse_usr_to_dct(self, mc_tracks):
        dct = defaultdict(list)
        for k in USR_MC_TRACKS_KEYS:
            dec_key = k.decode('utf_8')
            for i in range(mc_tracks.usr_names.shape[0]):
                value = np.nan
                if k in mc_tracks.usr_names[i]:
                    mask = mc_tracks.usr_names[i] == k
                    value = mc_tracks.usr[i][mask][0]
                dct[dec_key].append(value)
        return dct

    def _parse_mc_tracks(self, mc_tracks):
        dct = {
            'dir_x': mc_tracks.dir_x,
            'dir_y': mc_tracks.dir_y,
            'dir_z': mc_tracks.dir_z,
            'pos_x': mc_tracks.pos_x,
            'pos_y': mc_tracks.pos_y,
            'pos_z': mc_tracks.pos_z,
            'energy': mc_tracks.E,
            'time': mc_tracks.t,
            'type': mc_tracks.type,
            'id': mc_tracks.id,
            'length': mc_tracks.len
        }
        if self._read_usr_data:
            dct.update(self._parse_usr_to_dct(mc_tracks))
        return Table(dct, name='McTracks', h5loc='/mc_tracks', split_h5=True)

    def _parse_mc_hits(self, mc_hits):
        return Table({
            "a": mc_hits.a,
            "origin": mc_hits.origin,
            "pmt_id": mc_hits.pmt_id,
            "time": mc_hits.t
        },
                     h5loc="/mc_hits",
                     split_h5=self._split_hits,
                     name='McHits')

    def _parse_hits(self, hits):
        return Table({
            "channel_id": hits.channel_id,
            "dom_id": hits.dom_id,
            "time": hits.t,
            "tot": hits.tot,
            "triggered": hits.trig
        },
                     h5loc="/hits",
                     split_h5=self._split_hits,
                     name='Hits')
<|endoftext|>


<commit_before>VERSION = (0, 0, 9)
<commit_msg>bump 0.1.0<commit_after>VERSION = (0, 1, 0)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf8 -*-

import sys
import os
import numpy as np
import tensorflow as tf
import logging
import argparse

from gym.spaces import Discrete

from Learner import Learner
from utils import discount_rewards, save_config
from CartPole import make_predef_CartPole_envs, make_random_CartPole_envs
from Reporter import Reporter
from gradient_ops import create_accumulative_gradients_op, add_accumulative_gradients_op, reset_accumulative_gradients_op

class TaskLearner(Learner):
    """Learner for a specific environment and with its own action selection."""
    def __init__(self, env, action, master, **usercfg):
        super(TaskLearner, self).__init__(env, **usercfg)
        self.action = action
        self.master = master

    def choose_action(self, state):
        """Choose an action."""
        return self.master.session.run([self.action], feed_dict={self.master.states: [state]})[0]

class KnowledgeTransferLearner(Learner):
    """Learner for variations of a task."""
    def __init__(self, envs, monitor_dir, **usercfg):
        super(KnowledgeTransferLearner, self).__init__(envs[0], **usercfg)
        self.envs = envs
        self.n_tasks = len(envs)
        self.monitor_dir = monitor_dir
        self.nA = envs[0].action_space.n
        self.config.update(dict(
            timesteps_per_batch=2000,
            trajectories_per_batch=10,
            batch_update="timesteps",
            n_iter=400,
            gamma=0.99,
            learning_rate=0.005,
            n_hidden_units=20,
            repeat_n_actions=1,
            n_sparse_units=10
        ))
        self.config.update(usercfg)
        self.build_networks()
        self.task_learners = [TaskLearner(envs[i], action, self, **self.config) for i, action in enumerate(self.action_tensors)]
        if self.config["save_model"]:
            for action_tensor in self.action_tensors:
                tf.add_to_collection("action", action_tensor)
            tf.add_to_collection("states", self.states)
            self.saver = tf.train.Saver()

    def build_networks(self):
        self.session = tf.Session()

        self.states = tf.placeholder(tf.float32, name="states")
        self.action_taken = tf.placeholder(tf.float32, name="action_taken")
        self.advantage = tf.placeholder(tf.float32, name="advantage")

        W0 = tf.Variable(tf.random_normal([self.nO, self.config["n_hidden_units"]]) / np.sqrt(self.nO), name='W0')
        b0 = tf.Variable(tf.zeros([self.config["n_hidden_units"]]), name='b0')
        # Action probabilities
        L1 = tf.tanh(tf.matmul(self.states, W0) + b0[None, :])

        knowledge_base = tf.Variable(tf.random_normal([self.config["n_hidden_units"], self.config["n_sparse_units"]]))

        self.shared_vars = [W0, b0, knowledge_base]

        # Every task has its own sparse representation
        sparse_representations = [tf.Variable(tf.random_normal([self.config["n_sparse_units"], self.nA])) for _ in range(self.n_tasks)]

        self.probs_tensors = [tf.nn.softmax(tf.matmul(L1, tf.matmul(knowledge_base, s))) for s in sparse_representations]

        self.action_tensors = [tf.squeeze(tf.multinomial(tf.log(probs), 1)) for probs in self.probs_tensors]

        self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.config["learning_rate"], decay=0.9, epsilon=1e-9)
        net_vars = self.shared_vars + sparse_representations
        self.accum_grads = create_accumulative_gradients_op(net_vars, 1)

        # self.writers = []
        self.losses = []
        for i, probabilities in enumerate(self.probs_tensors):
            good_probabilities = tf.reduce_sum(tf.multiply(probabilities, tf.one_hot(tf.cast(self.action_taken, tf.int32), self.nA)), reduction_indices=[1])
            eligibility = tf.log(good_probabilities) * self.advantage
            # eligibility = tf.Print(eligibility, [eligibility], first_n=5)
            loss = -tf.reduce_sum(eligibility)
            self.losses.append(loss)
            # writer = tf.summary.FileWriter(self.monitor_dir + "/task" + str(i), self.session.graph)

        # An add op for every task & its loss
        # add_accumulative_gradients_op(net_vars, accum_grads, loss, identifier)
        self.add_accum_grads = [add_accumulative_gradients_op(
            self.shared_vars + [sparse_representations[i]],
            self.accum_grads,
            loss,
            i)
            for i, loss in enumerate(self.losses)]

        self.apply_gradients = self.optimizer.apply_gradients(
            zip(self.accum_grads, net_vars))
        self.reset_accum_grads = reset_accumulative_gradients_op(net_vars, self.accum_grads, 1)

        init = tf.global_variables_initializer()

        # Launch the graph.
        self.session.run(init)

    def learn(self):
        """Run learning algorithm"""
        reporter = Reporter()
        config = self.config
        total_n_trajectories = np.zeros(len(self.envs))
        for iteration in range(config["n_iter"]):
            self.session.run([self.reset_accum_grads])
            for i, learner in enumerate(self.task_learners):
                # Collect trajectories until we get timesteps_per_batch total timesteps
                trajectories = learner.get_trajectories()
                total_n_trajectories[i] += len(trajectories)
                all_state = np.concatenate([trajectory["state"] for trajectory in trajectories])
                # Compute discounted sums of rewards
                rets = [discount_rewards(trajectory["reward"], config["gamma"]) for trajectory in trajectories]
                max_len = max(len(ret) for ret in rets)
                padded_rets = [np.concatenate([ret, np.zeros(max_len - len(ret))]) for ret in rets]
                # Compute time-dependent baseline
                baseline = np.mean(padded_rets, axis=0)
                # Compute advantage function
                advs = [ret - baseline[:len(ret)] for ret in rets]
                all_action = np.concatenate([trajectory["action"] for trajectory in trajectories])
                all_adv = np.concatenate(advs)
                # Do policy gradient update step
                episode_rewards = np.array([trajectory["reward"].sum() for trajectory in trajectories])  # episode total rewards
                episode_lengths = np.array([len(trajectory["reward"]) for trajectory in trajectories])  # episode lengths
                self.session.run([self.add_accum_grads[i]], feed_dict={
                    self.states: all_state,
                    self.action_taken: all_action,
                    self.advantage: all_adv
                })
                # summary = self.session.run([self.master.summary_op], feed_dict={
                #     self.reward: reward
                #     # self.master.episode_length: trajectory["steps"]
                # })

                # self.writer.add_summary(summary[0], iteration)
                # self.writer.flush()
                print("Task:", i)
                reporter.print_iteration_stats(iteration, episode_rewards, episode_lengths, total_n_trajectories[i])

            # Apply accumulated gradient after all the gradients of each task are summed
            self.session.run([self.apply_gradients])

            # self.writer.add_summary(result[0], iteration)
            # self.writer.flush()

        if self.config["save_model"]:
            if not os.path.exists(self.monitor_dir):
                os.makedirs(self.monitor_dir)
            self.saver.save(self.session, os.path.join(self.monitor_dir, "model"))

parser = argparse.ArgumentParser()
parser.add_argument("environment", metavar="env", type=str, help="Gym environment to execute the experiment on.")
parser.add_argument("monitor_path", metavar="monitor_path", type=str, help="Path where Gym monitor files may be saved")
parser.add_argument("--iterations", default=100, type=int, help="Number of iterations to run the algorithm.")
parser.add_argument("--save_model", action="store_true", default=False, help="Save resulting model.")
parser.add_argument("--random_envs", type=int, help="Number of environments with random parameters to generate.")

def main():
    try:
        args = parser.parse_args()
    except:
        sys.exit()
    if not os.path.exists(args.monitor_path):
        os.makedirs(args.monitor_path)
    if args.environment != "CartPole-v0":
        raise NotImplementedError("Only the environment \"CartPole-v0\" is supported right now.")
    envs = make_random_CartPole_envs(args.random_envs) if args.random_envs else make_predef_CartPole_envs()
    if isinstance(envs[0].action_space, Discrete):
        agent = KnowledgeTransferLearner(envs, args.monitor_path, n_iter=args.iterations, save_model=args.save_model)
    else:
        raise NotImplementedError("Only environments with a discrete action space are supported right now.")
    save_config(args.monitor_path, agent.config, envs)
    try:
        agent.learn()
    except KeyboardInterrupt:
        pass

if __name__ == "__main__":
    main()
<commit_msg>Knowledge transfer learner now generates summaries<commit_after>#!/usr/bin/env python
# -*- coding: utf8 -*-

import sys
import os
import numpy as np
import tensorflow as tf
import logging
import argparse

from gym.spaces import Discrete

from Learner import Learner
from utils import discount_rewards, save_config
from CartPole import make_predef_CartPole_envs, make_random_CartPole_envs
from Reporter import Reporter
from gradient_ops import create_accumulative_gradients_op, add_accumulative_gradients_op, reset_accumulative_gradients_op

class TaskLearner(Learner):
    """Learner for a specific environment and with its own action selection."""
    def __init__(self, env, action, master, **usercfg):
        super(TaskLearner, self).__init__(env, **usercfg)
        self.action = action
        self.master = master

    def choose_action(self, state):
        """Choose an action."""
        return self.master.session.run([self.action], feed_dict={self.master.states: [state]})[0]

class KnowledgeTransferLearner(Learner):
    """Learner for variations of a task."""
    def __init__(self, envs, monitor_dir, **usercfg):
        super(KnowledgeTransferLearner, self).__init__(envs[0], **usercfg)
        self.envs = envs
        self.n_tasks = len(envs)
        self.monitor_dir = monitor_dir
        self.nA = envs[0].action_space.n
        self.config.update(dict(
            timesteps_per_batch=2000,
            trajectories_per_batch=10,
            batch_update="timesteps",
            n_iter=400,
            gamma=0.99,
            learning_rate=0.005,
            n_hidden_units=20,
            repeat_n_actions=1,
            n_sparse_units=10
        ))
        self.config.update(usercfg)
        self.build_networks()
        self.task_learners = [TaskLearner(envs[i], action, self, **self.config) for i, action in enumerate(self.action_tensors)]
        if self.config["save_model"]:
            for action_tensor in self.action_tensors:
                tf.add_to_collection("action", action_tensor)
            tf.add_to_collection("states", self.states)
            self.saver = tf.train.Saver()

    def build_networks(self):
        self.session = tf.Session()

        self.states = tf.placeholder(tf.float32, name="states")
        self.action_taken = tf.placeholder(tf.float32, name="action_taken")
        self.advantage = tf.placeholder(tf.float32, name="advantage")

        W0 = tf.Variable(tf.random_normal([self.nO, self.config["n_hidden_units"]]) / np.sqrt(self.nO), name='W0')
        b0 = tf.Variable(tf.zeros([self.config["n_hidden_units"]]), name='b0')
        # Action probabilities
        L1 = tf.tanh(tf.matmul(self.states, W0) + b0[None, :])

        knowledge_base = tf.Variable(tf.random_normal([self.config["n_hidden_units"], self.config["n_sparse_units"]]))

        self.shared_vars = [W0, b0, knowledge_base]

        # Every task has its own sparse representation
        sparse_representations = [tf.Variable(tf.random_normal([self.config["n_sparse_units"], self.nA])) for _ in range(self.n_tasks)]

        self.probs_tensors = [tf.nn.softmax(tf.matmul(L1, tf.matmul(knowledge_base, s))) for s in sparse_representations]

        self.action_tensors = [tf.squeeze(tf.multinomial(tf.log(probs), 1)) for probs in self.probs_tensors]

        self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.config["learning_rate"], decay=0.9, epsilon=1e-9)
        net_vars = self.shared_vars + sparse_representations
        self.accum_grads = create_accumulative_gradients_op(net_vars, 1)

        self.writers = []
        self.losses = []

        self.loss = tf.placeholder("float", name="loss")
        self.rewards = tf.placeholder("float", name="Rewards")
        self.episode_lengths = tf.placeholder("float", name="Episode_lengths")
        summary_loss = tf.summary.scalar("Loss", self.loss)
        summary_rewards = tf.summary.scalar("Rewards", self.rewards)
        summary_episode_lengths = tf.summary.scalar("Episode_lengths", self.episode_lengths)
        self.summary_op = tf.summary.merge([summary_loss, summary_rewards, summary_episode_lengths])

        for i, probabilities in enumerate(self.probs_tensors):
            good_probabilities = tf.reduce_sum(tf.multiply(probabilities, tf.one_hot(tf.cast(self.action_taken, tf.int32), self.nA)), reduction_indices=[1])
            eligibility = tf.log(good_probabilities) * self.advantage
            # eligibility = tf.Print(eligibility, [eligibility], first_n=5)
            loss = -tf.reduce_sum(eligibility)
            self.losses.append(loss)
            writer = tf.summary.FileWriter(self.monitor_dir + "/task" + str(i), self.session.graph)
            self.writers.append(writer)

        # An add op for every task & its loss
        # add_accumulative_gradients_op(net_vars, accum_grads, loss, identifier)
        self.add_accum_grads = [add_accumulative_gradients_op(
            self.shared_vars + [sparse_representations[i]],
            self.accum_grads,
            loss,
            i)
            for i, loss in enumerate(self.losses)]

        self.apply_gradients = self.optimizer.apply_gradients(
            zip(self.accum_grads, net_vars))
        self.reset_accum_grads = reset_accumulative_gradients_op(net_vars, self.accum_grads, 1)

        init = tf.global_variables_initializer()

        # Launch the graph.
        self.session.run(init)

    def learn(self):
        """Run learning algorithm"""
        reporter = Reporter()
        config = self.config
        total_n_trajectories = np.zeros(len(self.envs))
        for iteration in range(config["n_iter"]):
            self.session.run([self.reset_accum_grads])
            for i, learner in enumerate(self.task_learners):
                # Collect trajectories until we get timesteps_per_batch total timesteps
                trajectories = learner.get_trajectories()
                total_n_trajectories[i] += len(trajectories)
                all_state = np.concatenate([trajectory["state"] for trajectory in trajectories])
                # Compute discounted sums of rewards
                rets = [discount_rewards(trajectory["reward"], config["gamma"]) for trajectory in trajectories]
                max_len = max(len(ret) for ret in rets)
                padded_rets = [np.concatenate([ret, np.zeros(max_len - len(ret))]) for ret in rets]
                # Compute time-dependent baseline
                baseline = np.mean(padded_rets, axis=0)
                # Compute advantage function
                advs = [ret - baseline[:len(ret)] for ret in rets]
                all_action = np.concatenate([trajectory["action"] for trajectory in trajectories])
                all_adv = np.concatenate(advs)
                # Do policy gradient update step
                episode_rewards = np.array([trajectory["reward"].sum() for trajectory in trajectories])  # episode total rewards
                episode_lengths = np.array([len(trajectory["reward"]) for trajectory in trajectories])  # episode lengths
                results = self.session.run([self.losses[i], self.add_accum_grads[i]], feed_dict={
                    self.states: all_state,
                    self.action_taken: all_action,
                    self.advantage: all_adv
                })
                summary = self.session.run([self.summary_op], feed_dict={
                    self.loss: results[0],
                    self.rewards: np.mean(episode_rewards),
                    self.episode_lengths: np.mean(episode_lengths)
                })

                self.writers[i].add_summary(summary[0], iteration)
                self.writers[i].flush()
                print("Task:", i)
                reporter.print_iteration_stats(iteration, episode_rewards, episode_lengths, total_n_trajectories[i])

            # Apply accumulated gradient after all the gradients of each task are summed
            self.session.run([self.apply_gradients])

        if self.config["save_model"]:
            if not os.path.exists(self.monitor_dir):
                os.makedirs(self.monitor_dir)
            self.saver.save(self.session, os.path.join(self.monitor_dir, "model"))

parser = argparse.ArgumentParser()
parser.add_argument("environment", metavar="env", type=str, help="Gym environment to execute the experiment on.")
parser.add_argument("monitor_path", metavar="monitor_path", type=str, help="Path where Gym monitor files may be saved")
parser.add_argument("--iterations", default=100, type=int, help="Number of iterations to run the algorithm.")
parser.add_argument("--save_model", action="store_true", default=False, help="Save resulting model.")
parser.add_argument("--random_envs", type=int, help="Number of environments with random parameters to generate.")

def main():
    try:
        args = parser.parse_args()
    except:
        sys.exit()
    if not os.path.exists(args.monitor_path):
        os.makedirs(args.monitor_path)
    if args.environment != "CartPole-v0":
        raise NotImplementedError("Only the environment \"CartPole-v0\" is supported right now.")
    envs = make_random_CartPole_envs(args.random_envs) if args.random_envs else make_predef_CartPole_envs()
    if isinstance(envs[0].action_space, Discrete):
        agent = KnowledgeTransferLearner(envs, args.monitor_path, n_iter=args.iterations, save_model=args.save_model)
    else:
        raise NotImplementedError("Only environments with a discrete action space are supported right now.")
    save_config(args.monitor_path, agent.config, envs)
    try:
        agent.learn()
    except KeyboardInterrupt:
        pass

if __name__ == "__main__":
    main()
<|endoftext|>