# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_utils.ipynb.

# %% auto 0
__all__ = ['DecompressFunc', 'isiter', 'allinstance', 'allsametype', 'isin', 'arein', 'str_just_alpha', 'str_just_numeric',
           'strip_punc', 'str_split_get', 'str_part_pre', 'str_part_mid', 'str_part_end', 'filter_kwargs_for_func',
           'filter_kwargs_for_class', 'wrangle_kwargs_for_func', 'wrangle_kwargs_for_class', 'get_func_params',
           'copy_to_clipboard', 'is_pathlib', 'is_pathlike', 'is_path', 'str_to_path', 'path_to_str', 'to_abs_expanded',
           'sort_file_first', 'sort_directory_first', 'get_user', 'collapse_user', 'check_ext', 'drop_ext', 'is_tar',
           'is_gz', 'is_targz', 'is_tarball', 'filter_for_gz_files', 'get_gz_files_in_dir', 'decompress_tarball',
           'decompress_gunzip', 'undo_gz', 'RecursiveDecompressor', 'make_missing_dirs', 'dir_dirs',
           'decompress_directory_of_gunzipped_files', 'decompress_tarball_of_gunzipped_files', 'stream_file',
           'download_and_decompress_tarball_of_gunzipped_files', 'make_temp_file', 'urljoin', 'Slice', 'base_init_tree',
           'rich_init_tree', 'init_tree', 'base_entry_fn', 'rich_entry_fn', 'entry_fn', 'walk_dir_tree',
           'DirectoryTree', 'RichDirectory', 'Directory', 'is_matrix', 'undo_npmatrix', 'is_series', 'is_series_like',
           'is_np', 'is_device', 'is_cpu', 'is_mps', 'is_tensor', 'is_torch', 'undo_sparse', 'to_ndarray',
           'AdataExtractor', 'config_exp_logger', 'exp_log_filename', 'exp_param_filename', 'list_exps', 'gen_exp_name',
           'load_exp_params', 'save_exp_params', 'setup_exp', 'is_config_subset', 'find_exps', 'ArchiveDownloader']

# %% ../nbs/02_utils.ipynb 5
import os, inspect, string
from dataclasses import dataclass, field

import numpy as np, pandas as pd

from typing import List, Any, Optional, Callable, Union, Tuple, Iterable, Set, TypeAlias, Type, Dict

# %% ../nbs/02_utils.ipynb 7
def isiter(val: Any) -> bool:    
    return isinstance(val, Iterable)

def allinstance(vals:Any, dtype:Union[Type, TypeAlias]=Any) -> bool:
    return isiter(vals) and all(isinstance(i, dtype) for i in vals)

def allsametype(vals:Any) -> bool:
    if not isiter(vals) or len(vals) == 0: return True
    dtype = type(vals[0])
    return isiter(vals) and all(isinstance(i, dtype) for i in vals)

def isin(val:Any, vals:Iterable) -> bool:
    return val in vals

def arein(vals:Iterable, refs:Iterable) -> bool:                                             
    return isiter(vals) and isiter(refs) and all(isin(v, refs) for v in vals)

# %% ../nbs/02_utils.ipynb 8
def str_just_alpha(s: str) -> str:
    '''Filters a string for just alpha values'''
    return ''.join(list(filter(str.isalpha, s)))

def str_just_numeric(s: str) -> str:
    '''Filters a string for just numeric values'''
    return ''.join(list(filter(str.isnumeric, s)))

def strip_punc(s: str) -> str:
    return s.translate(str.maketrans('', '', string.punctuation))

def str_split_get(s: str, sep: str = '_', idx: int = 0) -> str:
    parts = s.split(sep)
    return parts[idx]    

def str_part_pre(s:str, sep: Optional[str] = ' ') -> str:
    pre, mid, end = s.partition(sep)
    return pre

def str_part_mid(s:str, sep: Optional[str] = ' ') -> str:
    pre, mid, end = s.partition(sep)
    return mid

def str_part_end(s:str, sep: Optional[str] = ' ') -> str:
    pre, mid, end = s.partition(sep)
    return end

# %% ../nbs/02_utils.ipynb 10
def filter_kwargs_for_func(fn: Callable, **kwargs:Optional[dict]):
    params = inspect.signature(fn).parameters
    return {k:v for k,v in kwargs.items() if k in params}

def filter_kwargs_for_class(cls: Callable, **kwargs:Optional[dict]):
    params = inspect.signature(cls.__init__).parameters
    return {k:v for k,v in kwargs.items() if k in params}

def wrangle_kwargs_for_func(
    fn: Callable, 
    defaults: Optional[dict]=None,
    **kwargs:Optional[dict]
) -> dict:
    # copy defaults
    params = (defaults or {}).copy()
    # update with kwargs of our function
    params.update(kwargs or {})
    # filter for only the params that other function accepts
    params = filter_kwargs_for_func(fn, **params)
    return params

def wrangle_kwargs_for_class(
    cls: Callable, 
    defaults: Optional[dict]=None,
    **kwargs:Optional[dict]
) -> dict:
    # copy defaults
    params = (defaults or {}).copy()
    # update with kwargs of our class
    params.update(kwargs or {})
    # filter for only the params that other class accepts
    params = filter_kwargs_for_class(cls, **params)
    return params

def get_func_params(
    fn: Callable, 
    drop_self: Optional[bool] = True,
    drop_before: Optional[int] = 0,
    drop_idxs: Optional[List[int]] = list(),
    drop_names: Optional[List[str]] = list(),
    drop_after: Optional[int] = None,
) -> Dict[str, Any]:
    params = inspect.signature(fn).parameters
    params = {k: v.default for k, v in params.items()}
    if drop_self and 'self' in params: 
        params.pop('self')

    params = {
        n: p for i, (n, p) in enumerate(params.items()) 
        if (
            # is before <= i < after
            (drop_before <= i or (drop_after is not None and i < drop_after))
            # i not in drop_idxs and n not in drop_names
            and (i not in drop_idxs and n not in drop_names)
        )
    }  
    return params

# %% ../nbs/02_utils.ipynb 11
def copy_to_clipboard(text: str):
    command = ""
    
    # macOS and Linux
    if os.name == 'posix':  
        command = 'echo "' + text + '" | pbcopy'
    
    # Windows
    elif os.name == 'nt':  
        command = 'echo ' + text + ' | clip'

    else:
        raise OSError("Unsupported operating system")

    os.system(command)

# %% ../nbs/02_utils.ipynb 13
import os, sys, pwd, atexit, tempfile, inspect, pathlib
import requests, tarfile, gzip, shutil
from dataclasses import dataclass, field
from tqdm.auto import tqdm

from typing import Optional, List, Union, Iterable, Tuple, Any, Callable

# %% ../nbs/02_utils.ipynb 14
from iza.static import (
    EXT_GZ, EXT_TAR, EXT_TAR_GZ,
)

from iza.types import (
    PathType, PathLike
)

