import pytest
from unittest.mock import patch, Mock
from _test_commons import _load_pytorch_transformer_model
from onnxruntime.training import amp, checkpoint, optim, orttrainer, _checkpoint_storage
import numpy as np
import onnx
import torch

# Helper functions


def _create_trainer(zero_enabled=False):
    """Cerates a simple ORTTrainer for ORTTrainer functional tests"""

    device = "cuda"
    optim_config = optim.LambConfig(lr=0.1)
    opts = {"device": {"id": device}, "debug": {"deterministic_compute": True}}
    if zero_enabled:
        opts["distributed"] = {
            "world_rank": 0,
            "world_size": 1,
            "horizontal_parallel_size": 1,
            "data_parallel_size": 1,
            "allreduce_post_accumulation": True,
            "deepspeed_zero_optimization": {"stage": 1},
        }
    model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(device)
    trainer = orttrainer.ORTTrainer(
        model, model_desc, optim_config, loss_fn=loss_fn, options=orttrainer.ORTTrainerOptions(opts)
    )

    return trainer


class _training_session_mock(object):
    """Mock object for the ORTTrainer _training_session member"""

    def __init__(self, model_states, optimizer_states, partition_info):
        self.model_states = model_states
        self.optimizer_states = optimizer_states
        self.partition_info = partition_info

    def get_model_state(self, include_mixed_precision_weights=False):
        return self.model_states

    def get_optimizer_state(self):
        return self.optimizer_states

    def get_partition_info_map(self):
        return self.partition_info


def _get_load_state_dict_strict_error_arguments():
    """Return a list of tuples that can be used as parameters for test_load_state_dict_errors_when_model_key_missing

    Construct a list of tuples (training_session_state_dict, input_state_dict, error_arguments)
    The load_state_dict function will compare the two state dicts (training_session_state_dict, input_state_dict) and
    throw a runtime error with the missing/unexpected keys. The error arguments capture these missing/unexpected keys.
    """

    training_session_state_dict = {
        "model": {"full_precision": {"a": np.arange(5), "b": np.arange(7)}},
        "optimizer": {
            "a": {"Moment_1": np.arange(5), "Moment_2": np.arange(7)},
            "shared_optimizer_state": {"step": np.arange(5)},
        },
    }

    # input state dictionaries
    precision_key_missing = {"model": {}, "optimizer": {}}
    precision_key_unexpected = {"model": {"full_precision": {}, "mixed_precision": {}}, "optimizer": {}}
    model_state_key_missing = {"model": {"full_precision": {}}, "optimizer": {}}
    model_state_key_unexpected = {"model": {"full_precision": {"a": 2, "b": 3, "c": 4}}, "optimizer": {}}
    optimizer_model_state_key_missing = {"model": {"full_precision": {"a": 2, "b": 3}}, "optimizer": {}}
    optimizer_model_state_key_unexpected = {
        "model": {"full_precision": {"a": 2, "b": 3}},
        "optimizer": {"a": {}, "shared_optimizer_state": {}, "b": {}},
    }
    optimizer_state_key_missing = {
        "model": {"full_precision": {"a": 2, "b": 3}},
        "optimizer": {"a": {}, "shared_optimizer_state": {"step": np.arange(5)}},
    }
    optimizer_state_key_unexpected = {
        "model": {"full_precision": {"a": 2, "b": 3}},
        "optimizer": {
            "a": {"Moment_1": np.arange(5), "Moment_2": np.arange(7)},
            "shared_optimizer_state": {"step": np.arange(5), "another_step": np.arange(1)},
        },
    }

    input_arguments = [
        (training_session_state_dict, precision_key_missing, ["full_precision"]),
        (training_session_state_dict, precision_key_unexpected, ["mixed_precision"]),
        (training_session_state_dict, model_state_key_missing, ["a", "b"]),
        (training_session_state_dict, model_state_key_unexpected, ["c"]),
        (training_session_state_dict, optimizer_model_state_key_missing, ["a", "shared_optimizer_state"]),
        (training_session_state_dict, optimizer_model_state_key_unexpected, ["b"]),
        (training_session_state_dict, optimizer_state_key_missing, ["Moment_1", "Moment_2"]),
        (training_session_state_dict, optimizer_state_key_unexpected, ["another_step"]),
    ]

    return input_arguments


