"""utils for vre"""
from typing import Any, T, Callable, Iterable
from types import ModuleType
from pathlib import Path
from datetime import datetime, timezone as tz
from collections import OrderedDict
from functools import partial
from math import sqrt
import random
import sys
from importlib.machinery import SourceFileLoader
import numpy as np
from vre.logger import vre_logger as logger

class FixedSizeOrderedDict(OrderedDict):
    """An OrderedDict with a fixed size. Useful for caching purposes."""
    def __init__(self, *args, maxlen: int = 0, **kwargs):
        self._maxlen = maxlen
        super().__init__(*args, **kwargs)

    def __setitem__(self, key, value):
        OrderedDict.__setitem__(self, key, value)
        if self._maxlen > 0:
            if len(self) > self._maxlen:
                self.popitem(False)

def get_project_root() -> Path:
    """gets the root of this project"""
    return Path(__file__).parents[2].absolute()

def parsed_str_type(item: Any) -> str:
    """Given an object with a type of the format: <class 'A.B.C.D'>, parse it and return 'A.B.C.D'"""
    return str(type(item)).rsplit(".", maxsplit=1)[-1][0:-2]

def is_dir_empty(dir_path: Path, pattern: str = "*") -> bool:
    """returns true if directory is not empty, false if it is empty"""
    assert pattern.startswith("*"), pattern
    return len(list(dir_path.glob(pattern))) == 0

def get_closest_square(n: int) -> tuple[int, int]:
    """
    Given a stack of N images, find the closest square X>=N*N and return that.
    Note: There are only 2 rows possible between x^2 and (x+1)^2 because (x+1)^2 = x^2 + 2*x + 1, thus we can add two
    columns at most. If a 3rd column is needed, then closest lower bound is (x+1)^2 and we must use that.
    Example: 9: 3*3; 12 -> 3*3 -> 3*4 (3 rows). 65 -> 8*8 -> 8*9. 73 -> 8*8 -> 8*9 -> 9*9
    """
    x = int(sqrt(n))
    r, c = x, x
    c = c + 1 if c * r < n else c
    r = r + 1 if c * r < n else r
    assert (c + 1) * r > n and c * (r + 1) > n
    return r, c

def now_fmt() -> str:
    """Returns now() as a UTC isoformat string"""
    return datetime.now(tz=tz.utc).replace(tzinfo=None).isoformat(timespec="milliseconds")

def reorder_dict(data: dict[str, Any], keys: list[str]) -> dict[str, Any]:
    """simply puts in front the desired keys from the original dict, keeping the others intact"""
    assert (diff := set(keys).difference(data.keys())) == set(),diff
    for k in keys[::-1]:
        data = {k: data[k], **{k: v for k, v in data.items() if data != k}}
    return data

def array_blend(x: np.ndarray, y: np.ndarray, alpha: float | np.ndarray) -> np.ndarray:
    """Blends two arrays of the same shape with an alpha: number of array"""
    alpha_arr = np.asarray(alpha)
    assert x.shape == y.shape, (x.shape, y.shape)
    assert np.issubdtype(x.dtype, np.floating) and np.issubdtype(y.dtype, np.floating), (x.dtype, y.dtype)
    assert (alpha_arr >= 0).all() and (alpha_arr <= 1).all(), (alpha_arr.min(), alpha_arr.max())
    try:
        return (1 - alpha_arr) * x + alpha_arr * y # the actual blend :)
    except Exception as e:
        logger.info(f"Exception thrown: {e}.\nShapes: {x.shape=} {y.shape=} {alpha_arr.shape=}")
        raise e

def make_batches(frames: list[int], batch_size: int) -> list[int]:
    """return 1D array [start_frame, start_frame+bs, start_frame+2*bs... end_frame]"""
    assert isinstance(batch_size, int) and batch_size > 0, batch_size
    if batch_size > len(frames):
        logger.warning(f"batch size {batch_size} is larger than #frames to process {len(frames)}.")
        batch_size = len(frames)
    if len(frames) == 0:
        return []
    batches, n_batches = [], len(frames) // batch_size + (len(frames) % batch_size > 0)
    for i in range(n_batches):
        batches.append(frames[i * batch_size: (i + 1) * batch_size])
    return batches

def clip(x: T, _min: T, _max: T) -> T:
    """clips a value between [min, max]"""
    return max(min(x, _max), _min)

def load_function_from_module(module_path: str | Path, function_name: str) -> Callable:
    """Usage: fn = load_function_from_module("/path/to/stuff.py", "function_name"); y = fn(args);"""
    module_path = Path(module_path).absolute()
    assert module_path.exists(), module_path
    module = ModuleType(module_path.stem)
    loader = SourceFileLoader(module_path.stem, str(module_path))
    loader.exec_module(module)
    sys.path.append(str(module_path.parent))
    sys.modules[module_path.stem] = module
    return getattr(module, function_name)

def random_chars(n: int) -> str:
    """returns a string of n random characters"""
    valid_chars = [*range(ord('A'), ord('Z')+1), *range(ord('a'), ord('z')+1), *range(ord('0'), ord('9')+1)]
    return "".join(map(chr, [random.choice(valid_chars) for _ in range(n)]))

def mean(l: list[int | float]) -> float:
    """the average of a list of ints or floats"""
    return sum(l) / len(l) if len(l) > 0 else 0

def str_topk(s: str, k: int) -> str:
    """returns a substring of the type 'first_part..last_part' of len(s) > s with the parts being s//2"""
    assert k >= 0, k
    if len(s) <= k:
        return s
    # 4 => 1..1 => (4//2-1) -(4//2-1)
    # 5 => 2..1 => (5//2-1+1) -(5//2-1)
    # 6 => 2..2 => (6//2-1) -(6//2-1)
    # 7 => 3..2 => (7//2-1+1)
    first_part = s[0: k // 2 - 1 + (k%2 == 1)]
    last_part = s[-(k//2 - 1):]
    return f"{first_part}..{last_part}"

def natsorted(seq: Iterable[T], key: Callable[[T], "SupportsGTAndLT"] | None = None, reverse: bool=False) -> list[T]:
    """wrapper on top of natsorted so we can properly remove it"""
    def _try_convert_to_num(x: str) -> str | int | float:
        try:
            return int(x)
        except ValueError:
            try:
                return float(x)
            except ValueError:
                return x

    def natsorted_key(item: T, key: Callable) -> "SupportsGTAndLT":
        item = key(item)
        if isinstance(item, str):
            ix_dot = item.rfind(".")
            item = item[0:ix_dot] if ix_dot != -1 else item
            item = _try_convert_to_num(item)
        return item

    key = key or (lambda item: item)
    return sorted(seq, key=partial(natsorted_key, key=key), reverse=reverse)
