"""Utility functions for optimizers and metric calculation.

Currently handles curvature metrics shared across ADAM-like optimizers,
and integrates a custom optimization monitor for additional metrics.
"""

import warnings
from collections.abc import Callable
from typing import Any

import torch
from composer.core import Callback, State, Time, TimeUnit
from composer.loggers import Logger
from composer.utils import dist
from llmfoundry.registry import callbacks
from torch import Tensor

__all__ = ["OptimizerMonitor", "get_report_curvature"]


def get_report_curvature() -> Callable[[Tensor, str], dict[str, Tensor]]:
    """Get a function that computes curvature metrics per-parameter.

    The function tracks differences in parameters and gradients between iterations,
    thus it allocates extra memory. To avoid interfering with training, it saves the
    previous parameter and gradient values for each named parameter on the cpu, and
    only sends them to the device when needed for calculation, for 2 H100s this
    causes a drop in average samples per second from ~96 to ~94.

    Metrics are primarily taken from AdaBB: Adaptive Barzilai-Borwein Method
    for Convex Optimization: https://arxiv.org/pdf/2401.08024 specifically
    equations 3,4,5 from page 2. Metrics include:
    - param_diff_norm: L2 norm of the parameter difference between iterations
    - grad_diff_norm: L2 norm of the gradient difference between iterations
    - local_lipschitz: Calculated later, it is the ratio of the gradient difference
        norm to the parameter difference norm and estimates the local Lipschitz
        constant. Its supposed to capture the local curvature of the loss function
        (https://arxiv.org/pdf/2401.08024)
    - second_derivative_estimate: L2 norm of the elementwise ratio of gradient
        difference to parameter difference, recommended by a collaborator
    - second_to_first_derivative_estimate_ratio: L2 norm of the elementwise ratio
        of the second derivative estimate to the first derivative,
        recommended by a collaborator
    - long_bb: Barzilai-Borwein "long" step size
        (equation 3 from https://arxiv.org/pdf/2401.08024)
    - short_bb: Barzilai-Borwein (https://arxiv.org/pdf/2401.08024)
        "short" step size (equation 4 from https://arxiv.org/pdf/2401.08024)

    Returns
    -------
    Callable[[Tensor, str], dict[str, Tensor]]
        A function that, given a
        parameter `param` and its string name `name`, returns a dictionary with
        curvature metrics for that parameter. The function internally tracks
        previous parameter and gradient values for each named parameter.

    """
    prev_params: dict[str, Tensor] = {}
    prev_grads: dict[str, Tensor] = {}

    def report_curvature(param: Tensor, name: str) -> dict[str, Tensor]:
        # If there's no gradient, we skip reporting
        if param.grad is None:
            return {}

        # If we've never seen this parameter before, initialize storage and skip
        # this round (can't compute diffs without history).
        if name not in prev_params or name not in prev_grads:
            prev_params[name] = param.detach().clone().cpu()
            prev_grads[name] = param.grad.detach().clone().cpu()
            return {}

        # Compute diffs
        grad_diff = param.grad - prev_grads[name].to(
            device=param.grad.device,
            dtype=param.grad.dtype,
        )
        param_diff = param - prev_params[name].to(
            device=param.device,
            dtype=param.dtype,
        )

        param_diff_norm: Tensor = torch.linalg.vector_norm(param_diff)
        grad_diff_norm: Tensor = torch.linalg.vector_norm(grad_diff)

        # Barzilai-Borwein "long" step size
        long_bb = param_diff_norm**2.0 / torch.sum(torch.mul(grad_diff, param_diff))

        # Barzilai-Borwein "short" step size
        short_bb = torch.sum(torch.mul(grad_diff, param_diff)) / grad_diff_norm**2.0

        # Second-derivative estimate (elementwise ratio)
        second_derivative_estimate = grad_diff / param_diff

        # Ratio of first to second derivative
        second_to_first_derivative_ratio = torch.linalg.vector_norm(
            second_derivative_estimate / param.grad,
        )

        # Update stored values for next iteration
        prev_params[name] = param.detach().clone().cpu()
        prev_grads[name] = param.grad.detach().clone().cpu()

        return {
            f"curvature/param_diff_norm/{name}": param_diff_norm,
            f"curvature/grad_diff_norm/{name}": grad_diff_norm,
            f"curvature/long_bb/{name}": long_bb,
            f"curvature/short_bb/{name}": short_bb,
            f"curvature/l2_norm/second_to_first_derivative_estimate_ratio/{name}": (
                second_to_first_derivative_ratio
            ),
            f"curvature/l2_norm/second_derivative_estimate/{name}": (
                torch.linalg.vector_norm(second_derivative_estimate)
            ),
        }

    return report_curvature