# Tests


def test_empty_state_dict_when_training_session_uninitialized():
    trainer = _create_trainer()
    with pytest.warns(UserWarning) as user_warning:
        state_dict = trainer.state_dict()

    assert len(state_dict.keys()) == 0
    assert (
        user_warning[0].message.args[0] == "ONNX Runtime training session is not initialized yet. "
        "Please run train_step or eval_step at least once before calling ORTTrainer.state_dict()."
    )


@patch("onnx.ModelProto")
def test_training_session_provides_empty_model_states(onnx_model_mock):
    trainer = _create_trainer()
    training_session_mock = _training_session_mock({}, {}, {})
    trainer._training_session = training_session_mock
    trainer._onnx_model = onnx_model_mock()

    state_dict = trainer.state_dict()
    assert len(state_dict["model"].keys()) == 0


@patch("onnx.ModelProto")
def test_training_session_provides_model_states(onnx_model_mock):
    trainer = _create_trainer()
    model_states = {"full_precision": {"a": np.arange(5), "b": np.arange(7)}}
    training_session_mock = _training_session_mock(model_states, {}, {})
    trainer._training_session = training_session_mock
    trainer._onnx_model = onnx_model_mock()

    state_dict = trainer.state_dict()
    assert (state_dict["model"]["full_precision"]["a"] == np.arange(5)).all()
    assert (state_dict["model"]["full_precision"]["b"] == np.arange(7)).all()


@patch("onnx.ModelProto")
def test_training_session_provides_model_states_pytorch_format(onnx_model_mock):
    trainer = _create_trainer()
    model_states = {"full_precision": {"a": np.arange(5), "b": np.arange(7)}}
    training_session_mock = _training_session_mock(model_states, {}, {})
    trainer._training_session = training_session_mock
    trainer._onnx_model = onnx_model_mock()

    state_dict = trainer.state_dict(pytorch_format=True)
    assert torch.all(torch.eq(state_dict["a"], torch.tensor(np.arange(5))))
    assert torch.all(torch.eq(state_dict["b"], torch.tensor(np.arange(7))))


@patch("onnx.ModelProto")
def test_onnx_graph_provides_frozen_model_states(onnx_model_mock):
    trainer = _create_trainer()
    model_states = {"full_precision": {"a": np.arange(5), "b": np.arange(7)}}
    training_session_mock = _training_session_mock(model_states, {}, {})
    trainer._training_session = training_session_mock
    trainer._onnx_model = onnx_model_mock()
    trainer.options.utils.frozen_weights = ["a_frozen_weight", "a_float16_weight"]
    trainer._onnx_model.graph.initializer = [
        onnx.numpy_helper.from_array(np.array([1, 2, 3], dtype=np.float32), "a_frozen_weight"),
        onnx.numpy_helper.from_array(np.array([4, 5, 6], dtype=np.float32), "a_non_fronzen_weight"),
        onnx.numpy_helper.from_array(np.array([7, 8, 9], dtype=np.float16), "a_float16_weight"),
    ]

    state_dict = trainer.state_dict()
    assert (state_dict["model"]["full_precision"]["a"] == np.arange(5)).all()
    assert (state_dict["model"]["full_precision"]["b"] == np.arange(7)).all()
    assert (state_dict["model"]["full_precision"]["a_frozen_weight"] == np.array([1, 2, 3], dtype=np.float32)).all()
    assert "a_non_fronzen_weight" not in state_dict["model"]["full_precision"]
    assert (state_dict["model"]["full_precision"]["a_float16_weight"] == np.array([7, 8, 9], dtype=np.float32)).all()


