"""Memory profiler module."""
import builtins
import gc
import inspect
import os
import operator
import psutil
import re
import runpy
import sys
import time

from collections import deque
from collections import Counter
from vprof import base_profiler

_BYTES_IN_MB = 1024 * 1024

def _remove_duplicates(objects):
    """Removes duplicate objects.

    http://www.peterbe.com/plog/uniqifiers-benchmark.
    """
    seen, uniq = set(), []
    for obj in objects:
        obj_id = id(obj)
        if obj_id in seen:
            continue
        seen.add(obj_id)
        uniq.append(obj)
    return uniq

def _get_in_memory_objects():
    """Returns all objects in memory."""
    gc.collect()
    return gc.get_objects()

def _process_in_memory_objects(objects):
    """Processes objects tracked by GC.

    Processing is done in separate function to avoid generating overhead.
    """
    return _remove_duplicates(obj for obj in objects
                              if not inspect.isframe(obj))

def _get_object_count_by_type(objects):
    """Counts Python objects by type."""
    return Counter(map(type, objects))

def _get_obj_count_difference(objs1, objs2):
    """Returns count difference in two collections of Python objects."""
    clean_obj_list1 = _process_in_memory_objects(objs1)
    clean_obj_list2 = _process_in_memory_objects(objs2)
    obj_count_1 = _get_object_count_by_type(clean_obj_list1)
    obj_count_2 = _get_object_count_by_type(clean_obj_list2)
    return obj_count_1 - obj_count_2

def _format_obj_count(objects):
    """Formats object count."""
    result = []
    regex = re.compile(r'<(?P<type>\w+) \'(?P<name>\S+)\'>')
    for obj_type, obj_count in objects.items():
        if obj_count != 0:
            match = re.findall(regex, repr(obj_type))
            if match:
                obj_type, obj_name = match[0]
                result.append(("%s %s" % (obj_type, obj_name), obj_count))
    return sorted(result, key=operator.itemgetter(1), reverse=True)

class _CodeEventsTracker:
    """Tracks specified events during code execution.

    Contains all logic related to measuring memory usage.
    """

    def __init__(self, target_modules):
        self._events_list = deque()
        self._original_trace_function = sys.gettrace()
        self._process = psutil.Process(os.getpid())
        self._resulting_events = []
        self.mem_overhead = None
        self.target_modules = target_modules

    def __enter__(self):
        """Enables events tracker."""
        sys.settrace(self._trace_memory_usage)
        return self

    def __exit__(self, exc_type, exc_val, exc_tbf):
        """Disables events tracker."""
        sys.settrace(self._original_trace_function)

    def _trace_memory_usage(self, frame, event, arg):  #pylint: disable=unused-argument
        """Checks memory usage when 'line' event occur."""
        if event == 'line' and frame.f_code.co_filename in self.target_modules:
            self._events_list.append(
                (frame.f_lineno, self._process.memory_info().rss,
                 frame.f_code.co_name, frame.f_code.co_filename))
        return self._trace_memory_usage

    @property
    def code_events(self):
        """Returns processed memory usage."""
        if self._resulting_events:
            return self._resulting_events
        for i, (lineno, mem, func, fname) in enumerate(self._events_list):
            mem_in_mb = float(mem - self.mem_overhead) / _BYTES_IN_MB
            if (self._resulting_events and
                    self._resulting_events[-1][0] == lineno and
                    self._resulting_events[-1][2] == func and
                    self._resulting_events[-1][3] == fname and
                    self._resulting_events[-1][1] < mem_in_mb):
                self._resulting_events[-1][1] = mem_in_mb
            else:
                self._resulting_events.append(
                    [i + 1, lineno, mem_in_mb, func, fname])
        return self._resulting_events

    @property
    def obj_overhead(self):
        """Returns all objects that are considered a profiler overhead.
        Objects are hardcoded for convenience.
        """
        overhead = [
            self,
            self._resulting_events,
            self._events_list,
            self._process
        ]
        overhead_count = _get_object_count_by_type(overhead)
        # One for reference to __dict__ and one for reference to
        # the current module.
        overhead_count[dict] += 2
        return overhead_count

    def compute_mem_overhead(self):
        """Returns memory overhead."""
        self.mem_overhead = (self._process.memory_info().rss -
                             builtins.initial_rss_size)

class MemoryProfiler(base_profiler.BaseProfiler):
    """Memory profiler wrapper.

    Runs memory profiler and processes collected stats.
    """

    def profile_package(self):
        """Returns memory stats for a package."""
        target_modules = base_profiler.get_pkg_module_names(self._run_object)
        try:
            with _CodeEventsTracker(target_modules) as prof:
                prof.compute_mem_overhead()
                runpy.run_path(self._run_object, run_name='__main__')
        except SystemExit:
            pass
        return prof, None

    def profile_module(self):
        """Returns memory stats for a module."""
        target_modules = {self._run_object}
        try:
            with open(self._run_object, 'rb') as srcfile,\
                _CodeEventsTracker(target_modules) as prof:
                code = compile(srcfile.read(), self._run_object, 'exec')
                prof.compute_mem_overhead()
                exec(code, self._globs, None)
        except SystemExit:
            pass
        return prof, None

    def profile_function(self):
        """Returns memory stats for a function."""
        target_modules = {self._run_object.__code__.co_filename}
        with _CodeEventsTracker(target_modules) as prof:
            prof.compute_mem_overhead()
            result = self._run_object(*self._run_args, **self._run_kwargs)
        return prof, result

    def run(self):
        """Collects memory stats for specified Python program."""
        existing_objects = _get_in_memory_objects()
        prof, result = self.profile()
        new_objects = _get_in_memory_objects()

        new_obj_count = _get_obj_count_difference(new_objects, existing_objects)
        result_obj_count = new_obj_count - prof.obj_overhead

        # existing_objects list is also profiler overhead
        result_obj_count[list] -= 1
        pretty_obj_count = _format_obj_count(result_obj_count)
        return {
            'objectName': self._object_name,
            'codeEvents': prof.code_events,
            'totalEvents': len(prof.code_events),
            'objectsCount': pretty_obj_count,
            'result': result,
            'timestamp': int(time.time())
        }


"""Base class for a profile wrapper."""
import inspect
import multiprocessing
import os
import pkgutil
import sys
import zlib

def get_pkg_module_names(package_path):
    """Returns module filenames from package.

    Args:
        package_path: Path to Python package.
    Returns:
        A set of module filenames.
    """
    module_names = set()
    for fobj, modname, _ in pkgutil.iter_modules(path=[package_path]):
        filename = os.path.join(fobj.path, '%s.py' % modname)
        if os.path.exists(filename):
            module_names.add(os.path.abspath(filename))
    return module_names

def hash_name(name):
    """Computes hash of the name."""
    return zlib.adler32(name.encode('utf-8'))

class ProcessWithException(multiprocessing.Process):
    """Process subclass that propagates exceptions to parent process.

    Also handles sending function output to parent process.
    Args:
        parent_conn: Parent end of multiprocessing.Pipe.
        child_conn: Child end of multiprocessing.Pipe.
        result: Result of the child process.
    """

    def __init__(self, result, *args, **kwargs):
        super(ProcessWithException, self).__init__(*args, **kwargs)
        self.parent_conn, self.child_conn = multiprocessing.Pipe()
        self.result = result

    def run(self):
        try:
            self.result.update(
                self._target(*self._args, **self._kwargs))
            self.child_conn.send(None)
        except Exception as exc:  # pylint: disable=broad-except
            self.child_conn.send(exc)

    @property
    def exception(self):
        """Returns exception from child process."""
        return self.parent_conn.recv()

    @property
    def output(self):
        """Returns target function output."""
        return self.result._getvalue()  # pylint: disable=protected-access

def run_in_separate_process(func, *args, **kwargs):
    """Runs function in separate process.

    This function is used instead of a decorator, since Python multiprocessing
    module can't serialize decorated function on all platforms.
    """
    manager = multiprocessing.Manager()
    manager_dict = manager.dict()
    process = ProcessWithException(
        manager_dict, target=func, args=args, kwargs=kwargs)
    process.start()
    process.join()
    exc = process.exception
    if exc:
        raise exc
    return process.output

class BaseProfiler:
    """Base class for a profiler."""

    def __init__(self, run_object):
        """Initializes profiler.

        Args:
            run_object: object to be profiled.
        """
        run_obj_type = self.get_run_object_type(run_object)
        if run_obj_type == 'module':
            self.init_module(run_object)
        elif run_obj_type == 'package':
            self.init_package(run_object)
        else:
            self.init_function(run_object)

    @staticmethod
    def get_run_object_type(run_object):
        """Determines run object type."""
        if isinstance(run_object, tuple):
            return 'function'
        run_object, _, _ = run_object.partition(' ')
        if os.path.isdir(run_object):
            return 'package'
        return 'module'

    def init_module(self, run_object):
        """Initializes profiler with a module."""
        self.profile = self.profile_module
        self._run_object, _, self._run_args = run_object.partition(' ')
        self._object_name = '%s (module)' % self._run_object
        self._globs = {
            '__file__': self._run_object,
            '__name__': '__main__',
            '__package__': None,
        }
        program_path = os.path.dirname(self._run_object)
        if sys.path[0] != program_path:
            sys.path.insert(0, program_path)
        self._replace_sysargs()

    def init_package(self, run_object):
        """Initializes profiler with a package."""
        self.profile = self.profile_package
        self._run_object, _, self._run_args = run_object.partition(' ')
        self._object_name = '%s (package)' % self._run_object
        self._replace_sysargs()

    def init_function(self, run_object):
        """Initializes profiler with a function."""
        self.profile = self.profile_function
        self._run_object, self._run_args, self._run_kwargs = run_object
        filename = inspect.getsourcefile(self._run_object)
        self._object_name = '%s @ %s (function)' % (
            self._run_object.__name__, filename)

    def _replace_sysargs(self):
        """Replaces sys.argv with proper args to pass to script."""
        sys.argv[:] = [self._run_object]
        if self._run_args:
            sys.argv += self._run_args.split()

    def profile_package(self):
        """Profiles package specified by filesystem path.

        Runs object self._run_object as a package specified by filesystem path.
        Must be overridden.
        """
        raise NotImplementedError

    def profile_module(self):
        """Profiles module.

        Runs object self._run_object as a Python module.
        Must be overridden.
        """
        raise NotImplementedError

    def profile_function(self):
        """Profiles function.

        Runs object self._run_object as a Python function.
        Must be overridden.
        """
        raise NotImplementedError

    def run(self):
        """Runs profiler and returns collected stats."""
        return self.profile()