# %% ../nbs/02_utils.ipynb 16
def is_pathlib(path:Any) -> bool:
    """Check if the input is a path"""
    return isinstance(path, PathType)

def is_pathlike(path:Any) -> bool:
    """Check if the input is a path"""
    return isinstance(path, (PathLike))

def is_path(path:Any, existance:Optional[bool]) -> bool:
    """Check if the input is a path"""
    exists_q = os.path.exists(os.path.expanduser(path))
    return is_pathlike(path) and (exists_q if existance else True)

def str_to_path(path:str) -> PathType:
    """Convert a string to a pathlib.Path object"""
    return pathlib.Path(path)

def path_to_str(path: PathType) -> str:
    """Convert a pathlib.Path object to a string"""
    return os.fspath(path)

def to_abs_expanded(dirname:Optional[PathLike]=None) -> PathType:
    if dirname is None:
        dirname = os.getcwd()

    dirname = os.path.expanduser(dirname) 
    dirname = os.path.abspath(dirname) 

    if not is_pathlib(dirname):
        dirname = pathlib.Path(dirname)

    return dirname

def sort_file_first(path: PathType):
    return (not path.is_file(), path.name.lower())

def sort_directory_first(path: PathType):
    return (path.is_file(), path.name.lower())


# %% ../nbs/02_utils.ipynb 18
def get_user() -> str:
    user = pwd.getpwuid(os.getuid())[0]
    return user

def collapse_user(path: str) -> str:
    _, rest = path.split(get_user())    
    return '~' + rest

# %% ../nbs/02_utils.ipynb 20
def check_ext(filename:str, extension:str) -> bool:
    if is_pathlib(filename):
        filename = path_to_str(filename)
    has_extension = extension in filename 
    splits = filename.split(extension)
    is_end_of_str = len(splits) >= 2 and splits[-1] == ''
    is_end_of_str = filename.endswith(extension)
    return has_extension and is_end_of_str

def drop_ext(filename:str, extension:Optional[str]=None) -> str:
    if is_pathlib(filename):
        filename = path_to_str(filename)
        
    file = os.path.basename(filename)
    if extension is None:
        file, *_ = file.split('.')
    else:
        file = filename.replace(extension, '')
    return os.path.join(os.path.dirname(filename), file)

# %% ../nbs/02_utils.ipynb 21
def is_tar(filename:str) -> bool:
    return check_ext(filename, EXT_TAR)

def is_gz(filename:str) -> bool:
    return check_ext(filename, EXT_GZ)

def is_targz(filename:str) -> bool:
    return check_ext(filename, EXT_TAR_GZ)

def is_tarball(filename:str) -> bool:
    return is_tar(filename) or is_targz(filename)


# %% ../nbs/02_utils.ipynb 22
def filter_for_gz_files(files:List[str]) -> List[str]:
    return list(filter(lambda f: is_gz(f), files))

def get_gz_files_in_dir(dirname:str) -> List[str]:
    all_files = []

    for (root, dirs, files) in os.walk(dirname):   
        fullpaths = [os.path.join(root, file) for file in files]
        all_files.extend(fullpaths)
    
    gz_files = filter_for_gz_files(all_files)
    return gz_files

# %% ../nbs/02_utils.ipynb 24
def decompress_tarball(filename:str, remove:bool=False) -> Tuple[str, Optional[EOFError]]:
    '''
    Returns
    -------
        dirname : str
            The name of the archive e.g. `~/Downloads/fluentbio.tar.gz` would
            yield `~/Downloads/fluentbio`


    Notes
    -----
    FluentBio has a weird gzip so it complains when it is 
        actually fine
    '''
    error = None
    decompress_dir = os.path.dirname(filename)
    ext = EXT_TAR_GZ if is_targz(filename) else EXT_TAR if is_tar(filename) else EXT_GZ

    dirname = drop_ext(filename, ext)
    try:
        with tarfile.open(filename) as tarball:
            tarball.extractall(decompress_dir)
            tarball.close()

    except EOFError as error:
        pass
    
    if os.path.isdir(dirname) and remove:
        os.remove(filename)

    return dirname, error


def decompress_gunzip(filename:str, remove:bool=False) -> Tuple[str, Optional[EOFError]]:
    '''
    Returns
    -------
        file : str
            The name of the decompressed file e.g. `~/Downloads/fluentbio.tsv.gz` would
            yield `~/Downloads/fluentbio.tsv`

    Notes
    -----
    FluentBio has a weird gzip so it complains when it is 
        actually fine
    '''
    error = None
    decompressed_file = drop_ext(filename, EXT_GZ)
    try:             
        with gzip.open(filename, 'rb') as gunzipped:
            with open(decompressed_file, 'wb') as unzipped:
                shutil.copyfileobj(gunzipped, unzipped)     
                   
    except EOFError as error:
        pass

    if os.path.isfile(decompressed_file) and remove:
        os.remove(filename)

    return decompressed_file, error

def undo_gz(filename: str, remove:bool=False) -> str:
    if is_targz(filename):
        filename, _ = decompress_tarball(filename, remove)
    elif is_gz(filename):
        filename, _ = decompress_gunzip(filename, remove)
    elif is_tar(filename):
        pass
    elif is_tarball(filename):
        filename, _ = decompress_tarball(filename, remove)
    return filename

# %% ../nbs/02_utils.ipynb 25
DecompressFunc = Callable[[str, bool], str]

@dataclass
class RecursiveDecompressor:
    dirname: PathType
    entries: Optional[List[str]] = field(default_factory=list, repr=False)
    strategy: Optional[DecompressFunc] = field(default=undo_gz, repr=False)
    remove: Optional[bool] = field(default=True, repr=False)
    progress: Optional[Any] = field(default=None, repr=False, init=True)

    total: Optional[int] = field(default=None, repr=False, init=False)
    done: Optional[int] = field(default=None, repr=False, init=False)
    task: Optional[Any] = field(default=None, repr=False, init=False)

    def __post_init__(self):
        if not is_pathlib(self.dirname): 
            self.dirname = pathlib.Path(self.dirname)
            
        if self.entries is None: 
            self.entries = os.listdir(self.dirname)

        self.done = 0
        self.total = len(self.entries)

    def execute(self):
        return self.decompress()

    def update_total(self, increment: int):
        self.total += increment
        if self.progress is not None:
            self.progress.update(self.task, total=self.total)
        else:
            self.tqdm_bar.total = self.total
            self.tqdm_bar.refresh()

    def decrease_total(self):
        self.total -= 1
        if self.progress is not None:
            self.progress.update(self.task, total=self.total)
        else:
            self.tqdm_bar.total = self.total
            self.tqdm_bar.refresh()

    def decompress(self):
        files = [self.dirname / entry for entry in self.entries]

        if self.progress is not None:
            with self.progress:
                self.task = self.progress.add_task("[cyan]Extracting...", total=self.total)
                for file in files:
                    self.recurse(file)
                    self.progress.advance(self.task)
        else:
            with tqdm(total=self.total, desc='Extracting') as self.tqdm_bar:
                for file in files:
                    self.recurse(file)
                    self.tqdm_bar.update()

    def recurse(self, current:PathLike):
        filename = path_to_str(current)
        nextfile = self.strategy(filename, self.remove)
        self.done += 1
        nextfile = pathlib.Path(nextfile)
        
        # If the decompressed file is a directory, recursively decompress it
        if nextfile.is_dir():
            for root, dirs, files in os.walk(nextfile):
                for f in files:
                    self.update_total(1)
                    self.recurse(pathlib.Path(root) / f)
                    
        elif self.remove:
            if filename in self.entries:
                self.entries.remove(filename)
            self.decrease_total()


