# examples/megatron.py
"""Offline serial MLP shard algebra from Megatron-LM (2019/2020), Section 3.

Authored CPU reference, not an upstream port or distributed implementation.
Exact erf GELU is an explicit choice; the existing figure uses tanh GELU.
Dropout, residuals, normalization, attention and training are omitted. Biases
are optional. One global squared-error loss supplies each shard's gradients.
"""

import argparse
import hashlib
import json
from pathlib import Path

import torch
import torch.nn.functional as functional


def validate_mlp(x, a, b, hidden_bias=None, output_bias=None):
    """Reject incompatible, empty, nonfinite or non-CPU floating tensors."""
    tensors = [x, a, b]
    tensors.extend(t for t in (hidden_bias, output_bias) if t is not None)
    if any(not isinstance(t, torch.Tensor) for t in tensors):
        raise ValueError("MLP inputs must be tensors")
    if any(t.device.type != "cpu" for t in tensors):
        raise ValueError("This serial reference supports CPU tensors only")
    if any(t.dtype not in (torch.float32, torch.float64) for t in tensors):
        raise ValueError("Use float32 or float64 tensors")
    if any(t.dtype != x.dtype for t in tensors):
        raise ValueError("All tensors must have the same dtype")
    if any(not torch.isfinite(t).all() for t in tensors):
        raise ValueError("All tensor values must be finite")
    if x.ndim != 3 or a.ndim != 2 or b.ndim != 2:
        raise ValueError("Expected X [batch, sequence, D], A [D,F], B [F,D]")
    if min(*x.shape, *a.shape, *b.shape) <= 0:
        raise ValueError("All dimensions must be positive")
    if a.shape[0] != x.shape[2] or b.shape != (a.shape[1], x.shape[2]):
        raise ValueError("Weight dimensions do not match X [batch, sequence,D]")
    if hidden_bias is not None and hidden_bias.shape != (a.shape[1],):
        raise ValueError("Hidden bias must have shape [F]")
    if output_bias is not None and output_bias.shape != (x.shape[2],):
        raise ValueError("Output bias must have shape [D]")


# region: dense_reference
def dense_reference(x, a, b, hidden_bias=None, output_bias=None):
    """Compute the unsplit MLP branch, with no dropout or residual add."""
    validate_mlp(x, a, b, hidden_bias, output_bias)
    preactivation = x @ a
    if hidden_bias is not None:
        preactivation = preactivation + hidden_bias
    # Exact GELU(t) = t * (1 + erf(t / sqrt(2))) / 2.
    hidden = functional.gelu(preactivation, approximate="none")
    output = hidden @ b
    if output_bias is not None:
        output = output + output_bias
    return output
# endregion: dense_reference


# region: partition_weights
def partition_weights(a, b, hidden_bias, shards):
    """Split A columns and matching B rows into independent leaf tensors."""
    if isinstance(shards, bool) or not isinstance(shards, int) or shards < 1:
        raise ValueError("shards must be a positive integer")
    if a.shape[1] % shards:
        raise ValueError("FFN width F must be divisible by shards")
    a_parts = a.chunk(shards, dim=1)  # [D, F/P]
    b_parts = b.chunk(shards, dim=0)  # [F/P, D]
    bias_parts = (
        hidden_bias.chunk(shards) if hidden_bias is not None
        else [None] * shards
    )
    def leaf(tensor):
        return tensor.detach().clone().requires_grad_(True)

    return [
        (leaf(ai), leaf(bi), leaf(ai_bias) if ai_bias is not None else None)
        for ai, bi, ai_bias in zip(a_parts, b_parts, bias_parts)
    ]
# endregion: partition_weights


# region: local_hidden
def local_hidden(x, partitions):
    """Replicate X as independent leaves and activate owned hidden units."""
    inputs = [x.detach().clone().requires_grad_(True) for _ in partitions]
    hidden = []
    for xi, (ai, _, ai_bias) in zip(inputs, partitions):
        preactivation = xi @ ai
        if ai_bias is not None:
            preactivation = preactivation + ai_bias
        hidden.append(functional.gelu(preactivation, approximate="none"))
    return inputs, hidden  # Each hidden shard is [batch, sequence, F/P].
# endregion: local_hidden


# region: local_partial_outputs
def local_partial_outputs(hidden, partitions):
    """Multiply each hidden slice by only its matching rows of B."""
    return [yi @ bi for yi, (_, bi, _) in zip(hidden, partitions)]
    # Every partial output has the full [batch, sequence, D] shape.
# endregion: local_partial_outputs


# region: sum_outputs
def sum_outputs(partials, output_bias=None):
    """Emulate forward g by summing, then add the unsplit bias once."""
    output = torch.stack(partials, dim=0).sum(dim=0)  # Never mean.
    if output_bias is not None:
        output = output + output_bias
    return output  # Dropout is disabled; no residual is included.
# endregion: sum_outputs


def serial_mlp(x, a, b, shards=2, hidden_bias=None, output_bias=None):
    """Return serial shard tensors and one shared output computation graph.

    Args:
        x: CPU input tensor of shape [batch, sequence, D].
        a: First weight matrix [D, F].
        b: Second weight matrix [F, D].
        shards: Positive integer evenly dividing F.
        hidden_bias: Optional bias [F], partitioned with A.
        output_bias: Optional unsplit bias [D], added once after summation.

    Returns:
        Dictionary holding independently owned leaves, partials and output.
    """
    validate_mlp(x, a, b, hidden_bias, output_bias)
    partitions = partition_weights(a, b, hidden_bias, shards)
    inputs, hidden = local_hidden(x, partitions)
    partials = local_partial_outputs(hidden, partitions)
    bias_leaf = (
        output_bias.detach().clone().requires_grad_(True)
        if output_bias is not None else None
    )
    return {
        "inputs": inputs, "partitions": partitions, "hidden": hidden,
        "partials": partials, "output_bias": bias_leaf,
        "output": sum_outputs(partials, bias_leaf),
    }