@patch("onnx.ModelProto")
def test_training_session_provides_empty_optimizer_states(onnx_model_mock):
    trainer = _create_trainer()
    training_session_mock = _training_session_mock({}, {}, {})
    trainer._training_session = training_session_mock
    trainer._onnx_model = onnx_model_mock()

    state_dict = trainer.state_dict()
    assert len(state_dict["optimizer"].keys()) == 0


@patch("onnx.ModelProto")
def test_training_session_provides_optimizer_states(onnx_model_mock):
    trainer = _create_trainer()
    optimizer_states = {
        "model_weight": {"Moment_1": np.arange(5), "Moment_2": np.arange(7)},
        "shared_optimizer_state": {"step": np.arange(1)},
    }
    training_session_mock = _training_session_mock({}, optimizer_states, {})
    trainer._training_session = training_session_mock
    trainer._onnx_model = onnx_model_mock()

    state_dict = trainer.state_dict()
    assert (state_dict["optimizer"]["model_weight"]["Moment_1"] == np.arange(5)).all()
    assert (state_dict["optimizer"]["model_weight"]["Moment_2"] == np.arange(7)).all()
    assert (state_dict["optimizer"]["shared_optimizer_state"]["step"] == np.arange(1)).all()


@patch("onnx.ModelProto")
def test_training_session_provides_optimizer_states_pytorch_format(onnx_model_mock):
    trainer = _create_trainer()
    model_states = {"full_precision": {"a": np.arange(5), "b": np.arange(7)}}
    optimizer_states = {
        "model_weight": {"Moment_1": np.arange(5), "Moment_2": np.arange(7)},
        "shared_optimizer_state": {"step": np.arange(1)},
    }
    training_session_mock = _training_session_mock(model_states, optimizer_states, {})
    trainer._training_session = training_session_mock
    trainer._onnx_model = onnx_model_mock()

    state_dict = trainer.state_dict(pytorch_format=True)
    assert "optimizer" not in state_dict


@patch("onnx.ModelProto")
def test_training_session_provides_empty_partition_info_map(onnx_model_mock):
    trainer = _create_trainer(zero_enabled=True)
    training_session_mock = _training_session_mock({}, {}, {})
    trainer._training_session = training_session_mock
    trainer._onnx_model = onnx_model_mock()

    state_dict = trainer.state_dict()
    assert len(state_dict["partition_info"].keys()) == 0


@patch("onnx.ModelProto")
def test_training_session_provides_partition_info_map(onnx_model_mock):
    trainer = _create_trainer(zero_enabled=True)
    partition_info = {"a": {"original_dim": [1, 2, 3]}}
    training_session_mock = _training_session_mock({}, {}, partition_info)
    trainer._training_session = training_session_mock
    trainer._onnx_model = onnx_model_mock()

    state_dict = trainer.state_dict()
    assert state_dict["partition_info"]["a"]["original_dim"] == [1, 2, 3]


@patch("onnx.ModelProto")
def test_training_session_provides_all_states(onnx_model_mock):
    trainer = _create_trainer(zero_enabled=True)
    model_states = {"full_precision": {"a": np.arange(5), "b": np.arange(7)}}
    optimizer_states = {
        "model_weight": {"Moment_1": np.arange(5), "Moment_2": np.arange(7)},
        "shared_optimizer_state": {"step": np.arange(1)},
    }
    partition_info = {"a": {"original_dim": [1, 2, 3]}}
    training_session_mock = _training_session_mock(model_states, optimizer_states, partition_info)
    trainer._training_session = training_session_mock
    trainer._onnx_model = onnx_model_mock()

    state_dict = trainer.state_dict()
    assert (state_dict["model"]["full_precision"]["a"] == np.arange(5)).all()
    assert (state_dict["model"]["full_precision"]["b"] == np.arange(7)).all()
    assert (state_dict["optimizer"]["model_weight"]["Moment_1"] == np.arange(5)).all()
    assert (state_dict["optimizer"]["model_weight"]["Moment_2"] == np.arange(7)).all()
    assert (state_dict["optimizer"]["shared_optimizer_state"]["step"] == np.arange(1)).all()
    assert state_dict["partition_info"]["a"]["original_dim"] == [1, 2, 3]


