import time
from itertools import product
from typing import Any, Callable, Dict, List, Optional, Union
import numpy
from onnx import AttributeProto, ModelProto, NodeProto, GraphProto, FunctionProto
from onnx.helper import make_model, make_node, make_graph, make_opsetid
from ...ext_test_case import measure_time


def has_subgraph(node: NodeProto) -> bool:
    """
    Tells if a node has a subgraph as an attribute.
    """
    for att in node.attribute:
        if att.type == AttributeProto.GRAPH:
            return True
    return False


def get_node_attribute(node: NodeProto, name: str) -> AttributeProto:
    """
    Returns the value of one attribute.

    :param node: node
    :param name: attribute name
    :return: value
    """
    for att in node.attribute:
        if att.name == name:
            return att
    raise KeyError(
        f"Unable to find {name!r} among {list(att.name for att in node.attribute)}."
    )


def change_onnx_operator_domain(
    onx: Union[ModelProto, GraphProto, FunctionProto],
    op_type: str,
    op_domain: str = "",
    new_op_type: Optional[str] = None,
    new_op_domain: Optional[str] = None,
    new_opset: Optional[int] = None,
    **kwargs: Dict[str, Any],
) -> Union[ModelProto, GraphProto, FunctionProto]:
    """
    Replaces an operator by another one in the same domain
    or another one.

    :param onx: proto to modify
    :param op_type: operator to look for
    :param op_domain: domain to look for
    :param new_op_type: new operator name or None for the same name
    :param new_op_domain: new domain name or None the for the same domain
    :param new_opset: new opset for the new domain
    :param kwargs: modified parameters, set it to None to remove them
    :return: same type as the input

    The function is not recursive yet.
    """

    def change_node(node):
        atts = []
        new_kwargs = {}
        for att in node.attribute:
            if att.name in kwargs:
                v = kwargs[att.name]
                if v is None:
                    continue
                new_kwargs[att.name] = v
                continue
            atts.append(att)
        for k, v in kwargs.items():
            if v is None or k in new_kwargs:
                continue
            new_kwargs[k] = v
        new_node = make_node(
            new_op_type or node.op_type,
            node.input,
            node.output,
            domain=new_op_domain or node.domain,
            **new_kwargs,
        )
        if len(atts) > 0:
            new_node.attribute.extend(atts)
        return new_node

    if isinstance(onx, GraphProto):
        new_nodes = []
        modified = False
        for node in onx.node:
            if has_subgraph(node):
                raise NotImplementedError(
                    f"The function is not recursive yet and cannot "
                    f"handle node {node.op_type!r} from domain "
                    f"{node.domain!r}."
                )
            if node.op_type == op_type and node.domain == op_domain:
                new_node = change_node(node)
                new_nodes.append(new_node)
                modified = True
                continue
            new_nodes.append(node)
        if not modified:
            return onx
        return make_graph(
            new_nodes,
            onx.name,
            onx.input,
            onx.output,
            onx.initializer,
            onx.sparse_initializer,
        )

    if isinstance(onx, FunctionProto):
        raise NotImplementedError()

    if not isinstance(onx, ModelProto):
        raise TypeError(f"Unexpected type for onx {type(onx)}.")

    new_graph = change_onnx_operator_domain(
        onx.graph,
        op_type=op_type,
        op_domain=op_domain,
        new_opset=new_opset,
        new_op_type=new_op_type,
        new_op_domain=new_op_domain,
        **kwargs,
    )
    if id(new_graph) == id(onx.graph):
        # no change
        return onx

    if new_op_domain is None:
        new_op_domain = op_domain
    if new_op_domain == op_domain and new_opset is not None:
        raise ValueError(
            f"If new_op_domain==domain=={new_op_domain!r}, "
            f"new_opset must be None not {new_opset}."
        )
    opsets = list(onx.opset_import)
    if new_op_domain != op_domain:
        opsets.append(make_opsetid(new_op_domain, new_opset or 1))

    new_model = make_model(
        new_graph,
        functions=onx.functions,
        ir_version=onx.ir_version,
        producer_name=onx.producer_name,
        producer_version=onx.producer_version,
        model_version=onx.model_version,
        doc_string=onx.doc_string,
        opset_imports=opsets,
        domain=onx.domain,
    )
    return new_model


