# Copyright (c) ONNX Project Contributors

# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

from typing import Any

import numpy as np

from onnx.reference.op_run import OpRun


def _concat_from_sequence(seq: list[Any], axis: int, new_axis: int = 0) -> np.ndarray:
    if new_axis == 1:
        if axis == -1:
            seq2 = [s[..., np.newaxis] for s in seq]
            res = np.concatenate(seq2, axis=-1)
        else:
            seq2 = [np.expand_dims(s, axis) for s in seq]
            res = np.concatenate(seq2, axis=axis)
    else:
        res = np.concatenate(seq, axis=axis)
    return res


class ConcatFromSequence(OpRun):
    def _run(self, seq, axis=None, new_axis=None):
        if seq is None:
            raise RuntimeError("A sequence cannot be null.")
        res = _concat_from_sequence(seq, axis, new_axis=new_axis)
        return (res,)