# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""Exact input and training-state parity probes."""

from __future__ import annotations

import hashlib
import json
import os
import pickle
import random
from collections.abc import Iterable

import numpy as np
import torch

_SCHEMA_VERSION = 3


def _to_cpu_tensor(tensor: torch.Tensor) -> torch.Tensor:
    """Materialize a detached, contiguous CPU snapshot of a tensor."""
    try:
        # Rank-filtered distributed probes must not launch a full-tensor
        # collective from only the selected rank. Compare the local FSDP shard.
        if hasattr(tensor, "to_local"):
            tensor = tensor.to_local()  # [*local_shape]
        elif hasattr(tensor, "full_tensor"):
            tensor = tensor.full_tensor()  # [*shape]
    except Exception:  # noqa: BLE001
        pass

    return tensor.detach().to("cpu").contiguous().clone()  # [*shape]


def _fingerprint_cpu(tensor: torch.Tensor) -> dict[str, object]:
    """Return an exact raw-byte hash plus useful numeric summaries."""
    raw_bytes = tensor.reshape(-1).view(torch.uint8).numpy().tobytes()
    info: dict[str, object] = {
        "shape": tuple(tensor.shape),
        "dtype": str(tensor.dtype),
        "numel": tensor.numel(),
        "nbytes": len(raw_bytes),
        "sha256": hashlib.sha256(raw_bytes).hexdigest(),
    }
    if tensor.numel() == 0:
        return info

    if tensor.is_floating_point():
        tensor_float = tensor.float()  # [*shape]
        info["sum"] = tensor_float.sum().item()
        info["mean"] = tensor_float.mean().item()
        info["absmax"] = tensor_float.abs().max().item()
    else:
        info["sum"] = int(tensor.to(torch.int64).sum().item())
    return info


def _exact_tensor_record(tensor: torch.Tensor) -> dict[str, object]:
    distribution: dict[str, object] = {}
    device_mesh = getattr(tensor, "device_mesh", None)
    if device_mesh is not None:
        mesh_ranks = device_mesh.mesh.detach().cpu()  # [*mesh_shape]
        distribution = {
            "mesh_dim_names": tuple(device_mesh.mesh_dim_names or ()),
            "mesh_shape": tuple(mesh_ranks.shape),
            "mesh_ranks": mesh_ranks.tolist(),
            "mesh_coordinate": device_mesh.get_coordinate(),
            "placements": tuple(str(placement) for placement in getattr(tensor, "placements", ())),
        }
    tensor_cpu = _to_cpu_tensor(tensor)  # [*shape]
    raw_bytes = tensor_cpu.reshape(-1).view(torch.uint8).numpy().tobytes()
    return {
        "shape": tuple(tensor_cpu.shape),
        "dtype": str(tensor_cpu.dtype),
        "numel": tensor_cpu.numel(),
        "nbytes": len(raw_bytes),
        "sha256": hashlib.sha256(raw_bytes).hexdigest(),
        **distribution,
    }


def _step_in_probe_range(step: int) -> bool:
    minimum = int(os.environ.get("COSMOS_PROBE_MIN_STEP", "0"))
    maximum = int(os.environ.get("COSMOS_PROBE_MAX_STEPS", "100"))
    return minimum <= step < maximum


def _probe_location(step: int, *, rank0_only: bool = False) -> tuple[str, int] | None:
    probe_dir = os.environ.get("COSMOS_PROBE_DIR")
    if not probe_dir or not _step_in_probe_range(step):
        return None
    rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
    if rank0_only and rank != 0:
        return None
    return probe_dir, rank


def _parity_probe_location(step: int) -> tuple[str, int] | None:
    return _probe_location(step, rank0_only=True)


def _deep_probe_enabled(step: int) -> bool:
    return _parity_probe_location(step) is not None


def _loss_probe_location(step: int) -> tuple[str, int] | None:
    return _probe_location(step)


def _small_tensor_record(tensor: torch.Tensor) -> dict[str, object]:
    record = _exact_tensor_record(tensor)
    tensor_cpu = _to_cpu_tensor(tensor)  # [*shape]
    raw_bytes = tensor_cpu.reshape(-1).view(torch.uint8).numpy().tobytes()  # [N_bytes]
    record["values"] = tensor_cpu.tolist()
    record["hex"] = raw_bytes.hex()
    return record