def test_load_state_dict_holds_when_training_session_not_initialized():
    trainer = _create_trainer()
    state_dict = {
        "model": {"full_precision": {"a": np.arange(5), "b": np.arange(7)}},
        "optimizer": {
            "a": {"Moment_1": np.arange(5), "Moment_2": np.arange(7)},
            "shared_optimizer_state": {"step": np.arange(5)},
        },
    }
    assert not trainer._load_state_dict
    state_dict = trainer.load_state_dict(state_dict)
    assert trainer._load_state_dict


@pytest.mark.parametrize(
    "state_dict, input_state_dict, error_key",
    [
        (
            {"model": {}, "optimizer": {}},
            {"model": {}, "optimizer": {}, "trainer_options": {"optimizer_name": "LambOptimizer"}},
            "train_step_info",
        ),
        (
            {"optimizer": {}, "train_step_info": {"optimization_step": 0, "step": 0}},
            {
                "optimizer": {},
                "trainer_options": {"optimizer_name": "LambOptimizer"},
                "train_step_info": {"optimization_step": 0, "step": 0},
            },
            "model",
        ),
        (
            {"model": {}, "train_step_info": {"optimization_step": 0, "step": 0}},
            {
                "model": {},
                "trainer_options": {"optimizer_name": "LambOptimizer"},
                "train_step_info": {"optimization_step": 0, "step": 0},
            },
            "optimizer",
        ),
    ],
)
def test_load_state_dict_warns_when_model_optimizer_key_missing(state_dict, input_state_dict, error_key):
    trainer = _create_trainer()
    trainer._training_session = _training_session_mock({}, {}, {})
    trainer.state_dict = Mock(return_value=state_dict)
    trainer._update_onnx_model_initializers = Mock()
    trainer._init_session = Mock()
    with patch("onnx.ModelProto") as onnx_model_mock:
        trainer._onnx_model = onnx_model_mock()
        trainer._onnx_model.graph.initializer = []
        with pytest.warns(UserWarning) as user_warning:
            trainer.load_state_dict(input_state_dict)

    assert user_warning[0].message.args[0] == "Missing key: {} in state_dict".format(error_key)


@pytest.mark.parametrize("state_dict, input_state_dict, error_keys", _get_load_state_dict_strict_error_arguments())
def test_load_state_dict_errors_when_state_dict_mismatch(state_dict, input_state_dict, error_keys):
    trainer = _create_trainer()
    trainer._training_session = _training_session_mock({}, {}, {})
    trainer.state_dict = Mock(return_value=state_dict)
    with pytest.raises(RuntimeError) as runtime_error:
        trainer.load_state_dict(input_state_dict)

    assert any(key in str(runtime_error.value) for key in error_keys)


@patch("onnx.ModelProto")
def test_load_state_dict_loads_the_states_and_inits_training_session(onnx_model_mock):
    trainer = _create_trainer()
    training_session_state_dict = {
        "model": {"full_precision": {"a": np.arange(5), "b": np.arange(7)}},
        "optimizer": {
            "a": {"Moment_1": np.arange(5), "Moment_2": np.arange(7)},
            "shared_optimizer_state": {"step": np.arange(1)},
        },
    }

    input_state_dict = {
        "model": {"full_precision": {"a": np.array([1, 2]), "b": np.array([3, 4])}},
        "optimizer": {
            "a": {"Moment_1": np.array([5, 6]), "Moment_2": np.array([7, 8])},
            "shared_optimizer_state": {"step": np.array([9])},
        },
        "trainer_options": {"optimizer_name": "LambOptimizer"},
    }
    trainer._training_session = _training_session_mock({}, {}, {})
    trainer.state_dict = Mock(return_value=training_session_state_dict)
    trainer._onnx_model = onnx_model_mock()
    trainer._onnx_model.graph.initializer = [
        onnx.numpy_helper.from_array(np.arange(20, dtype=np.float32), "a"),
        onnx.numpy_helper.from_array(np.arange(25, dtype=np.float32), "b"),
    ]
    trainer._update_onnx_model_initializers = Mock()
    trainer._init_session = Mock()

    trainer.load_state_dict(input_state_dict)

    loaded_initializers, _ = trainer._update_onnx_model_initializers.call_args
    state_dict_to_load, _ = trainer._init_session.call_args

    assert "a" in loaded_initializers[0]
    assert (loaded_initializers[0]["a"] == np.array([1, 2])).all()
    assert "b" in loaded_initializers[0]
    assert (loaded_initializers[0]["b"] == np.array([3, 4])).all()

    assert (state_dict_to_load[0]["a"]["Moment_1"] == np.array([5, 6])).all()
    assert (state_dict_to_load[0]["a"]["Moment_2"] == np.array([7, 8])).all()
    assert (state_dict_to_load[0]["shared_optimizer_state"]["step"] == np.array([9])).all()