"""Flame graph module."""
import inspect
import runpy
import signal
import time

from collections import defaultdict
from vprof import base_profiler

_SAMPLE_INTERVAL = 0.001

class _StatProfiler:
    """Statistical profiler.

    Samples call stack at regulal intervals specified by _SAMPLE_INTERVAL.
    """

    def __init__(self):
        self._stats = defaultdict(int)
        self._start_time = None
        self.base_frame = None
        self.run_time = None

    def __enter__(self):
        """Enables statistical profiler."""
        signal.signal(signal.SIGPROF, self.sample)
        signal.setitimer(signal.ITIMER_PROF, _SAMPLE_INTERVAL)
        self._start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tbf):
        """Disables statistical profiler."""
        self.run_time = time.time() - self._start_time
        signal.setitimer(signal.ITIMER_PROF, 0)

    def sample(self, signum, frame):  #pylint: disable=unused-argument
        """Samples current stack and adds result in self._stats.

        Args:
            signum: Signal that activates handler.
            frame: Frame on top of the stack when signal is handled.
        """
        stack = []
        while frame and frame != self.base_frame:
            stack.append((
                frame.f_code.co_name,
                frame.f_code.co_filename,
                frame.f_code.co_firstlineno))
            frame = frame.f_back
        self._stats[tuple(stack)] += 1
        signal.setitimer(signal.ITIMER_PROF, _SAMPLE_INTERVAL)

    @staticmethod
    def _insert_stack(stack, sample_count, call_tree):
        """Inserts stack into the call tree.

        Args:
            stack: Call stack.
            sample_count: Sample count of call stack.
            call_tree: Call tree.
        """
        curr_level = call_tree
        for func in stack:
            next_level_index = {
                node['stack']: node for node in curr_level['children']}
            if func not in next_level_index:
                new_node = {'stack': func, 'children': [], 'sampleCount': 0}
                curr_level['children'].append(new_node)
                curr_level = new_node
            else:
                curr_level = next_level_index[func]
        curr_level['sampleCount'] = sample_count

    def _fill_sample_count(self, node):
        """Counts and fills sample counts inside call tree."""
        node['sampleCount'] += sum(
            self._fill_sample_count(child) for child in node['children'])
        return node['sampleCount']

    @staticmethod
    def _get_percentage(sample_count, total_samples):
        """Return percentage of sample_count in total_samples."""
        if total_samples != 0:
            return 100 * round(float(sample_count) / total_samples, 3)
        return 0.0

    def _format_tree(self, node, total_samples):
        """Reformats call tree for the UI."""
        funcname, filename, _ = node['stack']
        sample_percent = self._get_percentage(
            node['sampleCount'], total_samples)
        color_hash = base_profiler.hash_name('%s @ %s' % (funcname, filename))
        return {
            'stack': node['stack'],
            'children': [self._format_tree(child, total_samples)
                         for child in node['children']],
            'sampleCount': node['sampleCount'],
            'samplePercentage': sample_percent,
            'colorHash': color_hash
        }

    @property
    def call_tree(self):
        """Returns call tree."""
        call_tree = {'stack': 'base', 'sampleCount': 0, 'children': []}
        for stack, sample_count in self._stats.items():
            self._insert_stack(reversed(stack), sample_count, call_tree)
        self._fill_sample_count(call_tree)
        if not call_tree['children']:
            return {}
        return self._format_tree(
            call_tree['children'][0], call_tree['sampleCount'])

class FlameGraphProfiler(base_profiler.BaseProfiler):
    """Statistical profiler wrapper.

    Runs statistical profiler and returns collected stats.
    """

    def _profile_package(self):
        """Runs statistical profiler on a package."""
        with _StatProfiler() as prof:
            prof.base_frame = inspect.currentframe()
            try:
                runpy.run_path(self._run_object, run_name='__main__')
            except SystemExit:
                pass

        call_tree = prof.call_tree
        return {
            'objectName': self._object_name,
            'sampleInterval': _SAMPLE_INTERVAL,
            'runTime': prof.run_time,
            'callStats': call_tree,
            'totalSamples': call_tree.get('sampleCount', 0),
            'timestamp': int(time.time())
        }

    def profile_package(self):
        """Runs package profiler in separate process."""
        return base_profiler.run_in_separate_process(self._profile_package)

    def _profile_module(self):
        """Runs statistical profiler on a module."""
        with open(self._run_object, 'rb') as srcfile, _StatProfiler() as prof:
            code = compile(srcfile.read(), self._run_object, 'exec')
            prof.base_frame = inspect.currentframe()
            try:
                exec(code, self._globs, None)
            except SystemExit:
                pass

        call_tree = prof.call_tree
        return {
            'objectName': self._object_name,
            'sampleInterval': _SAMPLE_INTERVAL,
            'runTime': prof.run_time,
            'callStats': call_tree,
            'totalSamples': call_tree.get('sampleCount', 0),
            'timestamp': int(time.time())
        }

    def profile_module(self):
        """Runs module profiler in separate process."""
        return base_profiler.run_in_separate_process(self._profile_module)

    def profile_function(self):
        """Runs statistical profiler on a function."""
        with _StatProfiler() as prof:
            result = self._run_object(*self._run_args, **self._run_kwargs)

        call_tree = prof.call_tree
        return {
            'objectName': self._object_name,
            'sampleInterval': _SAMPLE_INTERVAL,
            'runTime': prof.run_time,
            'callStats': call_tree,
            'totalSamples': call_tree.get('sampleCount', 0),
            'result': result,
            'timestamp': int(time.time())
        }


"""Profiler wrapper module."""
import cProfile
import operator
import pstats
import runpy
import time

from vprof import base_profiler

class Profiler(base_profiler.BaseProfiler):
    """Python profiler wrapper.

    Runs cProfile on specified program and returns collected stats.
    """

    @staticmethod
    def _transform_stats(prof):
        """Processes collected stats for UI."""
        records = []
        for info, params in prof.stats.items():
            filename, lineno, funcname = info
            cum_calls, num_calls, time_per_call, cum_time, _ = params
            if prof.total_tt == 0:
                percentage = 0
            else:
                percentage = round(100 * (cum_time / prof.total_tt), 4)
            cum_time = round(cum_time, 4)
            func_name = '%s @ %s' % (funcname, filename)
            color_hash = base_profiler.hash_name(func_name)
            records.append(
                (filename, lineno, funcname, cum_time, percentage, num_calls,
                 cum_calls, time_per_call, filename, color_hash))
        return sorted(records, key=operator.itemgetter(4), reverse=True)

    def _profile_package(self):
        """Runs cProfile on a package."""
        prof = cProfile.Profile()
        prof.enable()
        try:
            runpy.run_path(self._run_object, run_name='__main__')
        except SystemExit:
            pass
        prof.disable()
        prof_stats = pstats.Stats(prof)
        prof_stats.calc_callees()
        return {
            'objectName': self._object_name,
            'callStats': self._transform_stats(prof_stats),
            'totalTime': prof_stats.total_tt,
            'primitiveCalls': prof_stats.prim_calls,
            'totalCalls': prof_stats.total_calls,
            'timestamp': int(time.time())
        }

    def profile_package(self):
        """Runs package profiler in separate process."""
        return base_profiler.run_in_separate_process(self._profile_package)

    def _profile_module(self):
        """Runs cProfile on a module."""
        prof = cProfile.Profile()
        try:
            with open(self._run_object, 'rb') as srcfile:
                code = compile(srcfile.read(), self._run_object, 'exec')
            prof.runctx(code, self._globs, None)
        except SystemExit:
            pass
        prof_stats = pstats.Stats(prof)
        prof_stats.calc_callees()
        return {
            'objectName': self._object_name,
            'callStats': self._transform_stats(prof_stats),
            'totalTime': prof_stats.total_tt,
            'primitiveCalls': prof_stats.prim_calls,
            'totalCalls': prof_stats.total_calls,
            'timestamp': int(time.time())
        }

    def profile_module(self):
        """Runs module profiler in separate process."""
        return base_profiler.run_in_separate_process(self._profile_module)

    def profile_function(self):
        """Runs cProfile on a function."""
        prof = cProfile.Profile()
        prof.enable()
        result = self._run_object(*self._run_args, **self._run_kwargs)
        prof.disable()
        prof_stats = pstats.Stats(prof)
        prof_stats.calc_callees()
        return {
            'objectName': self._object_name,
            'callStats': self._transform_stats(prof_stats),
            'totalTime': prof_stats.total_tt,
            'primitiveCalls': prof_stats.prim_calls,
            'totalCalls': prof_stats.total_calls,
            'result': result,
            'timestamp': int(time.time())
        }


"""Remote profiling example with Flask.

First of all launch vprof in remote mode:

    vprof -r

and launch this script:

    python guestbook.py

Then you can profile '/' and '/add' handlers:

    curl http://127.0.0.1:5000/profile/main

and

    curl --data "name=foo&message=bar" http://127.0.0.1:5000/profile/add

Check profiler_handler source code below for details.
"""

import contextlib
import jinja2
import flask
import sqlite3

from vprof import runner