# %% ../nbs/02_utils.ipynb 27
def make_missing_dirs(dirs:List[str]):
    if isinstance(dirs, str):
        dirs = [dirs]
        
    elif is_pathlib(dirs):
        dirs = [dirs]
        
    for d in dirs:
        if not os.path.exists(d):
            os.makedirs(d)
            
def dir_dirs(dirname:str) -> List[str]:
    entries = os.listdir(dirname)
    is_subdir = lambda e : os.path.isdir(os.path.join(dirname, e))
    return list(filter(is_subdir, entries))

# %% ../nbs/02_utils.ipynb 28
def decompress_directory_of_gunzipped_files(
    dirname:str, desc:Optional[str]=None, remove:Optional[bool]=False
) -> None:
    if desc is None:
        desc = dirname.split('/')[-1]

    gz_files = get_gz_files_in_dir(dirname)
    for filename in tqdm(gz_files, desc=desc):
        decomp_filename, error = decompress_gunzip(filename, remove)   


def decompress_tarball_of_gunzipped_files(
    filename:str, desc:Optional[str]=None, remove:Optional[bool]=False
) -> None:
    # NOTE: initial decompress of .tar.gz
    dirname, error = decompress_tarball(filename)

    if desc is None:
        desc = dirname.split('/')[-1]

    # NOTE: decompress all internal .gz files
    decompress_directory_of_gunzipped_files(dirname, desc, remove)

# %% ../nbs/02_utils.ipynb 30
def stream_file(uri:str, filename:Optional[str]=None, desc:Optional[str]=None) -> None:
    '''
    Parameters
    ----------
    uri : str
        The URI to download

    filename : str, optional
        The fullpath name of the file to download. Defaults to 
        `~/Downloads/os.path.basename(uri)`.

    desc : str, optional
        The description of the `tqdm` progress bar. Defaults to 
        `os.path.basename(uri)`.
    '''
    if filename is None:
        download_dir = os.path.expanduser(f'~/Downloads')        
        filename = os.path.join(download_dir, os.path.basename(uri))

    basename = os.path.basename(filename)
    if desc is None:
        desc = basename

    response = requests.get(uri, stream=True)
    total = int(response.headers.get('content-length', 0))
    
    with tqdm.wrapattr(
        open(filename, 'wb'), 'write', 
        miniters=1, desc=desc, total=total
    ) as fout:
        for chunk in response.iter_content(chunk_size=4096):
            fout.write(chunk)

# %% ../nbs/02_utils.ipynb 32
def download_and_decompress_tarball_of_gunzipped_files(
    uri:str, download_dir:str=None, desc:Optional[str]=None, remove:Optional[bool]=False
):
    filename = os.path.basename(uri)
    fullpath = os.path.join(download_dir, filename)

    if desc is None:
        description = f'Downloading {filename}'


    # NOTE: Amazon --> filtered_matrix.tar.gz
    stream_file(uri, description)

    if desc is None:
        description = f'Decompressing {filename}'
    # NOTE: filtered_matrix.tar.gz --> filtered_matrix/**/file.tsv
    decompress_tarball_of_gunzipped_files(fullpath, desc, remove)


# %% ../nbs/02_utils.ipynb 34
def make_temp_file(**kwargs: Any) -> tempfile.NamedTemporaryFile:
    temp = tempfile.NamedTemporaryFile(**kwargs)
    @atexit.register
    def delete_temp() -> None:
        temp.close()
    return temp

# %% ../nbs/02_utils.ipynb 36
from urllib3.util.url import parse_url

# %% ../nbs/02_utils.ipynb 37
def urljoin(*parts: str) -> str:
    return parse_url('/'.join(s.strip('/') for s in parts)).url

# %% ../nbs/02_utils.ipynb 39
from dataclasses import dataclass, field
import numpy as np, pandas as pd

from typing import List, Union, Tuple

# %% ../nbs/02_utils.ipynb 42
@dataclass
class Slice:
    """A class for representing a slice and providing conversion to other formats."""
    slc: slice = field(default_factory=slice)

    @property
    def start(self):
        try:
            return self._start
        except AttributeError:
            self._start = self.slc.start
        return self._start
    
    @start.setter
    def start(self, value):
        """Sets the start index."""
        if value < 0:
            raise ValueError("Slice indices must be non-negative.")
        self._start = value

    @property
    def stop(self):
        try:
            return self._stop
        except AttributeError:
            self._stop = self.slc.stop
        return self._stop
    
    @stop.setter
    def stop(self, value):
        """Sets the stop index."""
        if value < 0:
            raise ValueError("Slice indices must be non-negative.")
        self._stop = value
    
    @property
    def step(self):
        """Gets the step index."""
        try:
            return self._step
        except AttributeError:
            self._step = self.slc.step
        return self._step
    
    @step.setter
    def step(self, value):
        """Sets the step index."""
        if value < 0:
            raise ValueError("Slice step must be non-negative.")
        self._step = value
    
    def __post_init__(self):
        if self.start is None:
            self.start = 0

        if self.stop is None:
            self.stop = min(0, self.start, max(1, self.start))

        if self.step is None:
            self.step = 1

    def totuple(self) -> Tuple[Union[int, float, None], Union[int, float, None], Union[int, float, None]]:
        """Converts the slice to a tuple."""
        return (self.start, self.stop, self.step)
    
    def toslice(self) -> slice:
        """Converts the updated slice."""
        return slice(self.start, self.stop, self.step)
    
    def tolist(self) -> List[Union[int, float]]:
        """Converts the slice to a list."""
        return list(range(self.start, self.stop, self.step))
    
    def todict(self) -> List[Union[int, float]]:
        """Converts the slice to a dict."""
        return dict(zip('start stop step'.split(), self.totuple()))
        

    def astype(self, dtype:str):
        """Converts the slice to a specified format."""
        if dtype in {'list', list}:
            return self.tolist()
        elif dtype in {'numpy', np.ndarray}:
            return np.array(self.tolist())
        elif dtype in {'pandas', pd.Series}:
            return pd.Series(self.tolist())
        elif dtype in {'tuple', tuple}:
            return self.totuple()
        elif dtype in {'dict', dict}:
            return self.todict()
        elif dtype in {'slice', slice}:
            return self.toslice()
        return self