def accumulate_curvature_metrics(
    optimizer_metrics: dict[str, torch.Tensor],
    metric_name_prefixes: tuple[str, ...] = (
        "curvature/param_diff_norm",
        "curvature/grad_diff_norm",
        "curvature/long_bb",
        "curvature/short_bb",
        "curvature/l2_norm/second_to_first_derivative_estimate_ratio",
        "curvature/l2_norm/second_derivative_estimate",
        "curvature/local_lipschitz",
    ),
) -> dict[str, Any]:
    """Accumulate curvature metrics from the optimizer_metrics dictionary.

    Parameters
    ----------
    optimizer_metrics : dict[str, torch.Tensor]
        The optimizer metrics dictionary.
    metric_name_prefixes : tuple[str, ...], optional
        The prefixes of the metrics to accumulate,
        by default ( "curvature/param_diff_norm", "curvature/grad_diff_norm",
        "curvature/long_bb", "curvature/short_bb",
        "curvature/l2_norm/second_to_first_derivative_estimate_ratio",
        "curvature/l2_norm/second_derivative_estimate", "curvature/local_lipschitz", )

    Returns
    -------
    dict[str, Any]
        The accumulated curvature metrics.

    """
    # Initialize accumulators
    sums_for_norms = {
        "curvature/param_diff_norm": 0.0,
        "curvature/grad_diff_norm": 0.0,
    }
    list_for_stats: dict = {
        "curvature/long_bb": [],
        "curvature/short_bb": [],
        "curvature/l2_norm/first_to_second_derivative_estimate_ratio": [],
        "curvature/l2_norm/second_derivative_estimate": [],
        "curvature/local_lipschitz": [],
    }

    # Accumulate
    for metric_name, metric_value in optimizer_metrics.items():
        for prefix in metric_name_prefixes:
            if metric_name.startswith(prefix):
                # Convert any Tensor to a Python float
                val = (
                    metric_value.item()
                    if isinstance(metric_value, torch.Tensor)
                    else metric_value
                )
                if prefix in sums_for_norms:
                    # Accumulate squared sum for these prefixes
                    sums_for_norms[prefix] += val**2
                elif prefix in list_for_stats:
                    # Just store the float
                    list_for_stats[prefix].append(val)

    return {
        "sums_for_norms": sums_for_norms,
        "list_for_stats": list_for_stats,
    }


