"""
    Load data stored as a Package to ram. Only suitable for
    small dataset which can fit into your ram.
    To make packages suitable for this method, use "dataProcess.setPack()".
    This method can also load ram cache generated by "setfromfolder" method.
    Copyright 2019 Yang Kaiyu https://github.com/yku12cn
"""
import os
from pathlib import Path

from . import VDPlus


class setfromPack(VDPlus):
    r"""Load dataset stored as a Package to ram.

    Args:
        root (str/Path): directory of your data.
        packname (str/Path, optional): name of your pack file (*.vdcache).
        metaname (str/Path, optional): by default metaname is the same as
                                    packname but with a different extention.
        transform (torchvision.transforms, optional): for image.
        target_transform (torchvision.transforms, optional): for lable.
    """
    def __init__(self, root, packname=None, metaname=None,
                 transform=None, target_transform=None):
        super(setfromPack, self).__init__(
            root, transform=transform, tags=None,
            target_transform=target_transform
        )
        # Look for the package
        if not packname:
            _root = Path(root)
            _pack = None
            for file in os.scandir(_root):
                if file.is_file() and file.name.endswith(".vdcache"):
                    _pack = Path(file)
                    if _pack.name.startswith(_root.name):
                        break
            if not _pack:
                raise FileNotFoundError('.vdcache not found')
            packname = _pack.relative_to(_root)
        else:
            _pack = self._prePath(packname, ".vdcache")

        if not _pack.is_file():
            raise FileNotFoundError(f'"{_pack}" not found')

        # Look for set metadata
        if not metaname:
            _meta = None
            for file in os.scandir(_pack.parent):
                if file.is_file() and file.name.endswith(".vdmeta"):
                    _meta = Path(file)
                    if _pack.name.startswith(_meta.name[0:-6]):
                        break
            if not _meta:
                raise FileNotFoundError('.vdmeta not found')
        else:
            _meta = self._prePath(metaname, ".vdmeta")
            if not _meta.is_file():
                raise FileNotFoundError(f'"{_meta}" not found')

        if _pack.name[0:-7] != _meta.name[0:-6]:
            print("Warning: Meta may not match with the dataset")

        _meta = _meta.relative_to(self.root)

        print(f'Loading "{packname}" to ram')
        self.img_type, self.classes, self.classes_count = self.loadMeta(_meta)
        self.loadCache(packname)
        print("Done")
