# BVH, Feb - May 2026.

import numpy as np
import torch
from torch.utils.data._utils.collate import default_collate

from imaginaire.utils import log
from custom.utils.logging import dump_batch as _dump_batch, fmt_path as _fmt_path

# Track which warnings have been emitted to avoid spamming every batch.
_warned_keys = set()

# On SageMaker, suppress collate warnings entirely to avoid flooding CloudWatch.
_QUIET = False


def set_collate_quiet(quiet: bool):
    '''
    Set collate warning verbosity. Called from trainer/model init with config value.
    '''
    global _QUIET
    _QUIET = quiet


def _should_log(key):
    '''
    Whether to log a warning for this key. Suppresses (t, c) tuple keys
    where t % 10 != 0 to reduce verbosity from per-timestep/camera warnings.
    '''
    if _QUIET:
        return False

    if isinstance(key, tuple) and len(key) >= 1 and isinstance(key[0], int):
        return key[0] % 10 == 0

    return True


# NOTE(bvh): any feature should be controllable as argument here
def any4d_collate(batch, pad_tensors=True, norm_str_lists=True, impute_missing_keys=True,
                  key_path=()):
    '''
    Collate that handles variable-length string lists and mismatched tensor shapes.

    Args:
        pad_tensors: Pad tensors with mismatched spatial dims to max shape.
        norm_str_lists: Normalize string/list-of-str fields into lists of lists.
        impute_missing_keys: Iterate the union of keys across batch elements. If a
            key is absent in any element, that key collates to None (and zero-fills
            inside _collate_tensor_dict). If False, only iterates over batch[0]'s
            keys (original behavior).
        key_path: tuple of parent keys (root -> ... -> here) used for diagnostics.
    '''
    if not isinstance(batch[0], dict):
        if norm_str_lists and any(_is_string_or_string_list(v) for v in batch):
            return _normalize_string_vals(batch, key=None, key_path=key_path)

        # Leaf-level batch (e.g. tensors, scalars): stack via default_collate,
        # fall back to a raw list with a warning if the type is unhandled.
        return _safe_default_collate(batch, key=None, key_path=key_path)

    # Guard: all batch elements must be dicts (not lists, tensors, etc.)
    non_dict = [(i, type(b).__name__) for i, b in enumerate(batch) if not isinstance(b, dict)]
    if non_dict:
        log.error(  # bypass _QUIET
            f"any4d_collate: expected all dicts but got mixed types at {_fmt_path(key_path)}. "
            f"Non-dict elements: {non_dict}. "
            f"Types: {[type(b).__name__ for b in batch]}. "
            f"Sample values: {[{k: type(v).__name__ for k, v in (b.items() if isinstance(b, dict) else [('_value_', b)])} for b in batch[:3]]}\n"
            f"Batch:\n{_dump_batch(batch)}",
            rank0_only=False,
        )
        return list(batch)

    # Collect keys: union across all batch elements, or just batch[0].
    if impute_missing_keys:
        all_keys = dict.fromkeys(k for b in batch for k in b)
    else:
        all_keys = batch[0].keys()

    result = {}
    for key in all_keys:
        child_path = key_path + (key,)
        try:
            result[key] = _collate_single_key(
                key, batch, pad_tensors, norm_str_lists, impute_missing_keys,
                key_path=child_path,
            )

        except Exception:
            # Collate has no retry mechanism (unlike dataloader workers), so swallow
            # errors per-key and return None. Training code must handle None gracefully.
            warn_key = ("collate_error", child_path)
            if warn_key not in _warned_keys:
                _warned_keys.add(warn_key)
                types = [type(b.get(key, None)).__name__ if isinstance(b, dict) else type(b).__name__ for b in batch]
                log.error(  # bypass _QUIET
                    f"any4d_collate: FAILED to collate at {_fmt_path(child_path)}, returning None. "
                    f"Types across batch: {types}\n"
                    f"Batch:\n{_dump_batch(batch)}",
                    rank0_only=False,
                )
            result[key] = None

    return result


