#!/usr/bin/env python3

def _configure(repl):
    from prompt_toolkit.formatted_text import HTML, ANSI
    from ptpython.prompt_style import PromptStyle
    prompt_color = '5'
    class PowerlinePrompt(PromptStyle):
        def in_prompt(self):
            self.state = 0
            return ANSI(f"\033[4{prompt_color}m  {repl.current_statement_index} \033[0m\033[3{prompt_color}m▓▒░\033[0m ")
        def in2_prompt(self, width):
            if self.state == 0:
                self.state = 1
                return ANSI(f"\033[3{prompt_color}m\033[0m\033[4{prompt_color}m \033[0m\033[3{prompt_color}m▓▒░\033[0m ".rjust(width+27))
            else:
                return ANSI(f"\033[4{prompt_color}m \033[0m\033[3{prompt_color}m▓▒░\033[0m ".rjust(width+18))
        def out_prompt(self):
            return []
    repl.all_prompt_styles["powerline"] = PowerlinePrompt()
    repl.prompt_style = "powerline"
    #repl.show_signature = True
    repl.show_docstring = True

if __name__ == '__main__':
    # Useful imports
    import builtins, sys, os, random, time, typing, json, io, string, pickle, functools, math, operator
    from collections import Counter, defaultdict
    from contextlib import contextmanager
    from functools import lru_cache, partial, wraps
    from itertools import chain, permutations, combinations, product, repeat
    from fractions import Fraction
    from inspect import getsource
    from pathlib import Path
    from math import sqrt
    
    for _module, _name in {
        'requests': 'requests',
        'numpy': 'np',
        'pandas': 'pd',
        'sklearn': 'sklearn',
        'skimage': 'skimage',
        'scipy': 'scipy',
        'cv2': 'cv2',
        'yaml': 'yaml',
        'torch': 'torch'
        #'tensorflow': 'tf', # takes too long
    }.items():
        try:
            globals()[_name] = __import__(_module)
        except ImportError:
            pass
    del _module, _name
    
    def _plt_show(_path='/tmp/matplotlib.png'):
        plt.savefig(_path)
        os.system(f'feh {_path}')
    
    for _code in [
        'from functools import cache',
        'from scipy import stats',
        'from tqdm import tqdm, trange',
        'import matplotlib.pyplot as plt; plt.show = _plt_show',
    ]:
        try:
            exec(_code)
        except:
            pass
    del _code
    
    del _plt_show
    
    def inspect(f):
        try:
            import prompt_toolkit as pt
            import pygments as pyg
            pt.print_formatted_text(pt.formatted_text.PygmentsTokens(pyg.lex(getsource(f), lexer=pyg.lexers.python.PythonLexer())))
        except ImportError:
            print(getsource(f))
    
    class A:
        def __class_getitem__(cls, k): return np.array(list(k) if isinstance(k, tuple) else [k])
        def __new__(cls, k): return np.arange(functools.reduce(operator.mul, k)).reshape(k) if isinstance(k, tuple) else np.array(k)
    class T:
        def __class_getitem__(cls, k): return torch.tensor(list(k) if isinstance(k, tuple) else [k])
        def __new__(cls, k): return torch.arange(functools.reduce(operator.mul, k)).reshape(k) if isinstance(k, tuple) else torch.tensor(k)
    
    # CVE-2020-10735 workaround
    sys.set_int_max_str_digits(0)
    
    try:
        sys.exit(__import__('ptpython').repl.embed(globals(), locals(), configure=_configure, history_filename=Path(__file__).parent / ".python_history"))
    except ImportError:
        print("ptpython is not available: falling back to standard prompt")

# vim:expandtab:tabstop=4:shiftwidth=4
