# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Specification for workloads, which defines the input tensors for a kernel."""

from __future__ import annotations

from typing import Literal, Optional, Union

from pydantic import Field, model_validator

from .base_model import BaseModelWithDocstrings, NonEmptyString, NonNegativeInt


class RandomInput(BaseModelWithDocstrings):
    """Random input generation descriptor.

    Represents a specification for generating random tensor input data
    during workload execution and benchmarking.
    """

    type: Literal["random"] = "random"
    """The input type identifier for random data generation."""


class ScalarInput(BaseModelWithDocstrings):
    """Scalar literal input specification.

    Represents a scalar value (integer, float, or boolean) that will be
    used as a direct input parameter to the computational workload.
    """

    type: Literal["scalar"] = "scalar"
    """The input type identifier for scalar values."""
    value: Union[int, float, bool]
    """The scalar value to be used as input. Must be int, float, or bool."""


class SafetensorsInput(BaseModelWithDocstrings):
    """Input specification for data loaded from safetensors files.

    Represents tensor data that will be loaded from a safetensors file
    using a specific tensor key within that file.
    """

    type: Literal["safetensors"] = "safetensors"
    """The input type identifier for safetensors data."""
    path: NonEmptyString
    """Path to the safetensors file containing the tensor data. The path is relative to the root
    path of the TraceSet."""
    tensor_key: NonEmptyString
    """Key identifier for the specific tensor within the safetensors file."""


class CustomInput(BaseModelWithDocstrings):
    """Custom input specification.

    Represents an that is generated from custom reference code. If a single input is this type, all inputs should be this type. Requires Definition.custom_inputs_entrypoint to be set.
    """

    type: Literal["custom"] = "custom"
    """The input type identifier for inputs generated by the reference code."""


InputSpec = Union[RandomInput, SafetensorsInput, ScalarInput, CustomInput]
"""Union type representing all possible input specification types."""


class ToleranceSpec(BaseModelWithDocstrings):
    """Tolerance specification for a workload.

    Specifies the conditions for a particular problem to be deemed numerically correct against the reference implementation.
    """

    max_atol: float = Field(default=1e-2)
    """The maximum absolute error allowed for the problem."""
    max_rtol: float = Field(default=1e-2)
    """The maximum relative error allowed for the problem."""
    required_matched_ratio: float = Field(default=0.99)
    """The ratio of elements that must pass the correctness bounds to be considered correct."""
    max_error_cap: Optional[float] = Field(default=None)
    """Hard ceiling on maximum absolute error. If set, correctness fails when any
    element's absolute error exceeds this cap, regardless of matched ratio."""
    allow_negative_inf: bool = Field(default=False)
    """When True, matching -inf values in both output and reference are treated as
    correct and excluded from error computation. Positions where only one tensor
    has -inf still fail. +inf and NaN are unaffected by this flag."""


class Workload(BaseModelWithDocstrings):
    """Concrete workload configuration for benchmarking.

    Defines a specific instance of a computational workload with concrete
    values for all variable axes and specifications for all input data.
    This represents an executable configuration that can be benchmarked.
    """

    axes: dict[str, NonNegativeInt]
    """Dictionary mapping axis names to their concrete integer values. All values must be
    positive."""
    inputs: dict[str, InputSpec]
    """Dictionary mapping input names to their data specifications."""
    uuid: NonEmptyString
    """Unique identifier for this specific workload configuration."""
    tolerance: ToleranceSpec = Field(default=ToleranceSpec())
    """Tolerance specification for the workload."""

    @model_validator(mode="after")
    def _validate_inputs(self) -> Workload:
        custom_inputs = [
            name
            for name, input in self.inputs.items()
            if isinstance(input, CustomInput)
        ]
        non_custom_inputs = [
            name
            for name, input in self.inputs.items()
            if not isinstance(input, CustomInput)
        ]
        if len(custom_inputs) > 0 and len(non_custom_inputs) > 0:
            raise ValueError(
                f"A workload cannot have both custom and non-custom inputs. Custom: {custom_inputs}. Non-custom: {non_custom_inputs}"
            )
        return self

    def get_scalar_inputs(self) -> dict[str, int]:
        return {
            name: input.value
            for name, input in self.inputs.items()
            if isinstance(input, ScalarInput)
        }