DB = '/tmp/guestbook.db'
DB_SCHEMA = """
DROP TABLE IF EXISTS entry;
CREATE TABLE entry (
	id INTEGER PRIMARY KEY,
	name TEXT NOT NULL,
	message TEXT NOT NULL
);
"""
LAYOUT = """

  <title>Guestbook</title>

    <form method=post action=/add>
      Name: <input type=text name=name>

      Message: <textarea name=message></textarea>

      <input type=submit value=Post>

    </form>

    {% for entry in entries %}

{{ entry.name }}

        {{ entry.message | safe }}

    {% endfor %}

"""

def connect_to_db():
    """Establishes connection to SQLite."""
    return sqlite3.connect(DB)

def init_db():
    """Initializes DB."""
    with contextlib.closing(connect_to_db()) as db:
        db.cursor().executescript(DB_SCHEMA)
        db.commit()

app = flask.Flask('guestbook')

@app.before_request
def before_request():
    """Establishes SQLite connection before request."""
    flask.g.db = connect_to_db()

@app.teardown_request
def teardown_request(exception):
    """Closes SQLite connection after request is processed."""
    db = getattr(flask.g, 'db', None)
    if db is not None:
        db.close()

@app.route('/')
def show_guestbook():
    """Returns all existing guestbook records."""
    cursor = flask.g.db.execute(
        'SELECT name, message FROM entry ORDER BY id DESC;')
    entries = [{'name': row[0], 'message': row[1]} for row in cursor.fetchall()]
    return jinja2.Template(LAYOUT).render(entries=entries)

@app.route('/add', methods=['POST'])
def add_entry():
    """Adds single guestbook record."""
    name, msg = flask.request.form['name'], flask.request.form['message']
    flask.g.db.execute(
        'INSERT INTO entry (name, message) VALUES (?, ?)', (name, msg))
    flask.g.db.commit()
    return flask.redirect('/')

@app.route('/profile/<uri>', methods=['GET', 'POST'])
def profiler_handler(uri):
    """Profiler handler."""
    # HTTP method should be GET.
    if uri == 'main':
        runner.run(show_guestbook, 'cmhp')
    # In this case HTTP method should be POST singe add_entry uses POST
    elif uri == 'add':
        runner.run(add_entry, 'cmhp')
    return flask.redirect('/')

if __name__ == '__main__':
    init_db()
    app.run()


"""Profiler server."""
import functools
import gzip
import io
import json
import os
import socketserver
import sys
import webbrowser

from http import server

_STATIC_DIR = 'ui'
_PROFILE_HTML = '%s/profile.html' % _STATIC_DIR

class StatsServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    """Declares multithreaded HTTP server."""
    allow_reuse_address = True

class StatsHandler(server.SimpleHTTPRequestHandler):
    """Program stats request handler."""

    def __init__(self, profile_json, *args, **kwargs):
        self._profile_json = profile_json
        self.uri_map = {
            '/': self._handle_root,
            '/profile': self._handle_profile,
        }
        # Since this class is old-style - call parent method directly.
        server.SimpleHTTPRequestHandler.__init__(
            self, *args, **kwargs)

    @staticmethod
    def _handle_root():
        """Handles index.html requests."""
        res_filename = os.path.join(
            os.path.dirname(__file__), _PROFILE_HTML)
        with io.open(res_filename, 'rb') as res_file:
            content = res_file.read()
        return content, 'text/html'

    def _handle_profile(self):
        """Handles profile stats requests."""
        return json.dumps(self._profile_json).encode(), 'text/json'

    def _handle_other(self):
        """Handles static files requests."""
        res_filename = os.path.join(
            os.path.dirname(__file__), _STATIC_DIR, self.path[1:])
        with io.open(res_filename, 'rb') as res_file:
            content = res_file.read()
        _, extension = os.path.splitext(self.path)
        return content, 'text/%s' % extension[1:]  # Skip dot in the extension.

    def do_GET(self):
        """Handles HTTP GET requests."""
        handler = self.uri_map.get(self.path) or self._handle_other
        content, content_type = handler()
        compressed_content = gzip.compress(content)
        self._send_response(
            200, headers=(('Content-type', '%s; charset=utf-8' % content_type),
                          ('Content-Encoding', 'gzip'),
                          ('Content-Length', len(compressed_content))))
        self.wfile.write(compressed_content)

    def do_POST(self):
        """Handles HTTP POST requests."""
        post_data = self.rfile.read(int(self.headers['Content-Length']))
        json_data = gzip.decompress(post_data)
        self._profile_json.update(json.loads(json_data.decode('utf-8')))
        self._send_response(
            200, headers=(('Content-type', '%s; charset=utf-8' % 'text/json'),
                          ('Content-Encoding', 'gzip'),
                          ('Content-Length', len(post_data))))

    def _send_response(self, http_code, message=None, headers=None):
        """Sends HTTP response code, message and headers."""
        self.send_response(http_code, message)
        if headers:
            for header in headers:
                self.send_header(*header)
            self.end_headers()

def start(host, port, profiler_stats, dont_start_browser, debug_mode):
    """Starts HTTP server with specified parameters.

    Args:
        host: Server host name.
        port: Server port.
        profiler_stats: A dict with collected program stats.
        dont_start_browser: Whether to open browser after profiling.
        debug_mode: Whether to redirect stderr to /dev/null.
    """
    stats_handler = functools.partial(StatsHandler, profiler_stats)
    if not debug_mode:
        sys.stderr = open(os.devnull, 'w')
    print('Starting HTTP server...')
    if not dont_start_browser:
        webbrowser.open('http://{}:{}/'.format(host, port))
    try:
        StatsServer((host, port), stats_handler).serve_forever()
    except KeyboardInterrupt:
        print('Stopping...')
        sys.exit(0)


"""Profiler's main module."""
# pylint: disable=wrong-import-position
import builtins
import os
import psutil

# Take process RSS in order to compute profiler memory overhead.
builtins.initial_rss_size = psutil.Process(os.getpid()).memory_info().rss

import argparse
import json
import sys

from vprof import runner
from vprof import stats_server

__version__ = '0.38'

_PROGRAN_NAME = 'vprof'
_MODULE_DESC = 'Visual profiler for Python'
_HOST, _PORT = 'localhost', 8000
_CONFIG_DESC = (
    """profile program SRC with configuration CONFIG
available CONFIG options
  c - flame graph
  m - memory graph
  h - code heatmap""")
_ERR_CODES = {
    'ambiguous_configuration': 1,
    'bad_option': 2,
    'input_file_error': 3
}

def main():
    """Main function of the module."""
    parser = argparse.ArgumentParser(
        prog=_PROGRAN_NAME, description=_MODULE_DESC,
        formatter_class=argparse.RawTextHelpFormatter)
    launch_modes = parser.add_mutually_exclusive_group(required=True)
    launch_modes.add_argument('-r', '--remote', dest='remote',
                              action='store_true', default=False,
                              help='launch in remote mode')
    launch_modes.add_argument('-i', '--input-file', dest='input_file',
                              type=str, default='',
                              help='render UI from file')
    launch_modes.add_argument('-c', '--config', nargs=2, dest='config',
                              help=_CONFIG_DESC, metavar=('CONFIG', 'SRC'))
    parser.add_argument('-H', '--host', dest='host', default=_HOST, type=str,
                        help='set internal webserver host')
    parser.add_argument('-p', '--port', dest='port', default=_PORT, type=int,
                        help='set internal webserver port')
    parser.add_argument('-n', '--no-browser', dest='dont_start_browser',
                        action='store_true', default=False,
                        help="don't start browser automatically")
    parser.add_argument('-o', '--output-file', dest='output_file',
                        type=str, default='', help='save profile to file')
    parser.add_argument('--debug', dest='debug_mode',
                        action='store_true', default=False,
                        help="don't suppress error messages")
    parser.add_argument('--version', action='version',
                        version='vprof %s' % __version__)
    args = parser.parse_args()

    # Render UI from file.
    if args.input_file:
        with open(args.input_file) as ifile:
            saved_stats = json.loads(ifile.read())
            if saved_stats['version'] != __version__:
                print('Incorrect profiler version - %s. %s is required.' % (
                    saved_stats['version'], __version__))
                sys.exit(_ERR_CODES['input_file_error'])
            stats_server.start(args.host, args.port, saved_stats,
                               args.dont_start_browser, args.debug_mode)
    # Launch in remote mode.
    elif args.remote:
        stats_server.start(args.host, args.port, {},
                           args.dont_start_browser, args.debug_mode)
    # Profiler mode.
    else:
        config, source = args.config
        try:
            program_stats = runner.run_profilers(source, config, verbose=True)
        except runner.AmbiguousConfigurationError:
            print('Profiler configuration %s is ambiguous. '
                  'Please, remove duplicates.' % config)
            sys.exit(_ERR_CODES['ambiguous_configuration'])
        except runner.BadOptionError as exc:
            print(exc)
            sys.exit(_ERR_CODES['bad_option'])

        if args.output_file:
            with open(args.output_file, 'w') as outfile:
                program_stats['version'] = __version__
                outfile.write(json.dumps(program_stats, indent=2))
        else:
            stats_server.start(
                args.host, args.port, program_stats,
                args.dont_start_browser, args.debug_mode)

if __name__ == "__main__":
    main()


"""Code heatmap module."""
import inspect
import fnmatch
import os
import runpy
import sys
import time

from collections import defaultdict
from collections import deque
from vprof import base_profiler

_STDLIB_PATHS = [
    os.path.abspath(path) for path in sys.path
    if os.path.isdir(path) and path.startswith(sys.prefix)]

def check_standard_dir(module_path):
    """Checks whether path belongs to standard library or installed modules."""
    if 'site-packages' in module_path:
        return True
    for stdlib_path in _STDLIB_PATHS:
        if fnmatch.fnmatchcase(module_path, stdlib_path + '*'):
            return True
    return False