# region: gradient_ownership
def gradient_ownership(run, target):
    """Differentiate ONE global loss and reassemble owned gradients."""
    output = run["output"]
    if (target.shape != output.shape or target.dtype != output.dtype
            or target.device != output.device
            or not torch.isfinite(target).all()):
        raise ValueError(
            "target must match output shape/dtype/device and be finite"
        )
    loss = 0.5 * (output - target).square().sum()
    leaves = list(run["inputs"])
    for ai, bi, ai_bias in run["partitions"]:
        leaves.extend([ai, bi] + ([] if ai_bias is None else [ai_bias]))
    if run["output_bias"] is not None:
        leaves.append(run["output_bias"])
    gradients = torch.autograd.grad(loss, leaves)
    count = len(run["inputs"])
    dx_parts = list(gradients[:count])
    da, db, dbias = [], [], []
    position = count
    for _, _, ai_bias in run["partitions"]:
        da.append(gradients[position])
        db.append(gradients[position + 1])
        position += 2
        if ai_bias is not None:
            dbias.append(gradients[position])
            position += 1
    return {
        "loss": loss.detach(), "dx_parts": dx_parts,
        # Backward f sums the independent shares of dX.
        "dx": torch.stack(dx_parts).sum(dim=0),
        "da": torch.cat(da, dim=1), "db": torch.cat(db, dim=0),
        "hidden_bias": torch.cat(dbias) if dbias else None,
        "output_bias": gradients[position]
        if run["output_bias"] is not None else None,
    }
# endregion: gradient_ownership


# region: wrong_split_counterexample
def wrong_split_counterexample():
    """Show that row-split preactivations must be summed BEFORE GELU."""
    x = torch.tensor([[[0.7, -1.2]]], dtype=torch.float64)
    a = torch.tensor([[1.1, -0.4], [0.6, 1.3]], dtype=torch.float64)
    first = x[..., :1] @ a[:1, :]
    second = x[..., 1:] @ a[1:, :]
    correct = functional.gelu(first + second, approximate="none")
    wrong = (
        functional.gelu(first, approximate="none")
        + functional.gelu(second, approximate="none")
    )
    return {"preactivation_parts": [first, second],
            "correct": correct, "wrong": wrong,
            "max_error": (correct - wrong).abs().max()}
# endregion: wrong_split_counterexample


def smoke(seed):
    """Report executed outputs, gradients and dense errors for 1/2/4 shards."""
    generator = torch.Generator(device="cpu").manual_seed(seed)

    def random(shape):
        return torch.randn(shape, generator=generator, dtype=torch.float64)

    x, a, b = random((2, 3, 4)), random((4, 8)), random((8, 4))
    hidden_bias, output_bias = random((8,)), random((4,))
    target = random((2, 3, 4))
    dense_leaves = [t.clone().requires_grad_(True)
                    for t in (x, a, b, hidden_bias, output_bias)]
    dense = dense_reference(*dense_leaves[:3], *dense_leaves[3:])
    dense_loss = 0.5 * (dense - target).square().sum()
    dense_grads = torch.autograd.grad(dense_loss, dense_leaves)
    observations = []
    for shards in (1, 2, 4):
        run = serial_mlp(x, a, b, shards, hidden_bias, output_bias)
        owned = gradient_ownership(run, target)
        names = ("dx", "da", "db", "hidden_bias", "output_bias")
        observations.append({
            "shards": shards,
            "hidden_shape": list(run["hidden"][0].shape),
            "partial_shape": list(run["partials"][0].shape),
            "loss": owned["loss"].item(),
            "output_norm": run["output"].norm().item(),
            "output_max_error": (dense - run["output"]).abs().max().item(),
            "gradient_norms": {k: owned[k].norm().item() for k in names},
            "gradient_max_errors": {
                k: (owned[k] - oracle).abs().max().item()
                for k, oracle in zip(names, dense_grads)
            },
            "dx_partial_norms": [t.norm().item() for t in owned["dx_parts"]],
        })
    counterexample = wrong_split_counterexample()
    return {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "seed": seed,
        "reductions": {
            "forward_g": "sum of partial outputs; output bias added once",
            "backward_f": "sum of independent input-gradient contributions",
            "loss": "one global 0.5 * sum squared error; no replicated losses",
            "scope": "offline serial CPU float64; exact erf GELU; no dropout",
        },
        "shapes": {"x": list(x.shape), "a": list(a.shape), "b": list(b.shape),
                   "hidden_bias": [8], "output_bias": [4]},
        "observations": observations,
        "row_split": {
            "preactivation_parts": [t.tolist() for t in
                                    counterexample["preactivation_parts"]],
            "correct": counterexample["correct"].tolist(),
            "wrong": counterexample["wrong"].tolist(),
            "max_error": counterexample["max_error"].item(),
        },
    }


def main():
    """Print finite smoke JSON on explicit invocation only."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", required=True)
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    print(json.dumps(smoke(args.seed), indent=2, allow_nan=False))


if __name__ == "__main__":
    main()