@patch("onnxruntime.training._checkpoint_storage.save")
def test_save_checkpoint_calls_checkpoint_storage_save(save_mock):
    trainer = _create_trainer()
    state_dict = {"model": {}, "optimizer": {}}
    trainer.state_dict = Mock(return_value=state_dict)

    trainer.save_checkpoint("abc")

    save_args, _ = save_mock.call_args
    assert "model" in save_args[0]
    assert not bool(save_args[0]["model"])
    assert "optimizer" in save_args[0]
    assert not bool(save_args[0]["optimizer"])
    assert save_args[1] == "abc"


@patch("onnxruntime.training._checkpoint_storage.save")
def test_save_checkpoint_exclude_optimizer_states(save_mock):
    trainer = _create_trainer()
    state_dict = {"model": {}, "optimizer": {}}
    trainer.state_dict = Mock(return_value=state_dict)

    trainer.save_checkpoint("abc", include_optimizer_states=False)

    save_args, _ = save_mock.call_args
    assert "model" in save_args[0]
    assert not bool(save_args[0]["model"])
    assert "optimizer" not in save_args[0]
    assert save_args[1] == "abc"


@patch("onnxruntime.training._checkpoint_storage.save")
def test_save_checkpoint_user_dict(save_mock):
    trainer = _create_trainer()
    state_dict = {"model": {}, "optimizer": {}}
    trainer.state_dict = Mock(return_value=state_dict)

    trainer.save_checkpoint("abc", user_dict={"abc": np.arange(4)})

    save_args, _ = save_mock.call_args
    assert "user_dict" in save_args[0]
    assert save_args[0]["user_dict"] == _checkpoint_storage.to_serialized_hex({"abc": np.arange(4)})


@patch("onnxruntime.training._checkpoint_storage.load")
@patch("onnxruntime.training.checkpoint.aggregate_checkpoints")
def test_load_checkpoint(aggregate_checkpoints_mock, load_mock):
    trainer = _create_trainer()
    trainer_options = {
        "mixed_precision": np.bool_(False),
        "world_rank": np.int64(0),
        "world_size": np.int64(1),
        "horizontal_parallel_size": np.int64(1),
        "data_parallel_size": np.int64(1),
        "zero_stage": np.int64(0),
    }
    state_dict = {
        "model": {},
        "optimizer": {},
        "trainer_options": {
            "mixed_precision": np.bool_(False),
            "world_rank": np.int64(0),
            "world_size": np.int64(1),
            "horizontal_parallel_size": np.int64(1),
            "data_parallel_size": np.int64(1),
            "zero_stage": np.int64(0),
        },
    }
    trainer.load_state_dict = Mock()

    load_mock.side_effect = [trainer_options, state_dict]
    trainer.load_checkpoint("abc")

    args_list = load_mock.call_args_list
    load_args, load_kwargs = args_list[0]
    assert load_args[0] == "abc"
    assert load_kwargs["key"] == "trainer_options"
    load_args, load_kwargs = args_list[1]
    assert load_args[0] == "abc"
    assert "key" not in load_kwargs
    assert not aggregate_checkpoints_mock.called


