#A dependency shotgun - installs a bunch of packages

#A set of packages that are tried and true to install together as one 
harmonized_packages = [
    "2to3",
    "Cython",
    "HTMLParser",
    "IPython",
    "Imath",
    "OpenEXR",
    "PyYAML",
    "Pyperclip",
    "SentencePiece",
    "accelerate",
    "addict",
    "aiohttp",
    "albumentations",
    "apriltag",
    "art",
    "ascii-magic",
    "asciinema",
    "astor",
    "austin-dist",
    "austin-tui",
    "av",
    "beautifulsoup4",
    "black",
    "bokeh",
    "boto3",
    "bpytop",
    "bs4",
    "click",
    "cloud-detect",
    "cmapy",
    "colorama",
    "datasets",
    "decord",
    "dill",
    "dploy",
    "drawille",
    "duviz",
    "easydict",
    "editor",
    "einops",
    "embeddify",
    "exofrills",
    "fbchat",
    "ffmpeg-python",
    "filelock",
    "fire",
    "flake8",
    "frogmouth",
    "ftfy",
    "func_timeout",
    "get-mac",
    "getmac",
    "glances",
    "googletrans",
    "gpustat",
    "gradio",
    "graphviz",
    "gtts_token",
    "hf_transfer",
    "hiddenlayer",
    "html2text",
    "huggingface",
    "huggingface_hub",
    "icecream",
    "imageio",
    "imageio-ffmpeg",
    "imagesize",
    "imgcat",
    "imgurpython",
    "inflect",
    "ipykernel",
    "ipynbname",
    "ipyplot",
    "ipython",
    "isort",
    "iterfzf",
    "js2py",
    "jtree",
    "jupyter",
    "lazyasd",
    "libtmux",
    "loguru",
    "mac_vendor_lookup",
    "macchiato",
    "matplotlib",
    "mediapy",
    "memray",
    "moms_apriltag",
    "more_itertools",
    "moviepy",
    "nbstripout",
    "nltk",
    "numba",
    "numpngw",
    "nvitop",
    "oldest-supported-numpy",
    "omegaconf",
    "openai",
    "opencv-contrib-python",
    "osascript",
    "pandas",
    "patool",
    "pillow",
    "pillow-avif-plugin",
    "pillow-jxl-plugin",
    "pillow_heif",
    "plotille",
    "psutil",
    "ptpython",
    "pudb",
    "py-spy",
    "py3nvml",
    "pyarrow",
    "pydantic",
    "pydub",
    "pyflann",  # TODO: This needs to be fixed. It's broken for python3. It needs to be fixed using 2to3, and then inlined into this rp package. This is needed for flann_dict
    "pygame",
    "pygments",
    "pyinstrument",
    "pyjpgclipboard",
    "pylcs",
    "pymatbridge",
    "pynput",
    "pypager",
    "pyqt5",
    "pyqtgraph",
    "pyscreenshot",
    "pyserial",
    "pytesseract",
    "python-minifier",
    "pytube",
    "pytubefix",
    "pytz",
    "pyunpack",
    "pyyaml",
    "qrcode",
    "ranger",
    "ray",
    "removestar",
    "requests",
    "rich",
    "scenedetect",
    "scikit-image",
    "scikit-learn",
    "scikit-video",
    "scipy",
    "scrapy",
    "screeninfo",
    "seaborn",
    "send2trash",
    "sentencepiece",
    "serial",
    "setproctitle",
    "sk-video",
    "sklearn",
    "smart_open",
    "snowy",
    "sounddevice",
    "stackprinter",
    "supervision",
    "suplemon",
    "sympy",
    "tabview",
    "taichi",
    "tensorboard",
    "textual",
    "textual-imageview",
    "textual-paint",
    "textual[syntax]",
    "timestring",
    "timg",
    "tmuxp",
    "tokenizers",
    "toml",
    "tqdm",
    "transliterate",
    "twilio",
    "ufmt",  # Combines black + usort: https://github.com/omnilib/ufmt
    "urllib3",
    "urwid",
    "usort",
    "usort",  # An alternative to isort: https://libraries.io/pypi/usort
    "uv", # Like pip...but much much faster! Might use as the default for pip_install if its reliable...
    "visidata",
    "vispy",
    "voluptuous",
    "wandb",
    "watchdog",
    "webcolors",
    "win_unicode_console",
    "word2number",
    "wordcloud",
    "xdoctest",
    "xonsh",
    "yapf",
    "ydiff",
    "youtube_dl",
    "zalgo-text",
]

#Packages that might not install together nicely with packages - must try individually
unreliable_packages = [
    "playsound",
    "pyaudio",
    "rtmidi",
    "blist",
    "git",
    "pyfx",
    "urllib",
    "skvideo",
    "peepdis",
]


#Packages that depend on pytorch. Note that some computers need a different torch version! So these are separated.
torch_packages = [
    "SwissArmyTransformer",
    "bitsandbytes",
    "deepspeed",
    "diffusers",
    "flash_attn", #This can take forever, see its instructions. Sorry, you gotta do this one yourself...
    "peft",
    "safetensors",
    "tensordict",
    "timm",
    "tomesd",
    "torch",
    "torchao",
    "torchaudio",
    "torchdiffeq",
    "torchsde",
    "torchvision",
    "torchviz",
    "transformers",
    "xformers",
]


    # 'pyflann',#TODO: This needs to be fixed. It's broken for python3. It needs to be fixed using 2to3, and then inlined into this rp package. This is needed for flann_dict

def install(*extra_packages):
    """ 
    Attempts to install as many packages as possible. Skips ones that error. 
    You can add additional packages, such as torch ones with shotgun.install(shotgun.torch_packages)
    """
    import rp

    for package in packages + rp.list_flatten(rp.enlist(x) for x in extra_packages):
        bar='――――――――――――――――――――――――――――――――――――――――――――――――――――――――――――'
        print(bar)
        print("ATTEMPTING TO INSTALL PACKAGE:".center(len(bar)))
        print(package.center(len(bar)))
        print(bar)
        rp.pip_install(package, use_uv=True)

if __name__=='__main__':
    try:
        install()
    except:
        print("Shotgun errored and is now cancelled. Try it again with shotgun.shotgun()")
        pass
