#!/usr/bin/env python3
"""
A PyImageTool helper to copy the packaged Numpy arrays from
privileged locations to a writeable cache, based on:
/usr/share/pyimagetool/examples/example_2D.py

This is only intended to be an example for your own scripts
based on PyImageTool.
"""

# -*- coding: utf-8 -*-
#
#  debian_pyimagetool.py
#
#  Copyright 2022 Neil Williams <codehelp@debian.org>
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; version 3 of the License.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#
# On Debian systems, the complete text of the GNU General
# Public License version 3 can be found in "/usr/share/common-licenses/GPL-3".

import os
import sys
import shutil
import xdg
import numpy
from pyimagetool import imagetool
from pyimagetool.DataMatrix import RegularDataArray


class PyImageToolHelper:
    """
    Provide a cache of PyImageTool Numpy data arrays in
    a writeable location using XDG. Launching PyImageTool
    requires Qt and hence a desktop.

    The DIRNAME defined here matches what is patched into
    /usr/lib/python3/dist-packages/pyimagetool/cmaps/CMap.py
    """

    DIRNAME = "PyImageTool"
    DATA_DIR = "data"
    CMAPS_DIR = "cmaps"

    def __init__(self):
        """
        Set the writeable location

        ~/.cache/PyImageTool/data/
        """
        self.cache_dir = os.path.join(xdg.xdg_cache_home(), self.DIRNAME)
        self.pkg_dir = "/usr/lib/python3/dist-packages/pyimagetool"

        self.write_dir = os.path.join(self.cache_dir, self.DATA_DIR)

    def cached(self):
        """Assume that if the cache_dir exists, all arrays also exist"""
        return os.path.exists(self.cache_dir)

    def create_cache(self):
        """
        Make the cache_dir and copy all pyimagetool data arrays.
        Remove the copied __init__.py which is replaced by functions
        in this class.

        IMPORTANT: the cmaps/data directory needs to be copied
        onto the same directory as the main data directory, as an overlay.
        """
        if self.cached():
            return
        os.makedirs(self.cache_dir, mode=0o755)
        shutil.copytree(os.path.join(self.pkg_dir, self.DATA_DIR), self.write_dir)
        shutil.copytree(
            os.path.join(self.pkg_dir, self.CMAPS_DIR, self.DATA_DIR),
            self.write_dir,  # not a typo
            dirs_exist_ok=True,
        )
        os.unlink(os.path.join(self.write_dir, "__init__.py"))

    def refresh_cache(self):
        """
        PyImageTool modifies the data in the arrays, refresh
        returns the cache to the data arrays in the package
        """
        if not self.cached():
            self.create_cache()
            return
        shutil.rmtree(self.cache_dir)
        self.create_cache()

    def arpes_data_3d(self, delta=None, coord_min=None):
        """Mimic & extend the examples.__init__.py support for 3D data"""
        if not delta:
            delta = [0.1, 0.0621, 0.000758]
        if not coord_min:
            coord_min = [-7.0, -22.5, 20.9]
        if not self.cached():
            self.create_cache()
        dat = numpy.load(os.path.join(self.write_dir, "arpes.npy"))
        return RegularDataArray(dat, delta=delta, coord_min=coord_min)

    def arpes_data_2d(self, delta=None, coord_min=None):
        """Mimic & extend the examples.__init__.py support for 3D data"""
        if not delta:
            delta = [0.1, 0.0621]
        if not coord_min:
            coord_min = [-7.0, -22.5]
        if not self.cached():
            self.create_cache()
        dat = numpy.load(os.path.join(self.write_dir, "example_2D.npy"))
        return RegularDataArray(dat, delta=delta, coord_min=coord_min)


def main():
    """
    Basic example

    This could be extended with argparse, for example.
    """
    cache = PyImageToolHelper()

    # Forcing the refresh of the cache is just for this example.
    cache.refresh_cache()

    # other data functions would need to be created
    data = cache.arpes_data_3d()

    # now raise the GUI window
    imagetool(data)
    return 0


if __name__ == "__main__":
    sys.exit(main())