# Copyright (c) 2023 Graphcore Ltd. All rights reserved.
from typing import Dict, List

import numpy as np
import onnx

from onnx import helper

from poprt import runtime
from poprt.compiler import Compiler, CompilerOptions

np.random.seed(1984)


def get_weights_info(infos: List[runtime.DataDesc]) -> Dict[str, runtime.DataDesc]:
    weights_info = {}
    for info in infos:
        if info.popef_contains_tensor_data:
            weights_info[info.name] = info

    return weights_info


def compile(model: onnx.ModelProto):
    """Compile ONNX to PopEF."""
    outputs = [o.name for o in model.graph.output]

    options = CompilerOptions()
    options.constant_weights = False
    options.enable_engine_caching = True
    options.cache_path = './cache'

    executable = Compiler.compile(model, outputs, options)
    return executable


def test_weights_update(executable):
    # get device
    dm = runtime.DeviceManager()
    replica_to_device = runtime.ReplicaIdToDevice()
    for device_id in range(dm.get_num_devices()):
        try:
            device = dm.get_specific_device(device_id)
            replica_to_device[0] = device
            break
        except:
            pass
    if not replica_to_device:
        raise RuntimeError("No IPU can be attached.")

    config = runtime.LightRunnerConfig()
    config.replica_to_device = replica_to_device

    runner = runtime.LightRunner(executable, config)

    # generate random inputs
    inputs = {}
    outputs = {}
    for info in runner.get_execute_inputs():
        inputs[info.name] = np.random.uniform(0, 1, info.shape).astype(
            info.numpy_data_type()
        )
    for info in runner.get_execute_outputs():
        outputs[info.name] = np.zeros(info.shape, dtype=info.numpy_data_type())

    # run with original weights(from popef)
    future = runner.execute_async(inputs, outputs)
    future.wait()
    print(outputs["output"].flatten()[1000:1020])

    # generate random weights
    weights_info = get_weights_info(runner.get_model_inputs())
    weights = {}
    for name, info in weights_info.items():
        weights[name] = np.random.uniform(-1, 1, info.shape).astype(
            info.numpy_data_type()
        )

    # update weights
    for k, v in weights.items():
        runner.set_input_anchor(k, v)
    runner.update_weights()
    print(f"fin update_weights")

    # run with updated weights
    future = runner.execute_async(inputs, outputs)
    future.wait()
    print(outputs["output"].flatten()[1000:1020])


def test_reload(executable):
    # get device
    dm = runtime.DeviceManager()
    replica_to_device = runtime.ReplicaIdToDevice()
    for device_id in range(dm.get_num_devices()):
        try:
            device = dm.get_specific_device(device_id)
            replica_to_device[0] = device
            break
        except:
            pass
    if not replica_to_device:
        raise RuntimeError("No IPU can be attached.")

    config = runtime.LightRunnerConfig()
    config.replica_to_device = replica_to_device

    runner = runtime.LightRunner(executable, config)

    # generate random inputs
    inputs = {}
    outputs = {}
    outputs_info = runner.get_execute_outputs()
    for info in runner.get_execute_inputs():
        inputs[info.name] = np.random.uniform(0, 1, info.shape).astype(
            info.numpy_data_type()
        )
    for info in outputs_info:
        outputs[info.name] = np.zeros(info.shape, dtype=info.numpy_data_type())

    # runner
    runner.execute(inputs, outputs)
    print(outputs["output"].flatten()[1000:1020])
    # reset outputs
    for info in outputs_info:
        outputs[info.name] = np.zeros(info.shape, dtype=info.numpy_data_type())

    # stop runner
    runner.stop()

    # another runner
    another_runner = runtime.LightRunner(executable, config)
    another_runner.execute(inputs, outputs)
    print(outputs["output"].flatten()[1000:1020])
    # reset outputs
    for info in outputs_info:
        outputs[info.name] = np.zeros(info.shape, dtype=info.numpy_data_type())

    # stop another_runner
    another_runner.stop()

    # reload runner
    runner.reload()
    runner.execute(inputs, outputs)
    print(outputs["output"].flatten()[1000:1020])


def default_model():
    TensorProto = onnx.TensorProto

    nodes = []
    num_matmuls = 4
    nodes.append(helper.make_node("Expand", ["input", "shape"], ["Act0"]))
    for i in range(num_matmuls):
        nodes.append(helper.make_node("MatMul", [f"Act{i}", "Weight"], [f"Act{i+1}"]))
    nodes.append(
        helper.make_node("ReduceMean", [f"Act{num_matmuls}"], ["output"], axes=[0, 1])
    )

    graph = helper.make_graph(
        nodes,
        "matmul_test",
        [
            helper.make_tensor_value_info("input", TensorProto.FLOAT, (256, 256)),
        ],
        [helper.make_tensor_value_info("output", TensorProto.FLOAT, (256, 256))],
        [
            helper.make_tensor(
                "shape",
                TensorProto.INT64,
                [4],
                np.array([4, 4, 256, 256], dtype=np.int64),
            ),
            helper.make_tensor(
                "Weight",
                TensorProto.FLOAT,
                (4, 4, 256, 256),
                np.random.randn(4, 4, 256, 256),
            ),
        ],
    )
    opset_imports = [helper.make_opsetid("", 11)]
    original_model = helper.make_model(graph, opset_imports=opset_imports)
    return original_model


if __name__ == '__main__':
    model = default_model()
    exe = compile(model)

    test_weights_update(exe)
    test_reload(exe)