def optimize_model(
    onx: ModelProto,
    feeds: Dict[str, numpy.ndarray],
    transform: Callable[[ModelProto], ModelProto],
    session: Callable[[ModelProto], Any],
    params: Dict[str, List[Any]],
    baseline: Optional[Callable[[ModelProto], Any]] = None,
    verbose: bool = False,
    number: int = 10,
    repeat: int = 10,
    warmup: int = 5,
    n_tries: int = 2,
    sleep: float = 0.1,
) -> List[Dict[str, Union[str, float]]]:
    """
    Optimizes a model by trying out many possibilities.

    :param onx: ModelProto
    :param feeds: inputs as a dictionary of numpy arrays
    :param transform: function taking a ModelProto and returning a ModelProto
        based on the values coming from *params*
    :param session: function which takes a modifed ModelProto
        and return a session
    :param params: dictionary of values to test `{ param_name: [ param_values ] }`
    :param baseline: function which takes a modifed ModelProto
        and return a session, identified as the baseline
    :param verbose: use :epkg:`tqdm` to show improvment
    :param number: parameter to :func:`measure_time
        <onnx_extended.ext_test_case.measure_time>`
    :param repeat: parameter to :func:`measure_time
        <onnx_extended.ext_test_case.measure_time>`
    :param warmup: parameter to :func:`measure_time
        <onnx_extended.ext_test_case.measure_time>`
    :param n_tries: number of times to measure, if the measurements returns
        very different results, values for *number* or *repeat* should
        be increased
    :param sleep: time to sleep between two measurements
    :return: list of results returned by :func:`measure_time
        <onnx_extended.ext_test_case.measure_time>`

    See example :ref:`l-plot-optim-tree-ensemble` for an example.
    """
    if sleep >= 1:
        raise ValueError(f"sleep={sleep} >= 1, probably a mistake.")
    keys = ["TRY"] + list(params.keys())
    sets = [list(range(n_tries))] + [params[k] for k in keys[1:]]
    loops = list(product(*sets))
    if verbose:
        from tqdm import tqdm

        loop = tqdm(loops)
    else:
        loop = loops

    res = []
    if baseline is not None:
        sess = baseline(onx)
        # one run to make run it is working
        sess.run(None, feeds)
        if sleep > 0:
            time.sleep(sleep)
        obs = measure_time(
            lambda sess=sess: sess.run(None, feeds),
            number=number,
            repeat=repeat,
            warmup=warmup,
        )
        obs["n_exp"] = 0
        obs["n_exp_name"] = "TRY=0,baseline"
        obs["short_name"] = "0,baseline"
        obs["TRY"] = 0
        obs["name"] = "baseline"
        res.append(obs)

    for it, values in enumerate(loop):
        if verbose:
            msg = [f"i={it+1}/{len(loops)}"]
            msg.extend([f"{k}={v}" for k, v in zip(keys, values)])
            loop.set_description(" ".join(msg))

        kwargs = dict(zip(keys, values))
        del kwargs["TRY"]
        onx_modified = transform(onx, **kwargs)
        sess = session(onx_modified)
        if sleep > 0:
            time.sleep(sleep)
        obs = measure_time(
            lambda sess=sess: sess.run(None, feeds),
            number=number,
            repeat=repeat,
            warmup=warmup,
        )
        obs.update(kwargs)
        obs["n_exp"] = it
        obs["n_exp_name"] = ",".join(f"{k}={v}" for k, v in zip(keys, values))
        obs["short_name"] = ",".join(f"{v}" for v in values)
        obs["name"] = ",".join(f"{v}" for v in values[1:])
        res.append(obs)

    if baseline is not None:
        for n in range(1, n_tries):
            sess = baseline(onx)
            if sleep > 0:
                time.sleep(sleep)
            obs = measure_time(
                lambda sess=sess: sess.run(None, feeds),
                number=number,
                repeat=repeat,
                warmup=warmup,
            )
            obs["n_exp"] = 0
            obs["n_exp_name"] = f"TRY={n},baseline"
            obs["short_name"] = f"{n},baseline"
            obs["name"] = "baseline"
            obs["TRY"] = n
            res.append(obs)

    return res