# %% ../nbs/02_utils.ipynb 44
import os, pathlib 
from pathlib import Path
from dataclasses import dataclass, field, KW_ONLY
from typing import Optional, List, ClassVar, Any, TypeAlias, Callable
from enum import StrEnum

# %% ../nbs/02_utils.ipynb 45
from iza.types import (
    PathLike, PathType,
    DirectoryTreeStrings, TreeEntryFunc,
    RichTree, RichText, RichConsole, RichProgress
)
from .static import EXT_PY

#| export
from .imp import RichImp
from ipos.imp import is_mod, is_var_imp, Module, Imp

# %% ../nbs/02_utils.ipynb 46
try:    
    from rich.tree import Tree
    from rich.text import Text    
    from rich.filesize import decimal
    from rich import get_console
    from rich.console import Console
    from rich.progress import Progress

except ImportError:
    Tree = None
    Text = None 
    Progress = None

# %% ../nbs/02_utils.ipynb 50
def base_init_tree(dirname: str):
    return None

def rich_init_tree(dirname):
    return Tree(f'[link file://{dirname}]{dirname}', guide_style='bold bright_blue')

def init_tree(dirname):
    try:
        return rich_init_tree(dirname)
    except:
        return base_init_tree(dirname)

# %% ../nbs/02_utils.ipynb 52
def base_entry_fn(path:Path, prefix:str='', pointer:str='', suffix:str='') -> str:
    name = path.name
    return f'{prefix}{pointer}{name}{suffix}'

def rich_entry_fn(path:Path, prefix:str='', pointer:str='', suffix:str='') -> str:
    text = Text(path.name)
    size = decimal(path.stat().st_size)
    text.stylize(f'link file://{path}')
    text.append(f' ({size})', 'cyan')
    return text

def entry_fn(path:Path, prefix:str='', pointer:str='', suffix:str='') -> str:
    try:
        return rich_entry_fn(path, prefix, pointer, suffix)
    except:
        return base_entry_fn(path, prefix, pointer, suffix)


# %% ../nbs/02_utils.ipynb 54
def walk_dir_tree(
    dirname: PathLike, 
    prefix: str = '',
    hidden: Optional[bool] = False, 
    tree: Optional[RichTree] = None,
    entry_fn: TreeEntryFunc = entry_fn,
):    
    '''
    A recursive generator, given a directory Path object
    will yield a visual tree structure line by line
    with each line prefixed by the same characters
    Notes
    -----
    Adapted from https://stackoverflow.com/a/59109706/5623899
    '''
    # NOTE: sort_directory_first, to_abs_expanded, rich_file defined in _02_utils/_08_modules.ipynb
    dirname = Path(to_abs_expanded(dirname))
    contents = sorted(Path(dirname).iterdir(), key=sort_directory_first)
    
    SPACE, BRANCH, TEE, LAST = DirectoryTreeStrings
    pointers = [TEE] * (len(contents) - 1) + [LAST]
    
    for pointer, path in zip(pointers, contents):
        # Remove hidden files
        if path.name.startswith('.') and not hidden:
            continue
        
        name = entry_fn(path, prefix, pointer, suffix='')
        yield name
        
        branch = None
        if tree is not None:
            branch = tree.add(entry_fn(path))
    
        if path.is_dir():
            # NOTE: space because last, └── , above so no more |
            extension = BRANCH if pointer == TEE else SPACE
            yield from walk_dir_tree(path, prefix=f'{prefix}{extension}', hidden=hidden, tree=branch, entry_fn=entry_fn)

# %% ../nbs/02_utils.ipynb 56
@dataclass
class DirectoryTree:
    dirname: str
    
    hidden: Optional[bool] = False
    entry_fn: Optional[TreeEntryFunc] = field(default=entry_fn)
    tree: Optional[RichTree] = field(init=False, default=None)

    coerce_str: ClassVar[Optional[bool]] = field(init=False, default=True, repr=False)
    
    def __post_init__(self):
        # NOTE: defined in _02_utils/_01_files.ipynb
        self.dirname = Path(to_abs_expanded(self.dirname))
        if self.tree is None:
            self.tree = init_tree(self.dirname)

    def tree_generator(self):
        yield from walk_dir_tree(
            self.dirname, hidden=self.hidden, 
            tree=self.tree, entry_fn=self.entry_fn
        )
        
    def get_tree_lines(self) -> List[str]:        
        tree_gen = self.tree_generator()
        lines = [
            str(line) if self.coerce_str else line 
            for line in tree_gen
        ]
        return lines
    
    def make_tree_str(self) -> str:        
        lines = self.get_tree_lines()
        tree_str = '\n'.join([str(self.dirname), *lines])
        return tree_str
   
    def print(self) -> None:        
        tree_str = self.make_tree_str()
        print(tree_str)
        return

    def __repr__(self):        
        tree_str = self.make_tree_str()
        return tree_str  

# %% ../nbs/02_utils.ipynb 57
@dataclass
class RichDirectory(DirectoryTree):
    _: KW_ONLY
    console: Optional[RichConsole] = field(default_factory=Console, init=True, repr=False)    
    _imp: Imp = field(default_factory=RichImp, init=False, repr=False)
    
    coerce_str: ClassVar[Optional[bool]] = field(init=False, default=False, repr=False)

    def __post_init__(self):
        super().__post_init__()
        self._imp = RichImp()
    
        if self.console is None and is_mod(self._imp._module):
            self.console = get_console()

    def print_rich(self) -> None:    
        lines = self.get_tree_lines(_coerce_str=False)
        self.console.print(self.tree)

# %% ../nbs/02_utils.ipynb 58
@dataclass
class Directory(DirectoryTree):
    dirname: PathLike
    pass

# %% ../nbs/02_utils.ipynb 60
import os
from pathlib import Path
from dataclasses import dataclass, field, KW_ONLY
from typing import Optional, List, ClassVar

# %% ../nbs/02_utils.ipynb 62
import pandas as pd

# %% ../nbs/02_utils.ipynb 63
from iza.static import (
    ADATA, MATRIX, BARCODES, FEATURES, EXT_H5, EXT_MTX, EXT_TSV,
    GENE_SYMBOL, ENSEMBL_ID
)
from iza.types import (
    AnnData
)
from tqdm.auto import tqdm