class _CodeHeatmapCalculator:
    """Calculates Python code heatmap.

    Class that contains all logic related to calculating code heatmap
    for a Python program.
    """

    def __init__(self):
        self.original_trace_function = sys.gettrace()
        self.prev_lineno = None
        self.prev_timestamp = None
        self.prev_path = None
        self.lines = deque()
        self._execution_count = defaultdict(lambda: defaultdict(int))
        self._heatmap = defaultdict(lambda: defaultdict(float))

    def __enter__(self):
        """Enables heatmap calculator."""
        sys.settrace(self.record_line)
        return self

    def __exit__(self, exc_type, exc_val, exc_tbf):
        """Disables heatmap calculator."""
        sys.settrace(self.original_trace_function)
        if self.prev_timestamp:
            runtime = time.time() - self.prev_timestamp
            self.lines.append([self.prev_path, self.prev_lineno, runtime])

    def record_line(self, frame, event, arg):  # pylint: disable=unused-argument
        """Records line execution time."""
        if event == 'line':
            if self.prev_timestamp:
                runtime = time.time() - self.prev_timestamp
                self.lines.append([self.prev_path, self.prev_lineno, runtime])
            self.prev_lineno = frame.f_lineno
            self.prev_path = frame.f_code.co_filename
            self.prev_timestamp = time.time()
        return self.record_line

    @property
    def lines_without_stdlib(self):
        """Filters code from standard library from self.lines."""
        prev_line = None
        current_module_path = inspect.getabsfile(inspect.currentframe())
        for module_path, lineno, runtime in self.lines:
            module_abspath = os.path.abspath(module_path)
            if not prev_line:
                prev_line = [module_abspath, lineno, runtime]
            else:
                if (not check_standard_dir(module_path) and
                        module_abspath != current_module_path):
                    yield prev_line
                    prev_line = [module_abspath, lineno, runtime]
                else:
                    prev_line[2] += runtime
        yield prev_line

    def fill_heatmap(self):
        """Fills code heatmap and execution count dictionaries."""
        for module_path, lineno, runtime in self.lines_without_stdlib:
            self._execution_count[module_path][lineno] += 1
            self._heatmap[module_path][lineno] += runtime

    @property
    def heatmap(self):
        """Returns heatmap with absolute path names."""
        if not self._heatmap:
            self.fill_heatmap()
        return self._heatmap

    @property
    def execution_count(self):
        """Returns execution count map with absolute path names."""
        if not self._execution_count:
            self.fill_heatmap()
        return self._execution_count

class CodeHeatmapProfiler(base_profiler.BaseProfiler):
    """Code heatmap wrapper."""

    SKIP_LINES = 10
    MIN_SKIP_SIZE = 100

    def _calc_skips(self, heatmap, num_lines):
        """Calculates skip map for large sources.
        Skip map is a list of tuples where first element of tuple is line
        number and second is length of the skip region:
            [(1, 10), (15, 10)] means skipping 10 lines after line 1 and
            10 lines after line 15.
        """
        if num_lines < self.MIN_SKIP_SIZE:
            return []
        skips, prev_line = [], 0
        for line in sorted(heatmap):
            curr_skip = line - prev_line - 1
            if curr_skip > self.SKIP_LINES:
                skips.append((prev_line, curr_skip))
            prev_line = line
        if num_lines - prev_line > self.SKIP_LINES:
            skips.append((prev_line, num_lines - prev_line))
        return skips

    @staticmethod
    def _skip_lines(src_code, skip_map):
        """Skips lines in src_code specified by skip map."""
        if not skip_map:
            return [['line', j + 1, l] for j, l in enumerate(src_code)]
        code_with_skips, i = [], 0
        for line, length in skip_map:
            code_with_skips.extend(
                ['line', i + j + 1, l] for j, l in enumerate(src_code[i:line]))
            if (code_with_skips
                    and code_with_skips[-1][0] == 'skip'):  # Merge skips.
                code_with_skips[-1][1] += length
            else:
                code_with_skips.append(['skip', length])
            i = line + length
        code_with_skips.extend(
            ['line', i + j + 1, l] for j, l in enumerate(src_code[i:]))
        return code_with_skips

    def _profile_package(self):
        """Calculates heatmap for package."""
        with _CodeHeatmapCalculator() as prof:
            try:
                runpy.run_path(self._run_object, run_name='__main__')
            except SystemExit:
                pass

        heatmaps = []
        for filename, heatmap in prof.heatmap.items():
            if os.path.isfile(filename):
                heatmaps.append(
                    self._format_heatmap(
                        filename, heatmap, prof.execution_count[filename]))

        run_time = sum(heatmap['runTime'] for heatmap in heatmaps)
        return {
            'objectName': self._run_object,
            'runTime': run_time,
            'heatmaps': heatmaps
        }

    def profile_package(self):
        """Runs package profiler in separate process."""
        return base_profiler.run_in_separate_process(self._profile_package)

    def _format_heatmap(self, filename, heatmap, execution_count):
        """Formats heatmap for UI."""
        with open(filename) as src_file:
            file_source = src_file.read().split('\n')
            skip_map = self._calc_skips(heatmap, len(file_source))
        run_time = sum(time for time in heatmap.values())
        return {
            'name': filename,
            'heatmap': heatmap,
            'executionCount': execution_count,
            'srcCode': self._skip_lines(file_source, skip_map),
            'runTime': run_time
        }

    def _profile_module(self):
        """Calculates heatmap for module."""
        with open(self._run_object, 'r') as srcfile:
            src_code = srcfile.read()
            code = compile(src_code, self._run_object, 'exec')
        try:
            with _CodeHeatmapCalculator() as prof:
                exec(code, self._globs, None)
        except SystemExit:
            pass

        heatmaps = []
        for filename, heatmap in prof.heatmap.items():
            if os.path.isfile(filename):
                heatmaps.append(
                    self._format_heatmap(
                        filename, heatmap, prof.execution_count[filename]))

        run_time = sum(heatmap['runTime'] for heatmap in heatmaps)
        return {
            'objectName': self._run_object,
            'runTime': run_time,
            'heatmaps': heatmaps
        }

    def profile_module(self):
        """Runs module profiler in separate process."""
        return base_profiler.run_in_separate_process(self._profile_module)

    def profile_function(self):
        """Calculates heatmap for function."""
        with _CodeHeatmapCalculator() as prof:
            result = self._run_object(*self._run_args, **self._run_kwargs)
        code_lines, start_line = inspect.getsourcelines(self._run_object)

        source_lines = []
        for line in code_lines:
            source_lines.append(('line', start_line, line))
            start_line += 1

        filename = os.path.abspath(inspect.getsourcefile(self._run_object))
        heatmap = prof.heatmap[filename]
        run_time = sum(time for time in heatmap.values())
        return {
            'objectName': self._object_name,
            'runTime': run_time,
            'result': result,
            'timestamp': int(time.time()),
            'heatmaps': [{
                'name': self._object_name,
                'heatmap': heatmap,
                'executionCount': prof.execution_count[filename],
                'srcCode': source_lines,
                'runTime': run_time
            }]
        }


"""Module with functions that run profilers."""
# pylint: disable=wrong-import-position
import builtins
import gzip
import os
import psutil
import urllib.request

# Take initial RSS in order to compute profiler memory overhead
# when profiling single functions.
if not hasattr(builtins, 'initial_rss_size'):
    builtins.initial_rss_size = psutil.Process(os.getpid()).memory_info().rss

import json

from collections import OrderedDict
from vprof import code_heatmap
from vprof import flame_graph
from vprof import memory_profiler
from vprof import profiler

_PROFILERS = (
    ('m', memory_profiler.MemoryProfiler),
    ('c', flame_graph.FlameGraphProfiler),
    ('h', code_heatmap.CodeHeatmapProfiler),
    ('p', profiler.Profiler)
)

class Error(Exception):
    """Base exception for current module."""
    pass  # pylint: disable=unnecessary-pass

class AmbiguousConfigurationError(Error):
    """Raised when profiler configuration is ambiguous."""
    pass  # pylint: disable=unnecessary-pass

class BadOptionError(Error):
    """Raised when unknown options are present in the configuration."""
    pass  # pylint: disable=unnecessary-pass

def run_profilers(run_object, prof_config, verbose=False):
    """Runs profilers on run_object.

    Args:
        run_object: An object (string or tuple) for profiling.
        prof_config: A string with profilers configuration.
        verbose: True if info about running profilers should be shown.
    Returns:
        An ordered dictionary with collected stats.
    Raises:
        AmbiguousConfigurationError: when prof_config is ambiguous.
        BadOptionError: when unknown options are present in configuration.
    """
    if len(prof_config) > len(set(prof_config)):
        raise AmbiguousConfigurationError(
            'Profiler configuration %s is ambiguous' % prof_config)

    available_profilers = {opt for opt, _ in _PROFILERS}
    for option in prof_config:
        if option not in available_profilers:
            raise BadOptionError('Unknown option: %s' % option)

    run_stats = OrderedDict()
    present_profilers = ((o, p) for o, p in _PROFILERS if o in prof_config)
    for option, prof in present_profilers:
        curr_profiler = prof(run_object)
        if verbose:
            print('Running %s...' % curr_profiler.__class__.__name__)
        run_stats[option] = curr_profiler.run()
    return run_stats

def run(func, options, args=(), kwargs={}, host='localhost', port=8000):  # pylint: disable=dangerous-default-value
    """Runs profilers on a function.

    Args:
        func: A Python function.
        options: A string with profilers configuration (i.e. 'cmh').
        args: func non-keyword arguments.
        kwargs: func keyword arguments.
        host: Host name to send collected data.
        port: Port number to send collected data.

    Returns:
        A result of func execution.
    """
    run_stats = run_profilers((func, args, kwargs), options)

    result = None
    for prof in run_stats:
        if not result:
            result = run_stats[prof]['result']
        del run_stats[prof]['result']  # Don't send result to remote host

    post_data = gzip.compress(
        json.dumps(run_stats).encode('utf-8'))
    urllib.request.urlopen('http://%s:%s' % (host, port), post_data)
    return result


