# -*- coding: utf-8 -*-
# !/usr/bin/env python

"""

Borrowed from https://gist.github.com/tgarc/21d2ffe00dab4b3c1075

1. Add the this script below to your ipython folder (by default $HOME(%HOME%)/.ipython). This script takes care of saving user variables on exit.

2. Add this line to your profile's ipython startup script (e.g., $HOME/.ipython/profile_default/startup/startup.py):
get_ipython().ex("import ipython_save_session;del ipython_save_session")

3. In your ipython profile config file (by default $HOME/.ipython/profile_default/ipython_config.py) find the following line:
# c.StoreMagics.autorestore = False
Uncomment it and set it to true. This automatically reloads stored variables on startup. Alternatively you can use reload the last session manually using %store -r.

"""


def get_response(quest, default=None, opts=('y', 'n'), please=None, fmt=None):
    # try:
    #     input = raw_input
    # except NameError:
    #     pass

    def get_input(_): return str(raw_input(_))

    quest += " ("
    quest += "/".join(['[' + o + ']' if o == default else o for o in opts])
    quest += "): "

    if default is not None: opts = list(opts) + ['']
    if please is None: please = quest
    if fmt is None: fmt = lambda x: x

    rin = get_input(quest)
    while fmt(rin) not in opts:
        print('not in opts(%s), try again' % (opts,))
        rin = get_input(please)

    return default if default is not None and rin == '' else fmt(rin)


def get_user_vars():
    """
    Get variables in user namespace (ripped directly from ipython namespace
    magic code)
    """
    import IPython
    ip = IPython.get_ipython()
    user_ns = ip.user_ns
    user_ns_hidden = ip.user_ns_hidden
    nonmatching = object()
    var_hist = [i for i in user_ns
                if not i.startswith('_') \
                and (user_ns[i] is not user_ns_hidden.get(i, nonmatching))]
    return var_hist


def shutdown_logger():
    """
    Prompts for saving the current session during shutdown
    """
    import IPython, pickle, cPickle
    var_hist = get_user_vars()
    ip = IPython.get_ipython()
    db = ip.db

    # collect any variables that need to be deleted from db
    keys = map(lambda x: x.split('/')[1], db.keys('autorestore/*'))
    todel = set(keys).difference(ip.user_ns)
    changed = [db[k] != ip.user_ns[k.split('/')[1]]
               for k in db.keys('autorestore/*') if k.split('/')[1] in ip.user_ns]

    try:
        if len(var_hist) == 0 and len(todel) == 0 and not any(changed): return
        if get_response("Save session?", 'n', fmt=lambda _ : str(_).lower()) == 'n': return
    except KeyboardInterrupt:
        return

    # Save interactive variables (ignore unsaveable ones)
    for name in var_hist:
        obj = ip.user_ns[name]
        try:
            db['autorestore/' + name] = obj
        except (pickle.PicklingError, cPickle.PickleError):
            print("Could not store variable '%s'. Skipping..." % name)
            del db['autorestore/' + name]

    # Remove any previously stored variables that were deleted in this session
    for k in todel:
        del db['autorestore/' + k]


import atexit

atexit.register(shutdown_logger)
del atexit