# NOTE(bvh): Debug utilities for NaN/Inf detection and VRAM tracking in training pipeline.

import gc
import torch
from imaginaire.utils import log


def check_nan_dict(d, label, extra=''):
    '''Check for NaN/Inf in all tensors within a (possibly nested) dict. Logs errors if found.
    Enable via config.check_nan = True.'''
    if not isinstance(d, dict):
        return
    for k, v in d.items():
        if isinstance(v, dict):
            for k2, v2 in v.items():
                if isinstance(v2, torch.Tensor) and (v2.isnan().any() or v2.isinf().any()):
                    nan_n = v2.isnan().sum().item()
                    inf_n = v2.isinf().sum().item()
                    log.error(f'NaN/Inf in {label}[{k}][{k2}]: nan={nan_n}, inf={inf_n}, '
                              f'shape={v2.shape}, {extra}', rank0_only=False)
        elif isinstance(v, torch.Tensor) and (v.isnan().any() or v.isinf().any()):
            nan_n = v.isnan().sum().item()
            inf_n = v.isinf().sum().item()
            log.error(f'NaN/Inf in {label}[{k}]: nan={nan_n}, inf={inf_n}, '
                      f'shape={v.shape}, {extra}', rank0_only=False)


def vram_snapshot(label='', rank0_only=True):
    '''Log current VRAM usage. Returns (allocated_gb, reserved_gb).'''  # DEBUG
    if rank0_only and torch.distributed.is_initialized() and torch.distributed.get_rank() != 0:
        alloc = torch.cuda.memory_allocated() / (1024 ** 3)
        res = torch.cuda.memory_reserved() / (1024 ** 3)
        return (alloc, res)
    alloc = torch.cuda.memory_allocated() / (1024 ** 3)
    res = torch.cuda.memory_reserved() / (1024 ** 3)
    peak = torch.cuda.max_memory_allocated() / (1024 ** 3)
    log.info(f'[VRAM {label}] alloc={alloc:.2f} GB, reserved={res:.2f} GB, peak={peak:.2f} GB',  # DEBUG
             rank0_only=rank0_only)
    return (alloc, res)


def vram_delta(label, prev_alloc, rank0_only=True):
    '''Log VRAM change since previous snapshot. Returns (allocated_gb, reserved_gb).'''  # DEBUG
    alloc, res = vram_snapshot(label, rank0_only=rank0_only)
    delta = alloc - prev_alloc
    if abs(delta) > 0.01:
        if rank0_only and torch.distributed.is_initialized() and torch.distributed.get_rank() != 0:
            pass
        else:
            log.info(f'[VRAM {label}] delta={delta:+.2f} GB', rank0_only=rank0_only)  # DEBUG
    return (alloc, res)


_LAST_CENSUS = {}  # rank -> {(shape, dtype): count} from previous census; for diffing


def _census_bytes(shape, dtype):
    n = 1
    for s in shape:
        n *= s
    s = str(dtype)
    if 'float16' in s or 'bfloat16' in s:
        return n * 2
    if 'float64' in s or 'int64' in s:
        return n * 8
    if 'int8' in s or 'uint8' in s or 'bool' in s:
        return n
    return n * 4  # default fp32 / int32


def gpu_tensor_census(label='', top_n=15, rank0_only=True, diff=True, growth_top_n=10):
    '''Count and summarize all GPU tensors reachable by gc. Optionally print
    *growth* (tensors whose count went up vs. last census on this rank), which
    is far more useful for leak hunting than the static top-N.

    Expensive (linear in # Python objects); use every ~100 iters at most.
    '''  # DEBUG
    if rank0_only and torch.distributed.is_initialized() and torch.distributed.get_rank() != 0:
        return
    gc.collect()
    tensors = {}
    for obj in gc.get_objects():
        try:
            if torch.is_tensor(obj) and obj.is_cuda:
                key = (tuple(obj.shape), str(obj.dtype))
                tensors[key] = tensors.get(key, 0) + 1
        except Exception:
            pass
    rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
    ranked = sorted(tensors.items(), key=lambda x: _census_bytes(x[0][0], x[0][1]) * x[1], reverse=True)

    lines = [f'[GPU TENSOR CENSUS {label}] {len(tensors)} unique (shape,dtype), top {top_n} by total MB:']
    for (shape, dtype), count in ranked[:top_n]:
        mb = _census_bytes(shape, dtype) * count / (1024 ** 2)
        lines.append(f'  {shape} {dtype} x{count} = {mb:.1f} MB')

    if diff:
        prev = _LAST_CENSUS.get(rank, {})
        # Compute growth (count delta)
        growth = []
        for k, c in tensors.items():
            d = c - prev.get(k, 0)
            if d > 0:
                growth.append((k, d, c))
        growth.sort(key=lambda x: _census_bytes(x[0][0], x[0][1]) * x[1], reverse=True)
        if growth:
            lines.append(f'[GPU TENSOR CENSUS {label}] GROWTH vs last census on rank{rank} (top {growth_top_n} by MB delta):')
            for (shape, dtype), d, c in growth[:growth_top_n]:
                d_mb = _census_bytes(shape, dtype) * d / (1024 ** 2)
                lines.append(f'  +{d} (now {c}) {shape} {dtype} delta={d_mb:.1f} MB')
        else:
            lines.append(f'[GPU TENSOR CENSUS {label}] no count growth vs last census on rank{rank}')
        _LAST_CENSUS[rank] = tensors

    log.info('\n'.join(lines), rank0_only=rank0_only)  # DEBUG