# -*- coding: utf-8 -*-

import operator

import numpy as np
import scipy.sparse as sp
from sklearn.base import copy
from sklearn.naive_bayes import (BaseDiscreteNB, BaseNB, BernoulliNB,
                                 GaussianNB, MultinomialNB)

from splearn import BlockRDD
from splearn.base import SparkClassifierMixin
from splearn.utils.validation import check_rdd

class SparkBaseNB(BaseNB, SparkClassifierMixin):

    """Abstract base class for distributed naive Bayes estimators"""

    def __radd__(self, other):
        return self if other == 0 else self.__add__(other)

    def predict(self, X):
        """
        Perform classification on an RDD containing arrays of test vectors X.

        Parameters
        ----------
        X : RDD containing array-like items, shape = [m_samples, n_features]

        Returns
        -------
        C : RDD with arrays, shape = [n_samples]
            Predicted target values for X
        """
        check_rdd(X, (sp.spmatrix, np.ndarray))
        return X.map(
            lambda X: super(SparkBaseNB, self).predict(X))

    def predict_proba(self, X):
        """
        Return probability estimates for the RDD containing test vector X.

        Parameters
        ----------
        X : RDD containing array-like items, shape = [m_samples, n_features]

        Returns
        -------
        C : RDD with array-like items , shape = [n_samples, n_classes]
            Returns the probability of the samples for each class in
            the models for each RDD block. The columns correspond to the classes
            in sorted order, as they appear in the attribute `classes_`.
        """
        check_rdd(X, (sp.spmatrix, np.ndarray))
        return X.map(
            lambda X: super(SparkBaseNB, self).predict_proba(X))

    def predict_log_proba(self, X):
        """
        Return log-probability estimates for the RDD containing the
        test vector X.

        Parameters
        ----------
        X : RDD containing array-like items, shape = [m_samples, n_features]

        Returns
        -------
        C : RDD with array-like items, shape = [n_samples, n_classes]
            Returns the log-probability of the samples for each class in
            the model for each RDD block. The columns correspond to the classes
            in sorted order, as they appear in the attribute `classes_`.
        """
        # required, scikit call self.predict_log_proba(X) in predict_proba
        # and thus this function is call, it must have the same behavior when
        # not called by sparkit-learn
        if not isinstance(X, BlockRDD):
            return super(SparkBaseNB, self).predict_log_proba(X)

        check_rdd(X, (sp.spmatrix, np.ndarray))
        return X.map(
            lambda X: super(SparkBaseNB, self).predict_log_proba(X))

    def _to_scikit(self, cls):
        new = cls()
        new.__dict__ = self.__dict__
        return new

class SparkGaussianNB(GaussianNB, SparkBaseNB):

    """
    Distributed Gaussian Naive Bayes (SparkGaussianNB)

    Based on sklearn's GaussianNB, model averaging is based on the work of
    the Stanford CS tech report STAN-CS-79-773 by Chan, Golub, and LeVeque:

        http://i.stanford.edu/pub/cstr/reports/cs/tr/79/773/CS-TR-79-773.pdf

    Attributes
    ----------
    class_prior_ : array, shape (n_classes,)
        probability of each class.

    class_count_ : array, shape (n_classes,)
        number of training samples observed in each class.

    theta_ : array, shape (n_classes, n_features)
        mean of each feature per class

    sigma_ : array, shape (n_classes, n_features)
        variance of each feature per class

    Examples
    --------
    TODO!
    """

    def fit(self, Z, classes=None):
        """Fit Gaussian Naive Bayes according to X, y

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training vectors, where n_samples is the number of samples
            and n_features is the number of features.

        y : array-like, shape (n_samples,)
            Target values.

        Returns
        -------
        self : object
            Returns self.
        """
        check_rdd(Z, {'X': (sp.spmatrix, np.ndarray), 'y': (sp.spmatrix, np.ndarray)})
        models = Z[:, ['X', 'y']].map(
            lambda X_y: self.partial_fit(X_y[0], X_y[1], classes))
        avg = models.reduce(operator.add)
        self.__dict__.update(avg.__dict__)
        return self

    def __add__(self, other):
        this = copy.deepcopy(self)
        this_class2idx = {cls: idx for idx, cls in enumerate(this.classes_)}
        other_class2idx = {cls: idx for idx, cls in enumerate(other.classes_)}
        for class_i in this.classes_:
            i = this_class2idx[class_i]
            j = other_class2idx[class_i]
            N_x = this.class_count_[i]
            N_y = other.class_count_[j]
            mu_x = this.theta_[i, :]
            mu_y = other.theta_[j, :]
            sigma_x = this.sigma_[i, :]
            sigma_y = this.sigma_[j, :]
            N_total = N_x + N_y

            mu_xy = N_x * mu_x + N_y * mu_y
            sigma_xy = (sigma_x * N_x + sigma_y * N_y +
                        (N_x * N_y * (mu_x - mu_y) ** 2) / N_total)

            this.theta_[i, :] = mu_xy / N_total
            this.sigma_[i, :] = sigma_xy / N_total
            this.class_count_[i] += N_y

        this.class_prior_[:] = this.class_count_ / np.sum(this.class_count_)
        return this

    def to_scikit(self):
        return self._to_scikit(GaussianNB)

class SparkBaseDiscreteNB(BaseDiscreteNB, SparkBaseNB):

    """
    Abstract base class for distributed naive Bayes on discrete/categorical
    data. It provides the necessary methods for model averaging.

    """

    def __add__(self, other):
        """
        Add method for DiscreteNB models.

        Parameters
        ----------
        other : fitted splearn multinomilal NB model with class_count_
                and feature_count_ attribute
            Model to add.

        Returns
        -------
        model : splearn Naive Bayes model
            Model with updated coefficients.
        """
        # The rdd operator add does not consider __radd__ :(
        if other == 0:
            return self
        model = copy.deepcopy(self)
        model.class_count_ += other.class_count_
        model.feature_count_ += other.feature_count_
        model._update_class_log_prior()
        model._update_feature_log_prob()
        return model

    def fit(self, Z, classes=None):
        """
        TODO fulibacsi fix docstring
        Fit Multinomial Naive Bayes according to (X,y) pair
        which is zipped into TupleRDD Z.

        Parameters
        ----------
        Z : TupleRDD containing X [array-like, shape (m_samples, n_features)]
            and y [array-like, shape (m_samples,)] tuple
            Training vectors, where ,_samples is the number of samples in the
            block and n_features is the number of features, and y contains
            the target values.

        Returns
        -------
        self : object
            Returns self.
        """
        check_rdd(Z, {'X': (sp.spmatrix, np.ndarray), 'y': (sp.spmatrix, np.ndarray)})
        if 'w' in Z.columns:
            models = Z[:, ['X', 'y', 'w']].map(
                lambda X_y_w: self.partial_fit(
                    X_y_w[0], X_y_w[1], classes, X_y_w[2]
                )
            )
        else:
            models = Z[:, ['X', 'y']].map(
                lambda X_y: self.partial_fit(X_y[0], X_y[1], classes))
        avg = models.sum()
        self.__dict__.update(avg.__dict__)
        return self

class SparkMultinomialNB(MultinomialNB, SparkBaseDiscreteNB):
    def to_scikit(self):
        return self._to_scikit(MultinomialNB)

class SparkBernoulliNB(BernoulliNB, SparkBaseDiscreteNB):
    def to_scikit(self):
        return self._to_scikit(BernoulliNB)


# -*- coding: utf-8 -*-

import numbers
import operator
from itertools import chain

import numpy as np
import scipy.sparse as sp
import six
from pyspark import AccumulatorParam
from sklearn.feature_extraction.text import (CountVectorizer,
                                             HashingVectorizer,
                                             TfidfTransformer,
                                             _document_frequency,
                                             _make_int_array)
from sklearn.utils.fixes import frombuffer_empty
from sklearn.utils.validation import check_is_fitted

from ..base import SparkBroadcasterMixin
from ..rdd import DictRDD
from ..utils.validation import check_rdd