def _collate_single_key(key, batch, pad_tensors, norm_str_lists, impute_missing_keys,
                        key_path=()):
    '''
    Collate a single key across batch elements. Raises on failure (caller catches).
    '''
    if impute_missing_keys:
        vals = [b.get(key, None) for b in batch]
    else:
        vals = [b[key] for b in batch]

    # If any value is None, return None. Downstream model code expects a tensor or None
    if impute_missing_keys and any(v is None for v in vals):
        absent_count = sum(1 for b in batch if key not in b)

        if absent_count > 0:
            warn_key = ("missing_key", key_path)

            if warn_key not in _warned_keys:
                _warned_keys.add(warn_key)
                _should_log(key) and log.warning(
                    f"any4d_collate {_fmt_path(key_path)} absent in {absent_count}/{len(batch)} "
                    f"batch elements, returning None\n"
                    f"Batch:\n{_dump_batch(batch)}",
                    rank0_only=False,
                )

        return None

    # Handle mixed dict/non-dict across batch elements.
    # This happens when a key maps to a dict in some samples (e.g. per-camera
    # metadata like {(0,0): tensor, (0,1): tensor}) but a plain value in others
    # (e.g. [720, 1280] when all cameras share the same resolution).
    # Fix: expand non-dict values into dicts with the same keys as the dict
    # elements, duplicating the shared value across all keys, then collate.
    has_dict = any(isinstance(v, dict) for v in vals)

    if has_dict:
        if not all(isinstance(v, dict) for v in vals):
            ref_dict = next(v for v in vals if isinstance(v, dict))
            ref_keys = list(ref_dict.keys())

            warn_key = ("mixed_dict_nondict", key_path)
            if warn_key not in _warned_keys:
                _warned_keys.add(warn_key)
                _should_log(key) and log.warning(
                    f"any4d_collate {_fmt_path(key_path)} has mixed dict/non-dict across batch. "
                    f"Expanding non-dict values to match dict keys {ref_keys[:4]}...\n"
                    f"Batch:\n{_dump_batch(batch)}",
                    rank0_only=False,
                )
            vals = [
                v if isinstance(v, dict) else {k: v for k in ref_keys}  #@IgnoreException
                for v in vals
            ]

        # Now all vals are dicts, check if they contain tensors for padding
        sample_val = next(iter(vals[0].values()), None) if vals[0] else None
        if (
            pad_tensors
            and isinstance(sample_val, torch.Tensor)
            and sample_val.dim() >= 2
        ):
            return _collate_tensor_dict(vals, impute_missing_keys=impute_missing_keys,
                                        key_path=key_path)

        return any4d_collate(
            vals, pad_tensors=pad_tensors, norm_str_lists=norm_str_lists,
            impute_missing_keys=impute_missing_keys, key_path=key_path,
        )

    elif norm_str_lists and any(_is_string_or_string_list(v) for v in vals):
        return _normalize_string_vals(vals, key=key, key_path=key_path)

    else:
        # Leaf value for this key (no dict / string-list branches matched):
        # stack via default_collate, fall back to a raw list with a warning.
        return _safe_default_collate(vals, key=key, key_path=key_path)


def _is_string_or_string_list(val):
    '''
    Check if val is a string, or a (possibly nested) list of strings.
    '''
    if isinstance(val, str):
        return True
    
    if isinstance(val, (list, tuple)) and len(val) > 0:
        return _is_string_or_string_list(val[0])
    
    return False


def _normalize_string_vals(vals, key=None, key_path=()):
    '''
    Normalize mixed str / list-of-str into a list of lists.

    If any element is a list (of strings), bare strings are wrapped in [str].
    If all are bare strings, returns list of strings as-is.
    Warns once per key on mixed types or ragged lengths.
    '''
    has_list = any(isinstance(v, (list, tuple)) for v in vals)
    if not has_list:
        # All bare strings: return as-is, no wrapping needed.
        return list(vals)

    # Normalize: wrap bare strings in lists
    has_str = any(isinstance(v, str) for v in vals)
    if has_str:
        warn_key = ("mixed_str_list", key_path or key)

        if warn_key not in _warned_keys:
            _warned_keys.add(warn_key)
            types = [type(v).__name__ for v in vals]
            _should_log(key) and log.warning(
                f"any4d_collate: mixed str/list types at {_fmt_path(key_path)}, "
                f"wrapping bare strings in lists. Types: {types}\n"
                f"Batch:\n{_dump_batch(vals)}",
                rank0_only=False,
            )

    normalized = [v if isinstance(v, (list, tuple)) else [v] for v in vals]

    # Warn on ragged lengths (lists of different sizes across batch)
    lengths = [len(v) for v in normalized]
    if len(set(lengths)) > 1:
        warn_key = ("ragged_str", key_path or key)

        if warn_key not in _warned_keys:
            _warned_keys.add(warn_key)
            _should_log(key) and log.warning(
                f"any4d_collate: ragged string lists at {_fmt_path(key_path)}, "
                f"lengths: {lengths}\n"
                f"Batch:\n{_dump_batch(vals)}",
                rank0_only=False,
            )

    return normalized