def finalize_curvature_metrics(acc: dict[str, Any]) -> dict[str, float]:
    """Process the accumulated curvature metrics to produce final metrics.

    Parameters
    ----------
    acc : dict[str, Any]
        The accumulated curvature metrics.

    Returns
    -------
    dict[str, float]
        The final curvature metrics.

    """
    final_metrics = {}

    # 1) Compute global norms: sqrt of the sums
    sums_for_norms = acc["sums_for_norms"]
    for prefix, squared_sum in sums_for_norms.items():
        global_val = squared_sum**0.5
        final_metrics[f"{prefix}/global"] = global_val

    # 2) Reproduce the local_lipschitz/global calculation exactly as before
    grad_diff_g = final_metrics.get("curvature/grad_diff_norm/global", 0.0)
    param_diff_g = final_metrics.get("curvature/param_diff_norm/global", 1.0)

    final_metrics["curvature/local_lipschitz/global"] = (
        grad_diff_g / param_diff_g
        if param_diff_g != 0.0
        else float(
            "inf",
        )
    )

    # 3) Compute min/mean/median/max for each metric in list_for_stats
    list_for_stats = acc["list_for_stats"]
    for prefix, values in list_for_stats.items():
        if not values:
            # No values were accumulated for this prefix
            continue

        t = torch.tensor(values, dtype=torch.float32)
        # Filter out infinite values
        finite_t = t[torch.isfinite(t)]
        if finite_t.numel() == 0:
            continue

        final_metrics[f"{prefix}/mean"] = finite_t.mean().item()
        final_metrics[f"{prefix}/median"] = finite_t.median().item()
        final_metrics[f"{prefix}/min"] = finite_t.min().item()
        final_metrics[f"{prefix}/max"] = finite_t.max().item()

    return final_metrics