class SparkCountVectorizer(CountVectorizer, SparkBroadcasterMixin):

    """Distributed implementation of CountVectorizer.

    Convert a collection of text documents to a matrix of token counts

    This implementation produces a sparse representation of the counts using
    scipy.sparse.coo_matrix.

    If you do not provide an a-priori dictionary and you do not use an analyzer
    that does some kind of feature selection then the number of features will
    be equal to the vocabulary size found by analyzing the data.

    Parameters
    ----------
    input : string {'filename', 'file', 'content'}
        If 'filename', the sequence passed as an argument to fit is
        expected to be a list of filenames that need reading to fetch
        the raw content to analyze.
        If 'file', the sequence items must have a 'read' method (file-like
        object) that is called to fetch the bytes in memory.
        Otherwise the input is expected to be the sequence strings or
        bytes items are expected to be analyzed directly.
    encoding : string, 'utf-8' by default.
        If bytes or files are given to analyze, this encoding is used to
        decode.
    decode_error : {'strict', 'ignore', 'replace'}
        Instruction on what to do if a byte sequence is given to analyze that
        contains characters not of the given `encoding`. By default, it is
        'strict', meaning that a UnicodeDecodeError will be raised. Other
        values are 'ignore' and 'replace'.
    strip_accents : {'ascii', 'unicode', None}
        Remove accents during the preprocessing step.
        'ascii' is a fast method that only works on characters that have
        an direct ASCII mapping.
        'unicode' is a slightly slower method that works on any characters.
        None (default) does nothing.
    analyzer : string, {'word', 'char', 'char_wb'} or callable
        Whether the feature should be made of word or character n-grams.
        Option 'char_wb' creates character n-grams only from text inside
        word boundaries.
        If a callable is passed it is used to extract the sequence of features
        out of the raw, unprocessed input.
    preprocessor : callable or None (default)
        Override the preprocessing (string transformation) stage while
        preserving the tokenizing and n-grams generation steps.
    tokenizer : callable or None (default)
        Override the string tokenization step while preserving the
        preprocessing and n-grams generation steps.
    ngram_range : tuple (min_n, max_n)
        The lower and upper boundary of the range of n-values for different
        n-grams to be extracted. All values of n such that min_n <= n <= max_n
        will be used.
    stop_words : string {'english'}, list, or None (default)
        If 'english', a built-in stop word list for English is used.
        If a list, that list is assumed to contain stop words, all of which
        will be removed from the resulting tokens.
        If None, no stop words will be used. max_df can be set to a value
        in the range [0.7, 1.0) to automatically detect and filter stop
        words based on intra corpus document frequency of terms.
    lowercase : boolean, True by default
        Convert all characters to lowercase before tokenizing.
    token_pattern : string
        Regular expression denoting what constitutes a "token", only used
        if `tokenize == 'word'`. The default regexp select tokens of 2
        or more alphanumeric characters (punctuation is completely ignored
        and always treated as a token separator).
    max_df : float in range [0.0, 1.0] or int, optional, 1.0 by default
        When building the vocabulary ignore terms that have a document
        frequency strictly higher than the given threshold (corpus-specific
        stop words).
        If float, the parameter represents a proportion of documents, integer
        absolute counts.
        This parameter is ignored if vocabulary is not None.
    min_df : float in range [0.0, 1.0] or int, optional, 1 by default
        When building the vocabulary ignore terms that have a document
        frequency strictly lower than the given threshold. This value is also
        called cut-off in the literature.
        If float, the parameter represents a proportion of documents, integer
        absolute counts.
        This parameter is ignored if vocabulary is not None.
    max_features : optional, None by default
        If not None, build a vocabulary that only consider the top
        max_features ordered by term frequency across the corpus.
        This parameter is ignored if vocabulary is not None.
    vocabulary : Mapping or iterable, optional
        Either a Mapping (e.g., a dict) where keys are terms and values are
        indices in the feature matrix, or an iterable over terms. If not
        given, a vocabulary is determined from the input documents. Indices
        in the mapping should not be repeated and should not have any gap
        between 0 and the largest index.
    binary : boolean, False by default.
        If True, all non zero counts are set to 1. This is useful for discrete
        probabilistic models that model binary events rather than integer
        counts.
    dtype : type, optional
        Type of the matrix returned by fit_transform() or transform().

    Attributes
    ----------
    vocabulary_ : dict
        A mapping of terms to feature indices.
    stop_words_ : set
        Terms that were ignored because
        they occurred in either too many
        (`max_df`) or in too few (`min_df`) documents.
        This is only available if no vocabulary was given.

    See also
    --------
    HashingVectorizer, TfidfVectorizer
    """

    __transient__ = ['vocabulary_']

    def to_scikit(self):
        obj = CountVectorizer([])
        obj.__dict__ = self.__dict__
        return obj

    def _init_vocab(self, analyzed_docs):
        """Create vocabulary
        """
        class SetAccum(AccumulatorParam):

            def zero(self, initialValue):
                return set(initialValue)

            def addInPlace(self, v1, v2):
                v1 |= v2
                return v1

        if not self.fixed_vocabulary_:
            accum = analyzed_docs._rdd.context.accumulator(set(), SetAccum())
            analyzed_docs.foreach(
                lambda x: accum.add(set(chain.from_iterable(x))))
            vocabulary = {t: i for i, t in enumerate(accum.value)}
        else:
            vocabulary = self.vocabulary_

        if not vocabulary:
            raise ValueError("empty vocabulary; perhaps the documents only"
                             " contain stop words")
        return vocabulary

    def _count_vocab(self, analyzed_docs):
        """Create sparse feature matrix, and vocabulary where fixed_vocab=False
        """
        vocabulary = self.vocabulary_
        j_indices = _make_int_array()
        indptr = _make_int_array()
        indptr.append(0)
        for doc in analyzed_docs:
            for feature in doc:
                try:
                    j_indices.append(vocabulary[feature])
                except KeyError:
                    # Ignore out-of-vocabulary items for fixed_vocab=True
                    continue
            indptr.append(len(j_indices))

        j_indices = frombuffer_empty(j_indices, dtype=np.intc)
        indptr = np.frombuffer(indptr, dtype=np.intc)
        values = np.ones(len(j_indices))

        X = sp.csr_matrix((values, j_indices, indptr),
                          shape=(len(indptr) - 1, len(vocabulary)),
                          dtype=self.dtype)
        X.sum_duplicates()

        if self.binary:
            X.data.fill(1)

        return X

    def _sort_features(self, vocabulary):
        """Sort features by name

        Returns a reordered matrix and modifies the vocabulary in place
        """
        sorted_features = sorted(six.iteritems(vocabulary))
        map_index = np.empty(len(sorted_features), dtype=np.int32)
        for new_val, (term, old_val) in enumerate(sorted_features):
            map_index[new_val] = old_val
            vocabulary[term] = new_val

        return map_index

    def _limit_features(self, X, vocabulary, high=None, low=None,
                        limit=None):
        """Remove too rare or too common features.

        Prune features that are non zero in more samples than high or less
        documents than low, modifying the vocabulary, and restricting it to
        at most the limit most frequent.

        This does not prune samples with zero features.
        """
        if high is None and low is None and limit is None:
            return X, set()

        # Calculate a mask based on document frequencies
        dfs = X.map(_document_frequency).sum()
        tfs = X.map(lambda x: np.asarray(x.sum(axis=0))).sum().ravel()
        mask = np.ones(len(dfs), dtype=bool)
        if high is not None:
            mask &= dfs <= high
        if low is not None:
            mask &= dfs >= low
        if limit is not None and mask.sum() > limit:
            mask_inds = (-tfs[mask]).argsort()[:limit]
            new_mask = np.zeros(len(dfs), dtype=bool)
            new_mask[np.where(mask)[0][mask_inds]] = True
            mask = new_mask

        new_indices = np.cumsum(mask) - 1  # maps old indices to new
        removed_terms = set()
        for term, old_index in list(six.iteritems(vocabulary)):
            if mask[old_index]:
                vocabulary[term] = new_indices[old_index]
            else:
                del vocabulary[term]
                removed_terms.add(term)
        kept_indices = np.where(mask)[0]

        if len(kept_indices) == 0:
            raise ValueError("After pruning, no terms remain. Try a lower"
                             " min_df or a higher max_df.")

        return kept_indices, removed_terms

    def fit(self, Z):
        """Learn a vocabulary dictionary of all tokens in the raw documents in
        the DictRDD's 'X' column.

        Parameters
        ----------
        raw_documents : iterable or DictRDD with column 'X'
            An iterable which yields either str, unicode or file objects; or a
            DictRDD with column 'X' containing such iterables.

        Returns
        -------
        self
        """
        self.fit_transform(Z)
        return self

    def fit_transform(self, Z):
        """Learn the vocabulary dictionary and return term-document matrix.

        This is equivalent to fit followed by transform, but more efficiently
        implemented.

        Parameters
        ----------
        Z : iterable or DictRDD with column 'X'
            An iterable of raw_documents which yields either str, unicode or
            file objects; or a DictRDD with column 'X' containing such
            iterables.

        Returns
        -------
        X : array, [n_samples, n_features] or DictRDD
            Document-term matrix.
        """
        self._validate_vocabulary()

        # map analyzer and cache result
        analyze = self.build_analyzer()
        A = Z.transform(lambda X: list(map(analyze, X)), column='X').persist()

        # create vocabulary
        X = A[:, 'X'] if isinstance(A, DictRDD) else A
        self.vocabulary_ = self._init_vocab(X)

        # transform according to vocabulary
        mapper = self.broadcast(self._count_vocab, A.context)
        Z = A.transform(mapper, column='X', dtype=sp.spmatrix)

        if not self.fixed_vocabulary_:
            X = Z[:, 'X'] if isinstance(Z, DictRDD) else Z

            max_df = self.max_df
            min_df = self.min_df
            max_features = self.max_features

            # limit features according to min_df, max_df parameters
            n_doc = X.shape[0]
            max_doc_count = (max_df
                             if isinstance(max_df, numbers.Integral)
                             else max_df * n_doc)
            min_doc_count = (min_df
                             if isinstance(min_df, numbers.Integral)
                             else min_df * n_doc)
            if max_doc_count < min_doc_count:
                raise ValueError(
                    "max_df corresponds to < documents than min_df")
            kept_indices, self.stop_words_ = self._limit_features(
                X, self.vocabulary_, max_doc_count, min_doc_count, max_features)

            # sort features
            map_index = self._sort_features(self.vocabulary_)

            # combined mask
            mask = kept_indices[map_index]

            Z = Z.transform(lambda x: x[:, mask], column='X', dtype=sp.spmatrix)
        A.unpersist()
        return Z

    def transform(self, Z):
        """Transform documents to document-term matrix.

        Extract token counts out of raw text documents using the vocabulary
        fitted with fit or the one provided to the constructor.

        Parameters
        ----------
        raw_documents : iterable
            An iterable which yields either str, unicode or file objects.

        Returns
        -------
        X : sparse matrix, [n_samples, n_features]
            Document-term matrix.
        """
        if not hasattr(self, 'vocabulary_'):
            self._validate_vocabulary()

        self._check_vocabulary()

        analyze = self.build_analyzer()
        mapper = self.broadcast(self._count_vocab, Z.context)

        Z = Z.transform(lambda X: list(map(analyze, X)), column='X') \
             .transform(mapper, column='X', dtype=sp.spmatrix)

        return Z

