"""Exact saved-tensor placement with a bounded GPU retention budget."""
import torch


def storage_key(tensor):
    return (str(tensor.device), tensor.untyped_storage().data_ptr())


class BudgetedSavedTensors(torch.autograd.graph.saved_tensors_hooks):
    """Keep existing backbone storage and a bounded set of additional storages.

    Count whole unique storages, not views. Stored tensors are detached aliases;
    no casting, quantization or recomputation changes backward values.
    """
    def __init__(self, resident_tensors=(), keep_bytes=0):
        if keep_bytes < 0:
            raise ValueError('saved-tensor GPU budget must be nonnegative')
        resident = {storage_key(t) for t in resident_tensors}
        kept = set()
        self.retained_bytes = 0
        self.resident_hits = 0
        self.cpu_copies = 0
        cpu = torch.autograd.graph.save_on_cpu(pin_memory=True, device_type='cuda')

        def pack(tensor):
            key = storage_key(tensor)
            size = tensor.untyped_storage().nbytes()
            if key in resident:
                self.resident_hits += 1
                return ('resident', tensor.detach())
            if key in kept or self.retained_bytes + size <= keep_bytes:
                if key not in kept:
                    kept.add(key)
                    self.retained_bytes += size
                return ('resident', tensor.detach())
            self.cpu_copies += 1
            return ('cpu', cpu.pack_hook(tensor))

        def unpack(packed):
            kind, value = packed
            return value if kind == 'resident' else cpu.unpack_hook(value)

        super().__init__(pack, unpack)


class AsyncSavedTensors(torch.autograd.graph.saved_tensors_hooks):
    """Pinned D2H without a host synchronization at every saved tensor.

    An event explicitly orders each H2D restore after its D2H copy, including
    when autograd executes that restore on a different CUDA stream.
    """
    def __init__(self):
        def pack(tensor):
            if tensor.device.type != 'cuda':
                return (tensor.device, tensor.detach(), None)
            packed = torch.empty(tensor.size(), dtype=tensor.dtype,
                                 layout=tensor.layout, pin_memory=True)
            stream = torch.cuda.current_stream(tensor.device)
            packed.copy_(tensor, non_blocking=True)
            tensor.record_stream(stream)
            ready = torch.cuda.Event()
            ready.record(stream)
            return (tensor.device, packed, ready)

        def unpack(saved):
            device, tensor, ready = saved
            if ready is None:
                return tensor
            torch.cuda.current_stream(device).wait_event(ready)
            return tensor.to(device, non_blocking=True)

        super().__init__(pack, unpack)