@callbacks.register_class("optimizer_monitor")
class OptimizerMonitor(Callback):
    """Compute and log the L2 norm of gradients.

    This is in addition to any optimizer-specific metrics implemented in the optimizer's
    `report_per_parameter_metrics` method.

    L2 norms are calculated after the reduction of gradients across GPUs. This function
    iterates over the parameters of the model and may cause a reduction in throughput
    while training large models. In order to ensure the correctness of the norm, this
    function should be called after gradient unscaling in cases where gradients are
    scaled.

    Examples
    --------
        >>> from composer import Trainer
        >>> from composer.callbacks import OptimizerMonitor
        >>> # constructing trainer object with this callback
        >>> trainer = Trainer(
        ...     model=model,
        ...     train_dataloader=train_dataloader,
        ...     eval_dataloader=eval_dataloader,
        ...     optimizers=optimizer,
        ...     max_duration="1ep",
        ...     callbacks=[OptimizerMonitor()],
        ... )

    Notes
    -----
    The metrics are logged by the :class:`.Logger` to the following keys as described
    below. `grad_l2_norm` and `layer_grad_l2_norm` are
    logged in addition to metrics logged by the optimizer's
    `report_per_parameter_metrics` method. For convenience we have listed the metrics
    logged by DecoupledAdamW below.

    +-----------------------------------------------+-----------------------------------------------------+
    | Key                                           | Logged data                                         |
    +===============================================+=====================================================+
    |                                               | L2 norm of the gradients of all parameters in       |
    | ``l2_norm/grad/global``                       | the model on the :attr:`.Event.AFTER_TRAIN_BATCH`   |
    |                                               | event.                                              |
    +-----------------------------------------------+-----------------------------------------------------+
    |                                               | Layer-wise L2 norms                                 |
    | ``l2_norm/grad/LAYER_NAME``                   |                                                     |
    |                                               |                                                     |
    +-----------------------------------------------+-----------------------------------------------------+
    |                                               | Layer-wise L2 norms of Adam first moment after      |
    | ``l2_norm/moment/LAYER_NAME``                 |  calling optimizer step.                            |
    |                                               |                                                     |
    +-----------------------------------------------
    +-----------------------------------------------------+
    |                                               | Layer-wise L2 norms of Adam second moment after      |
    | ``l2_norm/moment2/LAYER_NAME``                 |  calling optimizer step.                            |
    |                                               |                                                     |
    +-----------------------------------------------
    +-----------------------------------------------------+
    |                                               | Minimum element of the Adam second moment after      |
    | ``min/moment2/LAYER_NAME``                 |  calling optimizer step.                            |
    |                                               |                                                     |
    +-----------------------------------------------
    +-----------------------------------------------------+
    |                                               | Maximum element of the Adam second moment after      |
    | ``max/moment2/LAYER_NAME``                 |  calling optimizer step.                            |
    |                                               |                                                     |
    +-----------------------------------------------
    +-----------------------------------------------------+
    |                                               | Curvature metrics derived from the gradient and model parameters      |
    | ``curvature/METRIC_NAME``                 |  calling optimizer step.                            |
    |                                               |                                                     |
    +-----------------------------------------------
    +-----------------------------------------------------+
    |                                               | Layer-wise L2 norms of parameter weights            |
    | ``l2_norm/param/LAYER_NAME``                  |                                                     |
    |                                               |                                                     |
    +-----------------------------------------------+-----------------------------------------------------+
    |                                               | Layer-wise L2 norms of the step                     |
    | ``l2_norm/update/LAYER_NAME``                 |                                                     |
    |                                               |                                                     |
    +-----------------------------------------------+-----------------------------------------------------+

    """  # noqa: E501

    def __init__(
        self,
        interval: int | str | Time = "10ba",
        *,
        only_global: bool = False,
        log_optimizer_metrics: bool = True,
        report_curvature: bool = False,
    ) -> None:
        """Initialise the optimzier monitor.

        Parameters
        ----------
        only_global : bool, optional
            Whether only global parameters should be taken, by default False
        log_optimizer_metrics : bool, optional
            Whether optimizer metrics should be logged, by default True
        interval : Union[int, str, Time], optional
            The interval in which to log, by default "10ba"
        report_curvature : bool, optional
            Whether to report the curvature metrics, by default False

        Raises
        ------
        ValueError
            Invalid time unit for parameter interval if it is not in the batch or epoch.

        """
        self.log_optimizer_metrics = log_optimizer_metrics
        self.only_global = only_global

        # Check that the interval timestring is parsable and convert into time object
        if isinstance(interval, int):
            self.interval = Time(interval, TimeUnit.BATCH)
        elif isinstance(interval, str):
            self.interval = Time.from_timestring(interval)
        elif isinstance(interval, Time):
            self.interval = interval

        if (
            self.interval.unit == TimeUnit.BATCH
            and self.interval < Time.from_timestring("10ba")
        ):
            warnings.warn(
                f"Currently the ActivationMonitor`s interval is set to {self.interval} "
                f"which is below our recommended value of 10ba. We recommend you raise "
                f"the interval to at least 10ba, as the activation monitor adds extra overhead "  # noqa: E501
                f"and decreases throughput.",
                stacklevel=2,
            )

        # Verify that the interval has supported units
        if self.interval.unit not in {TimeUnit.BATCH, TimeUnit.EPOCH}:
            msg = f"Invalid time unit for parameter interval: {self.interval.unit}"
            raise ValueError(
                msg,
            )
        self.report_curvature = report_curvature

    def batch_end(self, state: State, logger: Logger) -> None:  # noqa: C901, PLR0915, PLR0912
        """Calculate the statistics at the end of the batch.

        Parameters
        ----------
        state : State
            The state of theTsrainer object
        logger : Logger
            The callback to log the optimizer metrics to the Trainer object.

        """
        current_time_value = state.timestamp.get(self.interval.unit).value

        if current_time_value % self.interval.value != 0:
            return

        optimizer_metrics: dict = {}

        for name, p in state.model.named_parameters():
            if p.grad is not None and p.requires_grad:
                metric_reporter: Callable[[Any, Any, Any], dict] = getattr(  # type: ignore[assignment]
                    state.optimizers[0],
                    "report_per_parameter_metrics",
                    None,
                )  # type: ignore[reportAssignmentType]
                if callable(metric_reporter) and self.log_optimizer_metrics:
                    optimizer_metrics.update(
                        metric_reporter(p, name, optimizer_metrics),
                    )

                # Always log grad norm as a default metric if it's not specified
                if f"l2_norm/grad/{name}" not in optimizer_metrics:
                    param_grad_norm = torch.linalg.vector_norm(p.grad)
                    optimizer_metrics[f"l2_norm/grad/{name}"] = param_grad_norm

        if (
            state.fsdp_enabled
            and dist.get_world_size() > 0
            and self.log_optimizer_metrics
        ):
            # If FSDP is enabled, the optimizer state lives on different ranks and must
            # be reduced and combined before we can compute metrics.
            # Each metric has a different way of being reduced, so the optimizer is
            # responsible for implementing the reduction process.
            # It occurs first via a pre-reduce, where the metric on each rank is
            # modified and prepared then an all-reduce where the modified metric on each
            # rank is combined into the correct metric across all ranks.
            #
            # For example, L2 norms are squared on each rank before we apply
            # all_reduce(SUM) and take the sqrt on each rank
            pre_reduce_metrics: Callable[[Any], dict] = getattr(  # type: ignore[assignment]
                state.optimizers[0],
                "pre_reduce_metrics",
                None,
            )  # type: ignore[reportAssignmentType]
            if callable(pre_reduce_metrics) and self.log_optimizer_metrics:
                optimizer_metrics = pre_reduce_metrics(optimizer_metrics)

            dist_reduce_metrics: Callable[[Any], dict] = getattr(  # type: ignore[assignment]
                state.optimizers[0],
                "dist_reduce_metrics",
                None,
            )  # type: ignore[reportAssignmentType]
            if callable(dist_reduce_metrics) and self.log_optimizer_metrics:
                optimizer_metrics = dist_reduce_metrics(optimizer_metrics)

        # Dynamically aggregate all metric names found in optimizer_metrics
        agg_types = ["l2_norm", "min", "max"]
        agg_dict = {}
        metric_parts_required = 2
        for metric in optimizer_metrics:
            parts = metric.split("/")
            if len(parts) < metric_parts_required:
                continue
            agg_type, metric_name = parts[0], parts[1]
            if agg_type not in agg_types:
                continue
            key = (agg_type, metric_name)
            if agg_type == "l2_norm" and key not in agg_dict:
                agg_dict[key] = 0.0
            elif agg_type == "min" and key not in agg_dict:
                agg_dict[key] = float("inf")
            elif agg_type == "max" and key not in agg_dict:
                agg_dict[key] = float("-inf")
        for metric in optimizer_metrics:
            parts = metric.split("/")
            if len(parts) < metric_parts_required:
                continue
            agg_type, metric_name = parts[0], parts[1]
            if agg_type not in agg_types:
                continue
            key = (agg_type, metric_name)
            value = optimizer_metrics[metric]
            if agg_type == "l2_norm":
                agg_dict[key] += value**2
            elif agg_type == "min":
                agg_dict[key] = min(agg_dict[key], value)
            elif agg_type == "max":
                agg_dict[key] = max(agg_dict[key], value)

        # Report all aggregated metrics as agg_type/metric_name/global
        for (agg_type, metric_name), agg_value in agg_dict.items():
            if agg_type == "l2_norm":
                optimizer_metrics[f"{agg_type}/{metric_name}/global"] = agg_value**0.5
            else:
                optimizer_metrics[f"{agg_type}/{metric_name}/global"] = agg_value

        # If only_global is set, remove all non-global metrics
        if self.only_global:
            optimizer_metrics = {
                k: v for k, v in optimizer_metrics.items() if k.endswith("/global")
            }

        curvature_acc = {}
        # Curvature metrics
        # trying to obtain them here
        # has minimal cost if the optimizer
        # does not report them
        if self.report_curvature:
            curvature_acc = accumulate_curvature_metrics(optimizer_metrics)

        if self.report_curvature:
            curvature_stats = finalize_curvature_metrics(curvature_acc)
            optimizer_metrics.update(curvature_stats)

        for metric in optimizer_metrics:
            if isinstance(optimizer_metrics[metric], torch.Tensor):
                optimizer_metrics[metric] = optimizer_metrics[metric].item()
        logger.log_metrics(optimizer_metrics)