class SparkHashingVectorizer(HashingVectorizer):

    """Distributed implementation of Hashingvectorizer.

    Convert a collection of text documents to a matrix of token occurrences

    It turns a collection of text documents into a scipy.sparse matrix holding
    token occurrence counts (or binary occurrence information), possibly
    normalized as token frequencies if norm='l1' or projected on the euclidean
    unit sphere if norm='l2'.

    This text vectorizer implementation uses the hashing trick to find the
    token string name to feature integer index mapping.

    This strategy has several advantages:
    - it is very low memory scalable to large datasets as there is no need to
      store a vocabulary dictionary in memory
    - it is fast to pickle and un-pickle as it holds no state besides the
      constructor parameters
    - it can be used in a streaming (partial fit) or parallel pipeline as there
      is no state computed during fit.

    There are also a couple of cons (vs using a CountVectorizer with an
    in-memory vocabulary):
    - there is no way to compute the inverse transform (from feature indices to
      string feature names) which can be a problem when trying to introspect
      which features are most important to a model.
    - there can be collisions: distinct tokens can be mapped to the same
      feature index. However in practice this is rarely an issue if n_features
      is large enough (e.g. 2 ** 18 for text classification problems).
    - no IDF weighting as this would render the transformer stateful.

    The hash function employed is the signed 32-bit version of Murmurhash3.

    Parameters
    ----------
    input: string {'filename', 'file', 'content'}
        If 'filename', the sequence passed as an argument to fit is
        expected to be a list of filenames that need reading to fetch
        the raw content to analyze.
        If 'file', the sequence items must have a 'read' method (file-like
        object) that is called to fetch the bytes in memory.
        Otherwise the input is expected to be the sequence strings or
        bytes items are expected to be analyzed directly.
    encoding : string, 'utf-8' by default.
        If bytes or files are given to analyze, this encoding is used to
        decode.
    decode_error : {'strict', 'ignore', 'replace'}
        Instruction on what to do if a byte sequence is given to analyze that
        contains characters not of the given `encoding`. By default, it is
        'strict', meaning that a UnicodeDecodeError will be raised. Other
        values are 'ignore' and 'replace'.
    strip_accents: {'ascii', 'unicode', None}
        Remove accents during the preprocessing step.
        'ascii' is a fast method that only works on characters that have
        an direct ASCII mapping.
        'unicode' is a slightly slower method that works on any characters.
        None (default) does nothing.
    analyzer: string, {'word', 'char', 'char_wb'} or callable
        Whether the feature should be made of word or character n-grams.
        Option 'char_wb' creates character n-grams only from text inside
        word boundaries.
        If a callable is passed it is used to extract the sequence of features
        out of the raw, unprocessed input.
    preprocessor: callable or None (default)
        Override the preprocessing (string transformation) stage while
        preserving the tokenizing and n-grams generation steps.
    tokenizer: callable or None (default)
        Override the string tokenization step while preserving the
        preprocessing and n-grams generation steps.
    ngram_range: tuple (min_n, max_n)
        The lower and upper boundary of the range of n-values for different
        n-grams to be extracted. All values of n such that min_n <= n <= max_n
        will be used.
    stop_words: string {'english'}, list, or None (default)
        If 'english', a built-in stop word list for English is used.
        If a list, that list is assumed to contain stop words, all of which
        will be removed from the resulting tokens.
    lowercase: boolean, default True
        Convert all characters to lowercase before tokenizing.
    token_pattern: string
        Regular expression denoting what constitutes a "token", only used
        if `analyzer == 'word'`. The default regexp selects tokens of 2
        or more alphanumeric characters (punctuation is completely ignored
        and always treated as a token separator).
    n_features : integer, optional, (2 ** 20) by default
        The number of features (columns) in the output matrices. Small numbers
        of features are likely to cause hash collisions, but large numbers
        will cause larger coefficient dimensions in linear learners.
    norm : 'l1', 'l2' or None, optional
        Norm used to normalize term vectors. None for no normalization.
    binary: boolean, False by default.
        If True, all non zero counts are set to 1. This is useful for discrete
        probabilistic models that model binary events rather than integer
        counts.
    dtype: type, optional
        Type of the matrix returned by fit_transform() or transform().
    non_negative : boolean, optional
        Whether output matrices should contain non-negative values only;
        effectively calls abs on the matrix prior to returning it.
        When True, output values can be interpreted as frequencies.
        When False, output values will have expected value zero.

    See also
    --------
    CountVectorizer, TfidfVectorizer
    """

    def transform(self, Z):
        """Transform an ArrayRDD (or DictRDD with column 'X') containing
        sequence of documents to a document-term matrix.

        Parameters
        ----------
        Z : ArrayRDD or DictRDD with raw text documents
            Samples. Each sample must be a text document (either bytes or
            unicode strings) which will be tokenized and hashed.

        Returns
        -------
        Z : SparseRDD/DictRDD containg scipy.sparse matrix
            Document-term matrix.
        """
        mapper = super(SparkHashingVectorizer, self).transform
        return Z.transform(mapper, column='X', dtype=sp.spmatrix)

    fit_transform = transform

    def to_scikit(self):
        obj = HashingVectorizer()
        obj.__dict__ = self.__dict__
        return obj

class SparkTfidfTransformer(TfidfTransformer, SparkBroadcasterMixin):

    """Distributed implementation of TfidfTransformer.

    Transform a count matrix to a normalized tf or tf-idf representation

    Tf means term-frequency while tf-idf means term-frequency times inverse
    document-frequency. This is a common term weighting scheme in information
    retrieval, that has also found good use in document classification.

    The goal of using tf-idf instead of the raw frequencies of occurrence of a
    token in a given document is to scale down the impact of tokens that occur
    very frequently in a given corpus and that are hence empirically less
    informative than features that occur in a small fraction of the training
    corpus.

    The actual formula used for tf-idf is tf * (idf + 1) = tf + tf * idf,
    instead of tf * idf. The effect of this is that terms with zero idf, i.e.
    that occur in all documents of a training set, will not be entirely
    ignored. The formulas used to compute tf and idf depend on parameter
    settings that correspond to the SMART notation used in IR, as follows:

    Tf is "n" (natural) by default, "l" (logarithmic) when sublinear_tf=True.
    Idf is "t" when use_idf is given, "n" (none) otherwise.
    Normalization is "c" (cosine) when norm='l2', "n" (none) when norm=None.

    Parameters
    ----------
    norm : 'l1', 'l2' or None, optional
        Norm used to normalize term vectors. None for no normalization.
    use_idf : boolean, optional
        Enable inverse-document-frequency reweighting.
    smooth_idf : boolean, optional
        Smooth idf weights by adding one to document frequencies, as if an
        extra document was seen containing every term in the collection
        exactly once. Prevents zero divisions.
    sublinear_tf : boolean, optional
        Apply sublinear tf scaling, i.e. replace tf with 1 + log(tf).

    References
    ----------
    .. [Yates2011] `R. Baeza-Yates and B. Ribeiro-Neto (2011). Modern
                   Information Retrieval. Addison Wesley, pp. 68-74.`
    .. [MRS2008] `C.D. Manning, P. Raghavan and H. Schuetze  (2008).
                   Introduction to Information Retrieval. Cambridge University
                   Press, pp. 118-120.`
    """

    __transient__ = ['_idf_diag']

    def fit(self, Z):
        """Learn the idf vector (global term weights)

        Parameters
        ----------
        Z : ArrayRDD or DictRDD containing (sparse matrices|ndarray)
            a matrix of term/token counts

        Returns
        -------
        self : TfidfVectorizer
        """

        X = Z[:, 'X'] if isinstance(Z, DictRDD) else Z
        check_rdd(X, (sp.spmatrix, np.ndarray))

        def mapper(X, use_idf=self.use_idf):
            if not sp.issparse(X):
                X = sp.csc_matrix(X)
            if use_idf:
                return _document_frequency(X)

        if self.use_idf:
            n_samples, n_features = X.shape
            df = X.map(mapper).treeReduce(operator.add)

            # perform idf smoothing if required
            df += int(self.smooth_idf)
            n_samples += int(self.smooth_idf)

            # log1p instead of log makes sure terms with zero idf don't get
            # suppressed entirely
            idf = np.log(float(n_samples) / df) + 1.0
            self._idf_diag = sp.spdiags(idf,
                                        diags=0, m=n_features, n=n_features)
        return self

    def transform(self, Z):
        """Transform an ArrayRDD (or DictRDD's 'X' column) containing count
        matrices to a tf or tf-idf representation

        Parameters
        ----------

        Z : ArrayRDD/DictRDD with sparse matrices
            a matrix of term/token counts

        Returns
        -------
        Z : SparseRDD/DictRDD containing sparse matrices
        """
        X = Z[:, 'X'] if isinstance(Z, DictRDD) else Z
        check_rdd(X, (sp.spmatrix, np.ndarray))

        mapper = super(SparkTfidfTransformer, self).transform

        if self.use_idf:
            check_is_fitted(self, '_idf_diag', 'idf vector is not fitted')
            mapper = self.broadcast(mapper, Z.context)

        return Z.transform(mapper, column='X', dtype=sp.spmatrix)

    def to_scikit(self):
        obj = TfidfTransformer()
        obj.__dict__ = self.__dict__
        return obj


# -*- coding: utf-8 -*-
"""

"""
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing.data import _handle_zeros_in_scale
from splearn.base import SparkTransformerMixin
import numpy as np
import scipy.sparse as sp
from sklearn.utils.sparsefuncs import mean_variance_axis, inplace_column_scale
from sklearn.utils import check_array

from ..rdd import DictRDD
from ..utils.validation import check_rdd, check_rdd_dtype