# %% ../nbs/02_utils.ipynb 64
try: 
    import scanpy as sc, scprep

    # NOTE: Directory defined in _02_utils/_02_directory.ipynb
    @dataclass
    class FilterMatrixDirectory(Directory):
        _: KW_ONLY
        ADATA_FILE: ClassVar[str] = f'{ADATA}{EXT_H5}'
        MATRIX_FILE: ClassVar[str] = f'{MATRIX}{EXT_MTX}'
        BARCODES_FILE: ClassVar[str] = f'{BARCODES}{EXT_TSV}'
        FEATURES_FILE: ClassVar[str] = f'{FEATURES}{EXT_TSV}'
        

        def __post_init__(self):    
            try:
                if not self.has_adata:
                    self.make_adata()
            except Exception as err:
                raise err

        def __repr__(self):
            base = os.path.basename(self.dirname)
            srep = f'FilteredMatrix(valid: {self.is_valid}, adata: {self.has_adata})'        
            srep += '\n'
            srep += super(FilterMatrixDirectory, self).__repr__()
            return srep
                    
        @property
        def adata_filename(self) -> str:
            return os.path.join(self.dirname, self.ADATA_FILE)

        @property
        def matrix_filename(self) -> str:
            return os.path.join(self.dirname, self.MATRIX_FILE)
        
        @property
        def barcodes_filename(self) -> str:
            return os.path.join(self.dirname, self.BARCODES_FILE)
        
        @property
        def features_filename(self) -> str:
            return os.path.join(self.dirname, self.FEATURES_FILE)

        @property
        def has_adata(self) -> bool:
            return os.path.isfile(self.adata_filename)

        @property
        def has_matrix(self) -> bool:
            return os.path.isfile(self.matrix_filename)

        @property
        def has_barcodes(self) -> bool:
            return os.path.isfile(self.barcodes_filename)

        @property
        def has_features(self) -> bool:
            return os.path.isfile(self.features_filename)

        @property
        def is_valid(self) -> bool:
            return all([self.has_matrix, self.has_barcodes, self.has_features])

        def make_adata(self) -> AnnData:
            if self.has_adata:
                return

            steps = (FEATURES, BARCODES, MATRIX, 'combine', ADATA)
            
            desc = os.path.basename(self.dirname)

            steps = tqdm(steps, desc=desc, leave=True)        
            for step in steps:
                steps.set_postfix(stage=step)
                match step:
                    case 'features':
                        features = pd.read_csv(self.features_filename, sep='\t', header=None)
                        features.columns = [ENSEMBL_ID, GENE_SYMBOL, 'feature_type']
                        features.index = pd.Series(features.ensembl_id.copy().values)

                    case 'barcodes':
                        barcodes = pd.read_csv(self.barcodes_filename, sep='\t', header=None)
                        barcodes.columns = [BARCODES]
                        barcodes.index = pd.Series(barcodes.barcodes.copy().values)

                    case 'matrix':
                        matrix = scprep.io.load_mtx(self.matrix_filename, sparse=True).T

                    case 'combine':
                        data = pd.DataFrame.sparse.from_spmatrix(
                            matrix, columns=features.index, index = barcodes.index
                        )
                        del matrix

                    case 'adata':
                        adata = sc.AnnData(X=data.values, obs=barcodes, var=features, dtype='float32')
                        adata.write(self.adata_filename)
                        
                    case _:
                        pass

            return adata

        def get_adata(self) -> AnnData:
            adata = sc.read_h5ad(self.adata_filename)
            return adata

except ImportError as err:
    @dataclass
    class FilterMatrixDirectory(Directory):
        _: KW_ONLY
        ADATA_FILE: ClassVar[str] = f'{ADATA}{EXT_H5}'
        MATRIX_FILE: ClassVar[str] = f'{MATRIX}{EXT_MTX}'
        BARCODES_FILE: ClassVar[str] = f'{BARCODES}{EXT_TSV}'
        FEATURES_FILE: ClassVar[str] = f'{FEATURES}{EXT_TSV}'
        

        def __post_init__(self):    
            raise ImportError('FilterMatrixDirectory requires scprep and scanpy to be installed')
        
    pass

# %% ../nbs/02_utils.ipynb 66
import os
import numpy as np, pandas as pd
from typing import Optional, List, ClassVar, Any

# %% ../nbs/02_utils.ipynb 67
from .types import Tensor, Device, SeriesLike, ndarray

# %% ../nbs/02_utils.ipynb 70
def is_matrix(arr: SeriesLike) -> bool:
    '''
    Checks whether or not `arr` is a np.matrix
    
    Parameters
    ----------    
    arr : SeriesLike
        object to check whether or not it is a np.matrix.
    
    Returns
    -------
    result : bool
    '''
    return isinstance(arr, np.matrix)

def undo_npmatrix(arr: SeriesLike) -> SeriesLike:  
    '''
    Given a tensor converts it to a numpy array
    
    Parameters
    ----------    
    tensor : Tensor
        
    
    Returns
    -------
    arr : ndarray
    
    Notes
    -----
    - several graphtool functions use dependencies which rely on
        the deprecated numpy class `np.matrix`.
        
    - these functions appear to be related to scipy sparse linalg
        methods.
    '''
    if is_matrix(arr):
        return np.array(arr)
    return arr

def is_series(arr: SeriesLike) -> bool:
    '''
    Checks whether or not `arr` is a pd.Series
    
    Parameters
    ----------    
    arr : SeriesLike
        object to check whether or not it is a pd.Series.
    
    Returns
    -------
    result : bool
    '''
    return isinstance(arr, pd.Series)

def is_series_like(series_q: SeriesLike) -> bool:
    '''
    Checks whether or not `series_q` is SeriesLike
    i.e. something that is probably data.
    
    Parameters
    ----------    
    series_q : SeriesLike
        object to check whether or not it is a SeriesLike.
    
    Returns
    -------
    result : bool
    '''
    return isinstance(series_q, SeriesLike)

def is_np(arr_q: SeriesLike) -> bool:
    '''
    Checks whether or not `arr_q` is a ndarray
    
    Parameters
    ----------    
    arr_q : SeriesLike
        object to check whether or not it is a SeriesLike.
    
    Returns
    -------
    result : bool
    '''
    return isinstance(arr_q, ndarray)


# %% ../nbs/02_utils.ipynb 71
def is_device(device_q: Device) -> bool:
    '''
    Checks whether or not `device_q` is a valid
    pytorch device type.
    
    Parameters
    ----------    
    device_q : Device
        object to check whether or not it is a pytorch device.
    
    Returns
    -------
    result : bool        
    
    Notes
    -----
    - There is an execption. `NoneType` is a valid
        pytorch device. Here we return `False` instead.
    '''
    if device_q is None:
        return False
    return isinstance(device_q, Device)

def is_cpu(tensor: Tensor) -> bool:
    '''
    Checks whether or not `tensor` is on cpu
    
    Parameters
    ----------    
    tensor : Tensor
        object to check whether or not it is on cpu.
    
    Returns
    -------
    result : bool
    '''
    # assert is_tensor(tensor)
    if not hasattr(tensor, 'device'):
        return True
    return tensor.device.type == 'cpu'

def is_mps(tensor: Tensor) -> bool:
    '''
    Checks whether or not `tensor` is on cpu
    
    Parameters
    ----------    
    tensor : Tensor
        object to check whether or not it is on cpu.
    
    Returns
    -------
    result : bool
    '''
    # assert is_tensor(tensor)
    if not hasattr(tensor, 'device'):
        return False
    return tensor.device.type == 'mps'