def maybe_dump_loss_reduction(
    *,
    step: int | None,
    tag: str | None,
    valid_counts: torch.Tensor,
    local_loss_sum: torch.Tensor,
    denominator_before: torch.Tensor,
    denominator_after: torch.Tensor,
    final_loss: torch.Tensor,
    exponent: float,
    loss_scaling_factor: float,
    world_size: int,
) -> None:
    """Save all local and globally reduced weighted-CE scalars on selected ranks."""
    if step is None or tag is None:
        return
    location = _loss_probe_location(step)
    if location is None:
        return
    probe_dir, rank = location
    payload = {
        "schema_version": _SCHEMA_VERSION,
        "step": step,
        "rank": rank,
        "tag": tag,
        "valid_counts": _small_tensor_record(valid_counts),
        "local_loss_sum": _small_tensor_record(local_loss_sum),
        "denominator_before": _small_tensor_record(denominator_before),
        "denominator_after": _small_tensor_record(denominator_after),
        "final_loss": _small_tensor_record(final_loss),
        "exponent": exponent,
        "loss_scaling_factor": loss_scaling_factor,
        "world_size": world_size,
    }
    os.makedirs(probe_dir, exist_ok=True)
    path = os.path.join(probe_dir, f"step{step:05d}_rank{rank:03d}_loss_reduction.pt")
    tmp = path + ".tmp"
    torch.save(payload, tmp)
    os.replace(tmp, path)


def _update_parity_record(step: int, tag: str, section: str, value: object) -> None:
    location = _parity_probe_location(step)
    if location is None:
        return
    probe_dir, rank = location
    os.makedirs(probe_dir, exist_ok=True)
    path = os.path.join(probe_dir, f"step{step:05d}_rank{rank:03d}_parity.pt")
    payload = (
        torch.load(path, map_location="cpu", weights_only=False)
        if os.path.exists(path)
        else {"schema_version": _SCHEMA_VERSION, "step": step, "rank": rank, "tag": tag, "sections": {}}
    )
    payload["sections"][section] = value
    tmp = path + ".tmp"
    torch.save(payload, tmp)
    os.replace(tmp, path)


def _hash_named_tensors(
    named_tensors: Iterable[tuple[str, torch.Tensor]],
    *,
    include_requires_grad: bool = False,
) -> dict[str, object]:
    result: dict[str, object] = {}
    for name, tensor in named_tensors:
        tensor_record = _exact_tensor_record(tensor)
        if include_requires_grad:
            tensor_record["requires_grad"] = tensor.requires_grad
        result[name] = tensor_record
    return result


def _rng_record() -> dict[str, object]:
    cuda_states = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else []  # [N_devices][N_state]
    return {
        "python": hashlib.sha256(pickle.dumps(random.getstate())).hexdigest(),
        "numpy": hashlib.sha256(pickle.dumps(np.random.get_state())).hexdigest(),
        "torch_cpu": _exact_tensor_record(torch.get_rng_state()),
        "torch_cuda": [_exact_tensor_record(state) for state in cuda_states],
    }


def _config_record(model: torch.nn.Module) -> dict[str, object]:
    config = getattr(model, "config", None)
    config_value = config.to_dict() if hasattr(config, "to_dict") else repr(config)
    serialized = json.dumps(config_value, sort_keys=True, default=str).encode()
    return {
        "class": type(model).__qualname__,
        "config_sha256": hashlib.sha256(serialized).hexdigest(),
        "training": model.training,
        "default_dtype": str(torch.get_default_dtype()),
        "deterministic_algorithms": torch.are_deterministic_algorithms_enabled(),
        "cudnn_deterministic": torch.backends.cudnn.deterministic,
        "cudnn_benchmark": torch.backends.cudnn.benchmark,
        "flash_sdp_enabled": torch.backends.cuda.flash_sdp_enabled(),
        "mem_efficient_sdp_enabled": torch.backends.cuda.mem_efficient_sdp_enabled(),
        "math_sdp_enabled": torch.backends.cuda.math_sdp_enabled(),
        "cudnn_sdp_enabled": torch.backends.cuda.cudnn_sdp_enabled(),
    }


def maybe_dump_pre_forward(
    model: torch.nn.Module,
    kwargs: dict[str, object],
    step: int | None,
    tag: str | None,
) -> None:
    """Hash effective raw-model inputs, model state, configuration, and RNG."""
    if step is None or tag is None or not _deep_probe_enabled(step):
        return
    effective_inputs = {
        key: _exact_tensor_record(value) if isinstance(value, torch.Tensor) else repr(value)
        for key, value in kwargs.items()
        if value is not None
    }
    _update_parity_record(
        step,
        tag,
        "pre_forward",
        {
            "effective_inputs": effective_inputs,
            "parameters": _hash_named_tensors(model.named_parameters(), include_requires_grad=True),
            "buffers": _hash_named_tensors(model.named_buffers()),
            "rng": _rng_record(),
            "model": _config_record(model),
        },
    )


def maybe_dump_forward_result(
    logits: torch.Tensor,
    losses: dict[str, torch.Tensor],
    step: int,
    tag: str,
) -> None:
    """Hash logits and every scalar loss participating in backward."""
    if not _deep_probe_enabled(step):
        return
    _update_parity_record(
        step,
        tag,
        "forward_result",
        {
            "logits": _exact_tensor_record(logits),
            "losses": {name: _exact_tensor_record(loss) for name, loss in losses.items()},
        },
    )