class SparkStandardScaler(StandardScaler, SparkTransformerMixin):
    """Standardize features by removing the mean and scaling to unit variance
    Centering and scaling happen independently on each feature by computing
    the relevant statistics on the samples in the training set. Mean and
    standard deviation are then stored to be used on later data using the
    `transform` method.
    Standardization of a dataset is a common requirement for many
    machine learning estimators: they might behave badly if the
    individual feature do not more or less look like standard normally
    distributed data (e.g. Gaussian with 0 mean and unit variance).
    For instance many elements used in the objective function of
    a learning algorithm (such as the RBF kernel of Support Vector
    Machines or the L1 and L2 regularizers of linear models) assume that
    all features are centered around 0 and have variance in the same
    order. If a feature has a variance that is orders of magnitude larger
    that others, it might dominate the objective function and make the
    estimator unable to learn from other features correctly as expected.
    This scaler can also be applied to sparse CSR or CSC matrices by passing
    `with_mean=False` to avoid breaking the sparsity structure of the data.
    Read more in the :ref:`User Guide <preprocessing_scaler>`.
    Parameters
    ----------
    with_mean : boolean, True by default
        If True, center the data before scaling.
        This does not work (and will raise an exception) when attempted on
        sparse matrices, because centering them entails building a dense
        matrix which in common use cases is likely to be too large to fit in
        memory.
    with_std : boolean, True by default
        If True, scale the data to unit variance (or equivalently,
        unit standard deviation).
    copy : boolean, optional, default True
        Ignored in Spark, used when converted to scikit
        If False, try to avoid a copy and do inplace scaling instead.
        This is not guaranteed to always work inplace; e.g. if the data is
        not a NumPy array or scipy.sparse CSR matrix, a copy may still be
        returned.
    Attributes
    ----------
    scale_ : ndarray, shape (n_features,)
        Per feature relative scaling of the data.
        .. versionadded:: 0.17
           *scale_* is recommended instead of deprecated *std_*.
    mean_ : array of floats with shape [n_features]
        The mean value for each feature in the training set.
    var_ : array of floats with shape [n_features]
        The variance for each feature in the training set. Used to compute
        `scale_`
    n_samples_seen_ : int
        The number of samples processed by the estimator. Will be reset on
        new calls to fit, but increments across ``partial_fit`` calls.
    See also
    --------
    :func:`sklearn.preprocessing.scale` to perform centering and
    scaling without using the ``Transformer`` object oriented API
    :class:`sklearn.decomposition.RandomizedPCA` with `whiten=True`
    to further remove the linear correlation across features.
    """

    def fit(self, Z):
        """Compute the mean and std to be used for later scaling.
        Parameters
        ----------
        Z : DictRDD containing (X, y) pairs
            X - Training vector.
                {array-like, sparse matrix}, shape [n_samples, n_features]
                The data used to compute the mean and standard deviation
                used for later scaling along the features axis.
            y - Target labels
                Passthrough for ``Pipeline`` compatibility.
        """

        # Reset internal state before fitting
        self._reset()
        X = Z[:, 'X'] if isinstance(Z, DictRDD) else Z
        check_rdd(X, (np.ndarray, sp.spmatrix))

        def mapper(X):
            """Calculate statistics for every numpy or scipy blocks."""
            X = check_array(X, ('csr', 'csc'), dtype=np.float64)
            if hasattr(X, "toarray"):   # sparse matrix
                mean, var = mean_variance_axis(X, axis=0)
            else:
                mean, var = np.mean(X, axis=0), np.var(X, axis=0)
            return X.shape[0], mean, var

        def reducer(a, b):
            """Calculate the combined statistics."""
            n_a, mean_a, var_a = a
            n_b, mean_b, var_b = b
            n_ab = n_a + n_b
            mean_ab = ((mean_a * n_a) + (mean_b * n_b)) / n_ab
            var_ab = (((n_a * var_a) + (n_b * var_b)) / n_ab) + \
                     ((n_a * n_b) * ((mean_b - mean_a) / n_ab) ** 2)
            return (n_ab, mean_ab, var_ab)

        if check_rdd_dtype(X, (sp.spmatrix)):
            if self.with_mean:
                raise ValueError(
                    "Cannot center sparse matrices: pass `with_mean=False` "
                    "instead. See docstring for motivation and alternatives.")
        self.n_samples_seen_, self.mean_, self.var_ = X.map(mapper).treeReduce(reducer)

        if self.with_std:
            self.scale_ = _handle_zeros_in_scale(np.sqrt(self.var_))
        else:
            self.scale_ = None

        return self

    def transform(self, Z):
        """Perform standardization by centering and scaling
        Parameters
        ----------
        Z : DictRDD containing (X, y) pairs
            X - Training vector
            y - Target labels
        Returns
        -------
        C : DictRDD containing (X, y) pairs
            X - Training vector standardized
            y - Target labels
        """
        X = Z[:, 'X'] if isinstance(Z, DictRDD) else Z
        check_rdd(X, (np.ndarray, sp.spmatrix))

        if check_rdd_dtype(X, (sp.spmatrix)):
            if self.with_mean:
                raise ValueError(
                    "Cannot center sparse matrices: pass `with_mean=False` "
                    "instead. See docstring for motivation and alternatives.")
            if self.scale_ is not None:
                def mapper(X):
                    inplace_column_scale(X, 1 / self.scale_)
                    return X
        else:
            if self.with_mean:
                if self.with_std:
                    def mapper(X):
                        X -= self.mean_
                        X /= self.scale_
                        return X
                else:
                    def mapper(X):
                        X -= self.mean_
                        return X
            else:
                if self.with_std:
                    def mapper(X):
                        X /= self.scale_
                        return X
                else:
                    raise ValueError("Need with_std or with_mean ")
        return Z.transform(mapper, column="X")

    def to_scikit(self):
        """
        Convert to equivalent StandardScaler
        """
        scaler = StandardScaler(with_mean=self.with_mean,
                                with_std=self.with_std,
                                copy=self.copy)
        scaler.__dict__ = self.__dict__
        return scaler


# encoding: utf-8

import operator

import numpy as np
import scipy.sparse as sp
from sklearn.base import copy
from sklearn.linear_model.base import LinearRegression

from ..utils.validation import check_rdd

class SparkLinearModelMixin(object):

    def __add__(self, other):
        """Add method for Linear models with coef and intercept attributes.

        Parameters
        ----------
        other : fitted sklearn linear model
            Model to add.

        Returns
        -------
        model : Linear model
            Model with updated coefficients.
        """
        model = copy.deepcopy(self)
        model.coef_ += other.coef_
        model.intercept_ += other.intercept_
        return model

    def __radd__(self, other):
        """Reverse add method for Linear models.

        Parameters
        ----------
        other : fitted sklearn linear model
            Model to add.

        Returns
        -------
        model : Linear model
            Model with updated coefficients.
        """
        return self if other == 0 else self.__add__(other)

    def __div__(self, other):
        """Division method for Linear models. Used for averaging.

        Parameters
        ----------
        other : integer
            Integer to divide with.

        Returns
        -------
        model : Linear model
            Model with updated coefficients.
        """
        self.coef_ /= other
        self.intercept_ /= other
        return self

    __truediv__ = __div__

    def _spark_fit(self, cls, Z, *args, **kwargs):
        """Wraps a Scikit-learn Linear model's fit method to use with RDD
        input.

        Parameters
        ----------
        cls : class object
            The sklearn linear model's class to wrap.
        Z : TupleRDD or DictRDD
            The distributed train data in a DictRDD.

        Returns
        -------
        self: the wrapped class
        """
        mapper = lambda X_y: super(cls, self).fit(
            X_y[0], X_y[1], *args, **kwargs
        )
        models = Z.map(mapper)
        avg = models.reduce(operator.add) / models.count()
        self.__dict__.update(avg.__dict__)
        return self

    def _spark_predict(self, cls, X, *args, **kwargs):
        """Wraps a Scikit-learn Linear model's predict method to use with RDD
        input.

        Parameters
        ----------
        cls : class object
            The sklearn linear model's class to wrap.
        Z : ArrayRDD
            The distributed data to predict in a DictRDD.

        Returns
        -------
        self: the wrapped class
        """
        return X.map(lambda X: super(cls, self).predict(X, *args, **kwargs))

    def _to_scikit(self, cls):
        new = cls()
        new.__dict__ = self.__dict__
        return new

class SparkLinearRegression(LinearRegression, SparkLinearModelMixin):

    """Distributed implementation of sklearn's Linear Regression.

    Parameters
    ----------
    fit_intercept : boolean, optional
        whether to calculate the intercept for this model. If set
        to false, no intercept will be used in calculations
        (e.g. data is expected to be already centered).
    normalize : boolean, optional, default False
        If True, the regressors X will be normalized before regression.
    copy_X : boolean, optional, default True
        If True, X will be copied; else, it may be overwritten.
    n_jobs : The number of jobs to use for the computation.
        If -1 all CPUs are used. This will only provide speedup for
        n_targets > 1 and sufficient large problems.

    Attributes
    ----------
    coef_ : array, shape (n_features, ) or (n_targets, n_features)
        Estimated coefficients for the linear regression problem.
        If multiple targets are passed during the fit (y 2D), this
        is a 2D array of shape (n_targets, n_features), while if only
        one target is passed, this is a 1D array of length n_features.
    intercept_ : array
        Independent term in the linear model.

    """

    def fit(self, Z):
        """
        Fit linear model.

        Parameters
        ----------
        Z : DictRDD with (X, y) values
            X containing numpy array or sparse matrix - The training data
            y containing the target values

        Returns
        -------
        self : returns an instance of self.
        """
        check_rdd(Z, {'X': (sp.spmatrix, np.ndarray)})
        return self._spark_fit(SparkLinearRegression, Z)

    def predict(self, X):
        """Distributed method to predict class labels for samples in X.

        Parameters
        ----------
        X : ArrayRDD containing {array-like, sparse matrix}
            Samples.

        Returns
        -------
        C : ArrayRDD
            Predicted class label per sample.
        """
        check_rdd(X, (sp.spmatrix, np.ndarray))
        return self._spark_predict(SparkLinearRegression, X)

    def to_scikit(self):
        return self._to_scikit(LinearRegression)