def is_tensor(tensor_q: SeriesLike) -> bool:
    '''
    Checks whether or not `tensor_q` is a pytorch tensor
    
    Parameters
    ----------    
    tensor_q : Tensor
        object to check whether or not it is a pytorch tensor
    
    Returns
    -------
    result : bool
    '''
    return isinstance(tensor_q, Tensor)


def is_torch(tensor_q: SeriesLike) -> bool:
    '''
    Alias for `is_tensor`.
    
    Parameters
    ----------    
    tensor_q : Tensor
        object to check whether or not it is a pytorch tensor
    
    Returns
    -------
    result : bool
    
    See Also
    --------
    is_tensor
    '''
    return is_tensor(tensor_q)

# %% ../nbs/02_utils.ipynb 72
def undo_sparse(arr: SeriesLike) -> SeriesLike:  
    '''
    Given a arr tries to make it a dense array
    
    Parameters
    ----------    
    arr : SeriesLike
        
    
    Returns
    -------
    arr : ndarray
    '''
    if hasattr(arr, 'toarray'):
        arr = arr.toarray()

    if hasattr(arr, 'todense'):
        arr = arr.todense()

    return arr

def to_ndarray(arr):
    if is_np(arr):
        return arr
        
    arr = undo_npmatrix(arr)
    
    arr = undo_sparse(arr)

    if not is_np(arr):
        arr = np.array(arr)

    return arr

# %% ../nbs/02_utils.ipynb 74
import os, random
import numpy as np

from dataclasses import dataclass, field, KW_ONLY
from typing import Optional, List, ClassVar, Any

# %% ../nbs/02_utils.ipynb 75
from .types import Tensor, Device, SeriesLike, ndarray

# %% ../nbs/02_utils.ipynb 77
try:
    import torch
    def ensure_device(device: Device) -> Device:
        '''
        Given a valid device type attempts to instantiant 
        a pytorch device object i.e. `device='cpu'` will
        return `torch.device('cpu')`.
        
        Parameters
        ----------    
        device : Device
            a valid pytorch device type, possible a string.
        
        Returns
        -------
        device : torch.device        
        
        Raises
        ------
        RuntimeError
            same error if `torch.device(device)` fails
        '''
        if device is None:
            return device    
        try:
            return torch.device(device)
        except RuntimeError as err:
            raise err
        return device
    
    def to_cuda(tensor: Tensor) -> Tensor:
        '''
        Given a tensor, ensures that it is on cuda.
        
        Parameters
        ----------    
        tensor : Tensor
            
        
        Returns
        -------
        tensor : Tensor
        '''
        return tensor.cuda()

    def to_mps(tensor: Tensor) -> Tensor:
        '''
        Given a tensor, ensures that it is on mac silicon.
        
        Parameters
        ----------    
        tensor : Tensor
            
        
        Returns
        -------
        tensor : Tensor
        '''
        return tensor.to(torch.device('mps'))
    
    
    def to_torch(
        arr: SeriesLike,
        cuda: Optional[bool] = False,
        mps: Optional[bool] = False,
        device: Optional[Device] = None,
        dtype: Optional[Any] = None
    ) -> Tensor:
        '''
        Given data, ensures that it is a pytorch Tensor.
        
        Parameters
        ----------    
        arr : SeriesLike
        
        cuda : bool, default=False
            whether to return the tensor on cuda
            
        mps : bool, default=False
            whether to return the tensor on mps
            
        device : Device, optional
            whether to return the tensor on given device
            
        Returns
        -------
        tensor : Tensor
            the input array as a pytorch tensor
            
        Notes
        -----
        - `device` takes priority over `cuda` and `mps`
        '''
        tensor = torch.as_tensor(arr)
        if device is not None:
            tensor = tensor.to(device)
        elif cuda:
            tensor = to_cuda(tensor)
        elif mps:
            tensor = to_mps(tensor)    
        
        if dtype is not None:
            dtype = coerce_mps_dtype(dtype, tensor.device, assume_on_mps=False)
            tensor = tensor.to(dtype)

        return tensor
    
    #| export
    def to_np(tensor:Tensor) -> ndarray:
        '''
        Given a tensor converts it to a numpy array
        
        Parameters
        ----------    
        tensor : Tensor
            
        
        Returns
        -------
        arr : ndarray
        '''
        assert is_tensor(tensor)
        if not hasattr(tensor, 'detach'):
            try:
                return np.array(tensor)
            except Exception as err:
                raise err
        try:
            return tensor.detach().clone().cpu().numpy()
        except Exception as err:
            raise err
    

    def is_mps_available() -> bool:
        '''
        Checks whether or not pytorch has mps availble (version) and was built with mps in mind.

        Returns
        -------
        result : bool
        '''
        maybe_mps = torch.backends.mps.is_available()
        built_mps = torch.backends.mps.is_built()
        return maybe_mps and built_mps
    
    def coerce_mps_dtype(
        dtype, 
        device: Optional[Device] = None, 
        assume_on_mps: Optional[bool] = True
    ):
        '''
        Makes sure `tensor` is `torch.float32` if `tensor.dtype` is `torch.float64`
        if `tensor.device` is `'mps'`.
        
        Parameters
        ----------    
        dtype : any
            dtype to check against
        
        device : Device, default=None
            the device of the tensor or model from which the `dtype` comes from. If provided
            will be used to detemine whether or not to make `torch.float64`, `torch.float32`
            only if the device is actually `'mps'`.

        assume_on_mps: bool, default=True
            whether or not to assume that the device of choice is `'mps'`. Setting this to
            `True` will result in `dtype` of `torch.float64` being converted to `torch.float32`
            to try and silently fix mps errors

        Returns
        -------
        dtype : any
            the dtype, corrected for mps if needed
        '''
        could_be_mps = is_mps_available()
        
        is_float64 = dtype == torch.float64

        if device is not None:
            is_device_mps = device.type == 'mps'
            if is_device_mps:
                assume_on_mps = True

            elif device.type == 'cuda':
                assume_on_mps = False

        
        # NOTE: float64 not availble on mps, coerce to float32
        # NOTE: could_be_mps and assume_on_mps both needed as
        #       device might not be provided.
        if could_be_mps and assume_on_mps and is_float64:
            return torch.float32
        
        return dtype
    
    def ensure_mps_dtype(tensor: Tensor) -> Tensor:
        '''
        Makes sure `tensor` is `torch.float32` if `tensor.dtype` is `torch.float64`
        if `tensor.device` is `'mps'`.
        
        Parameters
        ----------    
        tensor : Tensor
            pytorch tensor to maybe change dtype of
        
        Returns
        -------
        tensor : Tensor
        '''
        dtype = tensor.dtype

        # NOTE: we don't assume mps as we explicitly pass the device
        dtype = coerce_mps_dtype(dtype, tensor.device, assume_on_mps=False)

        tensor = tensor.to(dtype)
        return tensor

    def move_to(
        tensor: Tensor, other: Tensor, 
        dtype: Optional[Any] = None, do_dtype: Optional[bool] = True
    ) -> Tensor:
        '''
        Makes sure `tensor` is on the same device as `other`
        
        Parameters
        ----------    
        tensor : Tensor
            pytorch tensor to change device of
            
        other : Tensor
            pytorch tensor we want `tensor` to be on
            
        dtype : optional
            the data type to make `tensor`. If `None` will infer it
            from `other`
            
        do_dype: bool, default=True
            whether or not to just match the device of `other` or also
            the dtype
        
        Returns
        -------
        tensor : Tensor
        '''
        
        if not is_tensor(tensor):
            tensor = to_torch(tensor)
            
        # NOTE: dtype not provided, so we will infer it
        if dtype is None:
            # NOTE: this little line solves mps float64 issues since 
            #       infer our tensor types and move them accordingly
            other = ensure_mps_dtype(other)
            dtype = other.dtype

        if do_dtype:
            tensor = tensor.to(dtype)

        tensor = tensor.to(other.device)
        
        return tensor