def maybe_dump_gradients(model: torch.nn.Module, step: int, tag: str) -> None:
    """Hash every parameter gradient before clipping or optimizer mutation."""
    if not _deep_probe_enabled(step):
        return
    gradients = {
        name: None if parameter.grad is None else _exact_tensor_record(parameter.grad)
        for name, parameter in model.named_parameters()
    }
    _update_parity_record(step, tag, "gradients", gradients)


def maybe_dump_post_clip(
    model: torch.nn.Module,
    total_norm: torch.Tensor,
    max_norm: float,
    step: int,
    tag: str,
) -> None:
    """Hash the exact clipping scalar and every post-clipping gradient."""
    if not _deep_probe_enabled(step):
        return
    clip_coefficient = torch.clamp(max_norm / (total_norm + 1e-6), max=1.0)  # []
    norm_cpu = _to_cpu_tensor(total_norm)  # []
    gradients = {
        name: None if parameter.grad is None else _exact_tensor_record(parameter.grad)
        for name, parameter in model.named_parameters()
    }
    _update_parity_record(
        step,
        tag,
        "post_clip",
        {
            "total_norm": _exact_tensor_record(norm_cpu),
            "max_norm": max_norm,
            "clip_coefficient": _exact_tensor_record(clip_coefficient),
            "gradients": gradients,
        },
    )


def _iter_optimizers(optimizer: object) -> Iterable[object]:
    children = getattr(optimizer, "optimizers", None)
    if children is None:
        yield optimizer
        return
    for child in children:
        if isinstance(child, (list, tuple)):
            yield from child
        else:
            yield child


def _optimizer_record(model: torch.nn.Module, optimizer: object) -> dict[str, object]:
    parameter_name = {id(parameter): name for name, parameter in model.named_parameters()}
    state: dict[str, object] = {}
    for inner_optimizer in _iter_optimizers(optimizer):
        for parameter, parameter_state in getattr(inner_optimizer, "state", {}).items():
            name = parameter_name.get(id(parameter), f"<unmapped:{id(parameter)}>")
            state[name] = {
                key: _exact_tensor_record(value) if isinstance(value, torch.Tensor) else repr(value)
                for key, value in parameter_state.items()
            }
    return state


def maybe_dump_post_optimizer(
    model: torch.nn.Module,
    optimizer: object,
    step: int,
    tag: str,
) -> None:
    """Hash post-step parameters and Adam state independently of container layout."""
    if not _deep_probe_enabled(step):
        return
    _update_parity_record(
        step,
        tag,
        "post_optimizer",
        {
            "parameters": _hash_named_tensors(model.named_parameters(), include_requires_grad=True),
            "optimizer_state": _optimizer_record(model, optimizer),
        },
    )


def maybe_dump_model_inputs(
    kwargs: dict[str, object],
    step: int,
    tag: str,
    *,
    labels: torch.Tensor | None = None,
) -> None:
    """Dump forward kwargs for ``step`` for trainer label if the probe is enabled.

    Args:
        kwargs: the dict expanded into the model forward (before any wrapper
            filtering).
        step: current optimizer step / iteration.
        tag: trainer label, stored as metadata.
        labels: supervision tensor consumed by the external CE loss. It is not
            a model-forward kwarg, but is required for complete loss parity.
    """
    location = _probe_location(step)
    if location is None:
        return
    probe_dir, rank = location
    probe_values = dict(kwargs)
    auxiliary_keys: list[str] = []
    if labels is not None:
        probe_values["labels"] = labels
        auxiliary_keys.append("labels")

    record: dict[str, object] = {
        "schema_version": _SCHEMA_VERSION,
        "step": int(step),
        "rank": int(rank),
        "tag": tag,
        "keys": sorted(probe_values),
        "forward_keys": sorted(kwargs),
        "auxiliary_keys": auxiliary_keys,
    }
    for key, value in probe_values.items():
        if isinstance(value, torch.Tensor):
            try:
                tensor_cpu = _to_cpu_tensor(value)
                record[key] = _fingerprint_cpu(tensor_cpu)
            except Exception as e:  # noqa: BLE001
                record[key] = {"error": repr(e)}
        else:
            record[key] = {"type": type(value).__name__, "value": repr(value)}

    ii = kwargs.get("input_ids")
    if isinstance(ii, torch.Tensor) and ii.ndim >= 2:
        record["seq_len"] = int(ii.shape[1])

    os.makedirs(probe_dir, exist_ok=True)
    path = os.path.join(probe_dir, f"step{step:05d}_rank{rank:03d}.pt")
    tmp = path + ".tmp"
    torch.save({"record": record, "tensors": {}}, tmp)
    os.replace(tmp, path)