def _safe_default_collate(vals, key=None, key_path=()):
    '''
    default_collate (stacks same-shape tensors / nests dicts) with a raw-list +
    once-per-key warning fallback for unexpected datatypes -- e.g. numpy object-dtype
    arrays of dicts from HumanoidEveryday hand_pressure_state. No actual padding here;
    spatial padding for mismatched tensor shapes lives in _collate_tensor_dict.
    '''
    try:
        return default_collate(vals)  #@IgnoreException
    except (RuntimeError, TypeError) as e:
        _warn_leaf_fallback(vals, key=key, err=e, key_path=key_path)
        return list(vals)


def _warn_leaf_fallback(vals, key, err, key_path=()):
    '''
    Emit a once-per-key warning that we are returning the raw vals list.
    '''
    warn_key = ("leaf_fallback", key_path or key)
    if warn_key in _warned_keys:
        return

    _warned_keys.add(warn_key)
    types = [type(v).__name__ for v in vals]

    # Surface numpy object-dtype arrays explicitly (common HumEv-style edge case).
    obj_arr = any(isinstance(v, np.ndarray) and v.dtype == object for v in vals)
    extra = " [numpy object-dtype array detected]" if obj_arr else ""

    _should_log(key) and log.warning(
        f"any4d_collate: cannot collate leaf at {_fmt_path(key_path)}; "
        f"returning raw list of {len(vals)} elements.{extra} "
        f"Types: {types}. Error: {type(err).__name__}: {err}\n"
        f"Batch:\n{_dump_batch(vals)}",
        rank0_only=False,
    )


def _collate_tensor_dict(dict_list, impute_missing_keys=True, key_path=()):
    '''
    Collate a list of dicts mapping keys to tensors, padding spatial dims.
    NOTE(bvh): When these items represent images or videos, we have to think very carefully about
        intrinsics. Camera objects are constructed by prepare_batch() AFTER any4d_collate(), which
        means that BOTTOM RIGHT padding is the only "correct" operation to apply here, otherwise
        the K matrix would not be able to safely remain unchanged. Right now, cx, cy, fx, fy stay
        at the same pixel coordinates measured from the top left, and the pinhole model simply
        _extrapolates_ beyond the original field of view as needed to construct Plucker rays.
        WARNING: I believe gopro is also correct, but not sure about omni? (TODO verify)
    TODO(bvh): padded pixels are still included inside input/output/supervise masks, should be excluded
        by reading each sample's actual resolution in dataloaders
    '''
    if impute_missing_keys:
        all_keys = dict.fromkeys(k for d in dict_list for k in d)
    else:
        all_keys = dict_list[0].keys()
    result = {}

    for key in all_keys:
        child_path = key_path + (key,)
        # Collect tensors, zero-filling missing entries to match present ones.
        has_missing = impute_missing_keys and any(key not in d for d in dict_list)
        if has_missing:
            # Find a reference tensor from a sample that has this key.
            ref = next(d[key] for d in dict_list if key in d)
            warn_key = ("missing_tensor_key", child_path)

            if warn_key not in _warned_keys:
                _warned_keys.add(warn_key)
                present = sum(1 for d in dict_list if key in d)
                _should_log(key) and log.warning(
                    f"any4d_collate: tensor dict {_fmt_path(child_path)} missing in "
                    f"{len(dict_list) - present}/{len(dict_list)} batch elements, "
                    f"zero-filling with shape {list(ref.shape)}\n"
                    f"Batch:\n{_dump_batch(dict_list)}",
                    rank0_only=False,
                )
            tensors = [
                d[key] if key in d else torch.zeros_like(ref)
                for d in dict_list
            ]

        else:
            tensors = [d[key] for d in dict_list]

        # Check if all shapes match
        shapes = [t.shape for t in tensors]
        if all(s == shapes[0] for s in shapes):
            result[key] = torch.stack(tensors)

        else:
            # Pad each tensor to the per-dim max across the batch (all dims).
            warn_key = ("pad_tensor", child_path)
            if warn_key not in _warned_keys:
                _warned_keys.add(warn_key)
                _should_log(key) and log.warning(
                    f"any4d_collate: padding tensors at {_fmt_path(child_path)}, "
                    f"shapes: {shapes}\n"
                    f"Batch:\n{_dump_batch(dict_list)}",
                    rank0_only=False,
                )
            
            max_shape = list(shapes[0])
            for s in shapes[1:]:
                for i in range(len(max_shape)):
                    max_shape[i] = max(max_shape[i], s[i])

            padded = []
            for t in tensors:
                if list(t.shape) == max_shape:
                    padded.append(t)
                
                else:
                    pad_t = torch.zeros(max_shape, dtype=t.dtype, device=t.device)
                    slices = tuple(slice(0, s) for s in t.shape)
                    pad_t[slices] = t
                    padded.append(pad_t)
            
            result[key] = torch.stack(padded)

    return result