except ImportError as err:
    identity = lambda x: x
    ensure_device = identity
    to_cuda = identity
    to_mps = identity
    to_torch = lambda arr, cuda, mps, device, dtype: arr
    to_np = identity
    is_mps_available = lambda: False
    coerce_mps_dtype = lambda dtype, device, assume_on_mps: dtype
    ensure_mps_dtype = identity
    move_to = lambda tensor, other, dtype, do_dtype: tensor
    pass



# %% ../nbs/02_utils.ipynb 78
try:
    import torch, pytorch_lightning as pl
    def set_seeds(seed: int) -> None:
        '''
        Calls a bunch of seed functions with `seed`
        
        Parameters
        ----------
        seed : int
        '''    
        torch.manual_seed(seed)
        random.seed(seed)
        np.random.seed(seed)    
        pl.seed_everything(seed)
except ImportError as err:
     def set_seeds(seed: int) -> None:
         random.seed(seed)
         np.random.seed(seed)

# %% ../nbs/02_utils.ipynb 80
from dataclasses import dataclass, field
import numpy as np, pandas as pd

from typing import List, Any, Optional

# %% ../nbs/02_utils.ipynb 81
from .types import AnnData, ndarray, DataFrame
from .static import X_MAGIC, PHATE, X_PHATE

# %% ../nbs/02_utils.ipynb 83
@dataclass
class AdataExtractor:
    adata: AnnData
    layer: Optional[str] = X_MAGIC
    x_emb: Optional[str] = X_PHATE

    dim_str: Optional[str] = None
    use_hvg: Optional[bool] = True

    @property
    def has_hvg(self):
        return hasattr(self.adata, 'var') and hasattr(self.adata.var, 'highly_variable')
    
    @property
    def has_emb(self):
        return hasattr(self.adata, 'obsm') and self.x_emb in self.adata.obsm.keys()

    def get_layer(self) -> ndarray:
        layer = self.sdata().layers.get(self.layer, None)

        if layer is None:
            layer = self.sdata().X

        if hasattr(layer, 'toarray'):
            layer = layer.toarray()

        if hasattr(layer, 'todense'):
            layer = layer.todense()

        return layer
    
    def get_emb(self) -> ndarray:
        emb = self.sdata().obsm.get(self.x_emb, None)
        if emb is None:
            raise ValueError(f'No embedding found in adata.obsm {self.sdata().obsm.keys()}')

        # NOTE: defined in _02_utils/_05_guards.ipynb
        emb = to_ndarray(emb)
        return emb

    @property
    def axis_str(self):
        if self.dim_str:
            return self.dim_str
        return self.x_emb.replace('X_', '').upper()
    
    @property
    def emb_cols(self):
        ndim = self.get_emb().shape[1]
        cols = [f'{self.axis_str}_{i+1}' for i in range(ndim)]
        return cols
        
    def sdata(self):
        if self.use_hvg and self.has_hvg:
            return self.adata[:, self.adata.var.highly_variable]
        return self.adata
    
    def get_df_cnt(self) -> DataFrame:
        layer = self.get_layer()

        cols = self.sdata().var.index
        idxs = self.sdata().obs.index
        df = pd.DataFrame(layer, index=idxs, columns=cols)
        return df
    
    def get_df_emb(self) -> DataFrame:
        emb = self.get_emb()
        
        cols = self.emb_cols
        idxs = self.sdata().obs.index
        df = pd.DataFrame(emb, index=idxs, columns=cols)
        return df
    
    @property
    def df_cnt(self):
        return self.get_df_cnt()
    
    @property
    def df_emb(self):
        return self.get_df_emb()

# %% ../nbs/02_utils.ipynb 85
import os, yaml, datetime, logging

# %% ../nbs/02_utils.ipynb 86
def config_exp_logger(path):
    '''
    Arguments:
    ----------
        path (str): full path to experiment, i.e. 
            `<path-to-experiments-dir>/<experiment-timestamp>`
    Returns:
    ----------
        logger
    '''
    basename = os.path.basename(path)
    logger = logging.getLogger(basename)

    logging.basicConfig(
        filename=exp_log_filename(path), 
        encoding='utf-8',
        level=logging.DEBUG,
        format='%(asctime)s\t%(levelname)s:%(message)s',
        datefmt='%d/%m/%Y %I:%M:%S %p',
        filemode='w'
    )
    logger.info(f'Experiment path created {path}')
    return logger

def exp_log_filename(path):
    '''
    Arguments:
    ----------
        path (str): full path to experiment, i.e. 
            `<path-to-experiments-dir>/<experiment-timestamp>`
    Returns:
    ----------
        log_file (str): full path to provided experiment's log file
    '''
    return os.path.join(path, 'log.txt')

def exp_param_filename(path):
    '''
    Arguments:
    ----------
        path (str): full path to experiment, i.e. 
            `<path-to-experiments-dir>/<experiment-timestamp>`
    Returns:
    ----------
        param_file (str): full path to provided experiment's parameter file
    '''
    return os.path.join(path, 'params.yml')

def list_exps(path):
    '''
    Notes:
    ----------
        - an experiment is defined as a directory containing a `'params.yml'` file.
    Arguments:
    ----------
        path (str): full path to experiments directory, i.e. 
            `<path-to-experiments-dir>`
    Returns:
    ----------
        experiments (str[]): experiments (subdirectories) in the specified 
            `path`.
    '''
    test_fn = lambda el: os.path.isdir(el) and os.path.isfile(exp_param_filename(el))
    return list(filter(test_fn, os.listdir(path)))

def gen_exp_name(name=None):
    '''    
    Returns:
    ----------
        exp_name (str): timestamp to serve as experiment name
    '''
    if name is None:
        now =  datetime.datetime.now()
        return now.strftime("%Y_%m_%d-%I_%M_%S_%p")
    return name
        
def load_exp_params(path):
    '''
    Arguments:
    ----------
        path (str): full path to experiment, i.e. 
            `<path-to-experiments-dir>/<experiment-timestamp>`
    Returns:
    ----------
        params (dict): the loaded parameters
    '''
    with open(exp_param_filename(path)) as f:
        return yaml.safe_load(f)
    