@patch("onnxruntime.training._checkpoint_storage.load")
@patch("onnxruntime.training.checkpoint.aggregate_checkpoints")
@pytest.mark.parametrize(
    "trainer_options",
    [
        {
            "mixed_precision": np.bool_(False),
            "world_rank": np.int64(0),
            "world_size": np.int64(4),
            "horizontal_parallel_size": np.int64(1),
            "data_parallel_size": np.int64(4),
            "zero_stage": np.int64(1),
        },
        {
            "mixed_precision": np.bool_(True),
            "world_rank": np.int64(0),
            "world_size": np.int64(1),
            "horizontal_parallel_size": np.int64(1),
            "data_parallel_size": np.int64(1),
            "zero_stage": np.int64(1),
        },
        {
            "mixed_precision": np.bool_(True),
            "world_rank": np.int64(0),
            "world_size": np.int64(1),
            "horizontal_parallel_size": np.int64(1),
            "data_parallel_size": np.int64(1),
            "zero_stage": np.int64(1),
        },
    ],
)
def test_load_checkpoint_aggregation_required_zero_enabled(aggregate_checkpoints_mock, load_mock, trainer_options):
    trainer = _create_trainer()
    trainer.load_state_dict = Mock()

    load_mock.side_effect = [trainer_options]
    trainer.load_checkpoint("abc")

    args_list = load_mock.call_args_list
    load_args, load_kwargs = args_list[0]
    assert load_args[0] == "abc"
    assert load_kwargs["key"] == "trainer_options"
    assert aggregate_checkpoints_mock.called
    call_args, _ = aggregate_checkpoints_mock.call_args
    assert call_args[0] == tuple(["abc"])


@patch("onnxruntime.training._checkpoint_storage.load")
@patch("onnxruntime.training.checkpoint.aggregate_checkpoints")
def test_load_checkpoint_user_dict(aggregate_checkpoints_mock, load_mock):
    trainer = _create_trainer()
    trainer_options = {
        "mixed_precision": np.bool_(False),
        "world_rank": np.int64(0),
        "world_size": np.int64(1),
        "horizontal_parallel_size": np.int64(1),
        "data_parallel_size": np.int64(1),
        "zero_stage": np.int64(0),
    }
    state_dict = {
        "model": {},
        "optimizer": {},
        "trainer_options": {
            "mixed_precision": np.bool_(False),
            "world_rank": np.int64(0),
            "world_size": np.int64(1),
            "horizontal_parallel_size": np.int64(1),
            "data_parallel_size": np.int64(1),
            "zero_stage": np.int64(0),
        },
        "user_dict": _checkpoint_storage.to_serialized_hex({"array": torch.tensor(np.arange(5))}),
    }
    trainer.load_state_dict = Mock()

    load_mock.side_effect = [trainer_options, state_dict]
    user_dict = trainer.load_checkpoint("abc")

    assert torch.all(torch.eq(user_dict["array"], torch.tensor(np.arange(5))))