def save_exp_params(path, params, logger=None):
    '''
    Arguments:
    ----------
        path (str): full path to experiment, i.e. 
            `<path-to-experiments-dir>/<experiment-timestamp>`
        params (dict): dictionary of parameters to save
        logger (logging.Logger): Defaults to None.
    '''
    with open(exp_param_filename(path), 'w') as f:
        yaml.dump(params, f, default_flow_style=False)
    if logger: 
        logger.info('Experiment parameters saved.')

def setup_exp(path, params, name=None):
    '''
    Arguments:
    ----------
        path (str): full path to where to create experiments, i.e. 
            `<path-to-experiments-dir>`
        params (dict): dictionary of parameters to save
    Returns:
    ----------
        exp_dir (str): full path to experiment, i.e. 
            `<path-to-experiments-dir>/<experiment-timestamp>`
        logger (logging.Logger)
    '''
    exp_name = gen_exp_name(name)
    exp_dir = os.path.join(path, exp_name)
    if not os.path.isdir(exp_dir):
        os.makedirs(exp_dir)

    logger = config_exp_logger(exp_dir)    
    save_exp_params(exp_dir, params, logger)
    return exp_dir, logger

    
def is_config_subset(truth, params):
    '''
    Arguments:
    ----------
        truth (dict): dictionary of parameters to compare to
        params (dict): dictionary of parameters to test
    Returns:
    ----------
        result (bool) whether or not `params` is a subset of `truth`
    '''
    if not type(truth) == type(params): return False
    for key, val in params.items():
        if key not in truth: return False
        if type(val) is dict:
            if not is_config_subset(truth[key], val): 
                return False
        else:            
            if not truth[key] == val: return False
    return True


def find_exps(path, params):
    '''
    Arguments:
    ----------
        path (str): full path to where to create experiments, i.e. 
            `<path-to-experiments-dir>`
        params (dict): dictionary of parameters to test
    Returns:
    ----------
        results (str[]) list of experiment names where their parameters are 
            supersets of the provided `params`
    '''
    exps = list_exps(path)
    results = []
    for exp in exps:
        exp_name = os.path.join(path, exp)
        exp_params = load_exp_params(exp_param_filename(exp_name))
        if is_config_subset(exp_params, params):
            results.append(exp)
    return results

# %% ../nbs/02_utils.ipynb 88
import os, pathlib, itertools
from pathlib import Path
from dataclasses import dataclass, field, KW_ONLY
from typing import Optional, List, ClassVar, Any, TypeAlias, Union

from ipos.imp import is_mod, is_var_imp, Module

# %% ../nbs/02_utils.ipynb 89
from iza.types import (
    PathLike, PathType,
    RichConsole, RichProgress, RichText, RichTree
)
from .static import EXT_PY
from .imp import RichImp

# %% ../nbs/02_utils.ipynb 93
@dataclass
class ArchiveDownloader:    
    _: KW_ONLY
    rootdir: str
    archive: str
    entries: Union[str, list[str]]
    savedir: str
    extract: bool = False
    cleanup: bool = False
    compound_archive: bool = False
    archives: Optional[list[str]] = None
    console: Optional[RichConsole] = None
    progress: Optional[RichProgress] = None

    _rich: Module = field(init=False, repr=False, default=None)
    has_rich: bool = field(init=False, repr=False, default=None)

    

    def __post_init__(self):    
        r = RichImp()
        self._rich = r._module
        self.has_rich = is_mod(self._rich)

        self.entries = self.entries if isinstance(self.entries, list) else [self.entries]        
        self.console = get_console()        
        self.progress = self.get_progress()
    
        self.savedir = Path(self.savedir).expanduser()
        make_missing_dirs(self.savedir)

    def get_progress(self):
        progress = getattr(self, 'progress', None)
        if progress is not None:
            return progress

        if self.has_rich and is_var_imp('Progress'):
                self.progress = Progress(console=self.console)
                return self.progress

        return None

    @property
    def path(self) -> str:
        return urljoin(self.rootdir, self.archive)

    @property
    def urls(self) -> list[str]:
        urls = []
        if self.compound_archive and self.archives is not None:
            for archive, entry in itertools.product(self.archives, self.entries):
                urls.append(urljoin(self.rootdir, archive, entry))
        else:
            urls = [urljoin(self.path, entry) for entry in self.entries]
        return urls

    def download_missing_files(self) -> None:        
        total_files = len(self.urls)
        if self.has_rich and self.progress:
            with self.progress:
                task = self.progress.add_task("[cyan]Downloading...", total=total_files)
                for url in self.urls:
                    filename = Path(parse_url(url).path).name
                    fullpath = self.savedir / filename
                    if not fullpath.exists():
                        stream_file(url, str(fullpath))
                        self.progress.advance(task)
                    else:
                        self.progress.advance(task)
        else:            
            for url in tqdm(self.urls, desc='Downloading'):       
                filename = Path(parse_url(url).path).name
                fullpath = self.savedir / filename
                if not fullpath.exists():
                    stream_file(url, str(fullpath))
                

    def calc_n_to_extract(self) -> int:
        files = [self.savedir / entry for entry in self.entries]
        files = get_gz_files_in_dir(self.savedir)
        total = 0
        for file in files:
            if is_tarball(file):
                total += 1
            elif is_gz(file):
                total += 1
        return total                

    def extract_files(self) -> None:
        recurser = RecursiveDecompressor(
            dirname=self.savedir, 
            entries=self.entries,
            strategy=undo_gz,
            remove=self.cleanup,             
            progress=self.progress
        )
        recurser.decompress()
        return
        files = [self.savedir / entry for entry in self.entries]
        files = get_gz_files_in_dir(self.savedir)
        total = self.calc_n_to_extract()
        if self.has_rich and self.progress:
            with self.progress:
                task = self.progress.add_task("[cyan]Extracting...", total=total)
                for file in files:
                    undo_gz(file, remove=self.cleanup)                    
                    self.progress.advance(task)
        else:
            for file in tqdm(files, desc='Extracting'):
                undo_gz(file, remove=self.cleanup)

    def execute(self) -> None:
        if self.has_rich and self.console:
            self.console.print(f"Processing archive: [bold cyan]{self.archive}[/bold cyan]")
        else:
            print(f"Processing archive: {self.archive}")

        self.download_missing_files()
        if self.extract:
            self.extract_files()

        if self.has_rich and self.console:
            dir = RichDirectory(self.savedir, console=self.console)
            dir.print_rich()
        else:
            dir = Directory(self.savedir)
            dir.print()

# %% ../nbs/02_utils.ipynb 97
#| eval: False
try:
    import typer
    app = typer.Typer()

    @app.command()
    def download(rootdir: str, archive: str, entries: List[str], savedir: str, extract: bool = False, cleanup: bool = False):
        downloader = ArchiveDownloader(rootdir, archive, entries, savedir, extract, cleanup)
        downloader.execute()

except ImportError:
    pass