@patch("onnxruntime.training._checkpoint_storage.load")
def test_checkpoint_aggregation(load_mock):
    trainer_options1 = {
        "mixed_precision": np.bool_(False),
        "world_rank": np.int64(0),
        "world_size": np.int64(2),
        "horizontal_parallel_size": np.int64(1),
        "data_parallel_size": np.int64(2),
        "zero_stage": np.int64(1),
        "optimizer_name": b"Adam",
    }
    trainer_options2 = {
        "mixed_precision": np.bool_(False),
        "world_rank": np.int64(1),
        "world_size": np.int64(2),
        "horizontal_parallel_size": np.int64(1),
        "data_parallel_size": np.int64(2),
        "zero_stage": np.int64(1),
        "optimizer_name": b"Adam",
    }

    state_dict1 = {
        "model": {"full_precision": {"optimizer_sharded": np.array([1, 2, 3]), "non_sharded": np.array([11, 22, 33])}},
        "optimizer": {
            "optimizer_sharded": {
                "Moment_1": np.array([9, 8, 7]),
                "Moment_2": np.array([99, 88, 77]),
                "Step": np.array([5]),
            },
            "non_sharded": {
                "Moment_1": np.array([666, 555, 444]),
                "Moment_2": np.array([6666, 5555, 4444]),
                "Step": np.array([55]),
            },
        },
        "trainer_options": {
            "mixed_precision": np.bool_(False),
            "world_rank": np.int64(0),
            "world_size": np.int64(1),
            "horizontal_parallel_size": np.int64(1),
            "data_parallel_size": np.int64(1),
            "zero_stage": np.int64(0),
            "optimizer_name": b"Adam",
        },
        "partition_info": {"optimizer_sharded": {"original_dim": np.array([2, 3])}},
    }

    state_dict2 = {
        "model": {"full_precision": {"optimizer_sharded": np.array([1, 2, 3]), "non_sharded": np.array([11, 22, 33])}},
        "optimizer": {
            "optimizer_sharded": {
                "Moment_1": np.array([6, 5, 4]),
                "Moment_2": np.array([66, 55, 44]),
                "Step": np.array([5]),
            },
            "non_sharded": {
                "Moment_1": np.array([666, 555, 444]),
                "Moment_2": np.array([6666, 5555, 4444]),
                "Step": np.array([55]),
            },
        },
        "trainer_options": {
            "mixed_precision": np.bool_(False),
            "world_rank": np.int64(1),
            "world_size": np.int64(1),
            "horizontal_parallel_size": np.int64(1),
            "data_parallel_size": np.int64(1),
            "zero_stage": np.int64(0),
            "optimizer_name": b"Adam",
        },
        "partition_info": {"optimizer_sharded": {"original_dim": np.array([2, 3])}},
    }

    load_mock.side_effect = [trainer_options1, trainer_options2, trainer_options1, state_dict1, state_dict2]
    state_dict = checkpoint.aggregate_checkpoints(["abc", "def"], pytorch_format=False)

    assert (state_dict["model"]["full_precision"]["optimizer_sharded"] == np.array([1, 2, 3])).all()
    assert (state_dict["model"]["full_precision"]["non_sharded"] == np.array([11, 22, 33])).all()
    assert (state_dict["optimizer"]["optimizer_sharded"]["Moment_1"] == np.array([[9, 8, 7], [6, 5, 4]])).all()
    assert (state_dict["optimizer"]["optimizer_sharded"]["Moment_2"] == np.array([[99, 88, 77], [66, 55, 44]])).all()
    assert (state_dict["optimizer"]["optimizer_sharded"]["Step"] == np.array([5])).all()
    assert (state_dict["optimizer"]["non_sharded"]["Moment_1"] == np.array([666, 555, 444])).all()
    assert (state_dict["optimizer"]["non_sharded"]["Moment_2"] == np.array([6666, 5555, 4444])).all()
    assert (state_dict["optimizer"]["non_sharded"]["Step"] == np.array([55])).all()

    assert state_dict["trainer_options"]["mixed_precision"] == False
    assert state_dict["trainer_options"]["world_rank"] == 0
    assert state_dict["trainer_options"]["world_size"] == 1
    assert state_dict["trainer_options"]["horizontal_parallel_size"] == 1
    assert state_dict["trainer_options"]["data_parallel_size"] == 1
    assert state_dict["trainer_options"]["zero_stage"] == 0
    assert state_dict["trainer_options"]["optimizer_name"] == b"Adam"


@patch("onnxruntime.training._checkpoint_storage.load")
def test_checkpoint_aggregation_mixed_precision(load_mock):
    trainer_options1 = {
        "mixed_precision": np.bool_(True),
        "world_rank": np.int64(0),
        "world_size": np.int64(2),
        "horizontal_parallel_size": np.int64(1),
        "data_parallel_size": np.int64(2),
        "zero_stage": np.int64(1),
        "optimizer_name": b"Adam",
    }
    trainer_options2 = {
        "mixed_precision": np.bool_(True),
        "world_rank": np.int64(1),
        "world_size": np.int64(2),
        "horizontal_parallel_size": np.int64(1),
        "data_parallel_size": np.int64(2),
        "zero_stage": np.int64(1),
        "optimizer_name": b"Adam",
    }

    state_dict1 = {
        "model": {"full_precision": {"sharded": np.array([1, 2, 3]), "non_sharded": np.array([11, 22, 33])}},
        "optimizer": {
            "sharded": {"Moment_1": np.array([9, 8, 7]), "Moment_2": np.array([99, 88, 77]), "Step": np.array([5])},
            "non_sharded": {
                "Moment_1": np.array([666, 555, 444]),
                "Moment_2": np.array([6666, 5555, 4444]),
                "Step": np.array([55]),
            },
        },
        "trainer_options": {
            "mixed_precision": np.bool_(True),
            "world_rank": np.int64(0),
            "world_size": np.int64(1),
            "horizontal_parallel_size": np.int64(1),
            "data_parallel_size": np.int64(1),
            "zero_stage": np.int64(0),
            "optimizer_name": b"Adam",
        },
        "partition_info": {"sharded": {"original_dim": np.array([2, 3])}},
    }

    state_dict2 = {
        "model": {"full_precision": {"sharded": np.array([4, 5, 6]), "non_sharded": np.array([11, 22, 33])}},
        "optimizer": {
            "sharded": {"Moment_1": np.array([6, 5, 4]), "Moment_2": np.array([66, 55, 44]), "Step": np.array([5])},
            "non_sharded": {
                "Moment_1": np.array([666, 555, 444]),
                "Moment_2": np.array([6666, 5555, 4444]),
                "Step": np.array([55]),
            },
        },
        "trainer_options": {
            "mixed_precision": np.bool_(True),
            "world_rank": np.int64(1),
            "world_size": np.int64(1),
            "horizontal_parallel_size": np.int64(1),
            "data_parallel_size": np.int64(1),
            "zero_stage": np.int64(0),
            "optimizer_name": b"Adam",
        },
        "partition_info": {"sharded": {"original_dim": np.array([2, 3])}},
    }

    load_mock.side_effect = [trainer_options1, trainer_options2, trainer_options1, state_dict1, state_dict2]
    state_dict = checkpoint.aggregate_checkpoints(["abc", "def"], pytorch_format=False)

    assert (state_dict["model"]["full_precision"]["sharded"] == np.array([[1, 2, 3], [4, 5, 6]])).all()
    assert (state_dict["model"]["full_precision"]["non_sharded"] == np.array([11, 22, 33])).all()
    assert (state_dict["optimizer"]["sharded"]["Moment_1"] == np.array([[9, 8, 7], [6, 5, 4]])).all()
    assert (state_dict["optimizer"]["sharded"]["Moment_2"] == np.array([[99, 88, 77], [66, 55, 44]])).all()
    assert (state_dict["optimizer"]["sharded"]["Step"] == np.array([5])).all()
    assert (state_dict["optimizer"]["non_sharded"]["Moment_1"] == np.array([666, 555, 444])).all()
    assert (state_dict["optimizer"]["non_sharded"]["Moment_2"] == np.array([6666, 5555, 4444])).all()
    assert (state_dict["optimizer"]["non_sharded"]["Step"] == np.array([55])).all()

    assert state_dict["trainer_options"]["mixed_precision"] == True
    assert state_dict["trainer_options"]["world_rank"] == 0
    assert state_dict["trainer_options"]["world_size"] == 1
    assert state_dict["trainer_options"]["horizontal_parallel_size"] == 1
    assert state_dict["trainer_options"]["data_parallel_size"] == 1
    assert state_dict["trainer_options"]["zero_stage"] == 0
    assert state_dict["trainer_options"]["optimizer_name"] == b"Adam"
