import collections
import random
from typing import Tuple

import tests.resources.reuse_subdag
from hamilton import ad_hoc_utils, graph
from hamilton.function_modifiers import config, parameterized_subdag, recursive, subdag, value
from hamilton.function_modifiers.dependencies import source


def test_collect_function_fns():
    val = random.randint(0, 100000)

    def test_fn(out: int = val) -> int:
        return out

    assert recursive.subdag.collect_functions(load_from=[test_fn])[0]() == test_fn()


def test_collect_functions_module():
    val = random.randint(0, 100000)

    def test_fn(out: int = val) -> int:
        return out

    assert (
        recursive.subdag.collect_functions(
            load_from=[ad_hoc_utils.create_temporary_module(test_fn)]
        )[0]()
        == test_fn()
    )


def test_assign_namespaces():
    assert recursive.assign_namespace(node_name="foo", namespace="bar") == "bar.foo"


def foo(a: int) -> int:
    return a


@config.when_not(some_config_param=True)
def bar(b: int) -> int:
    return b


@config.when(some_config_param=True)
def bar__alt() -> int:
    return 10


def test_reuse_subdag_validate_succeeds():
    def baz(foo: int, bar: str) -> str:
        return bar * foo

    decorator = recursive.subdag(
        foo,
        bar,
        inputs={},
        config={},
    )
    decorator.validate(baz)


def test_reuse_subdag_basic_no_parameterization():
    def baz(foo: int, bar: int) -> int:
        return foo + bar

    decorator = recursive.subdag(foo, bar, inputs={}, config={})
    nodes = {node_.name: node_ for node_ in decorator.generate_nodes(baz, {})}
    # subdags have prefixed names
    assert "baz.foo" in nodes
    assert "baz.bar" in nodes
    # but we expect our outputs to exist as well
    assert "baz" in nodes
    assert nodes["baz"].callable(**{"baz.foo": 10, "baz.bar": 20}) == 30


def test_reuse_subdag_basic_simple_parameterization():
    def baz(foo: int, bar: int) -> int:
        return foo + bar

    decorator = recursive.subdag(
        foo,
        bar,
        config={},
        inputs={"a": value(1), "b": value(2)},
    )
    nodes = {node_.name: node_ for node_ in decorator.generate_nodes(baz, {})}
    # This doesn't necessarily have to be part of the contract, but we're testing now to ensure that it works
    assert "baz.a" in nodes
    assert nodes["baz.a"]() == 1
    assert "baz.b" in nodes
    assert nodes["baz.b"]() == 2


def test_reuse_subdag_basic_source_parameterization():
    def baz(foo: int, bar: int) -> int:
        return foo + bar

    decorator = recursive.subdag(foo, bar, inputs={"a": source("c"), "b": source("d")}, config={})
    nodes = {node_.name: node_ for node_ in decorator.generate_nodes(baz, {})}
    # These aren't entirely part of the contract, but they're required for the
    # way we're currently implementing it. See https://github.com/dagworks-inc/hamilton/issues/201
    assert "baz.a" in nodes
    assert nodes["baz.a"](c=1) == 1
    assert "baz.b" in nodes
    assert nodes["baz.b"](d=2) == 2


def test_reuse_subdag_handles_config_assignment():
    def baz(foo: int, bar: int) -> int:
        return foo + bar

    decorator = recursive.subdag(
        foo,
        bar,
        bar__alt,
        inputs={"a": value(1)},
        config={"some_config_param": True},
    )
    nodes = {node_.name: node_ for node_ in decorator.generate_nodes(baz, {})}
    assert nodes["baz.bar"]() == 10


def test_reuse_subdag_allows_config_as_input():
    def baz(foo: int, bar: int) -> int:
        return foo + bar

    decorator = recursive.subdag(
        foo,
        bar,
        bar__alt,
        inputs={},
        config={"some_config_param": False, "b": 100},
    )
    nodes = {node_.name: node_ for node_ in decorator.generate_nodes(baz, {})}
    assert nodes["baz.b"]() == 100
    assert nodes["baz.bar"](**{"baz.b": 100}) == 100


def test_subdag_empty_config():
    @config.when_not(some_config_param=False)
    def baz(foo: int, bar: int) -> int:
        return foo + bar

    def subdag(baz: int) -> int:
        return baz

    decorator = recursive.subdag(baz, inputs={})
    nodes = {node_.name: node_ for node_ in decorator.generate_nodes(subdag, {})}
    assert "subdag.baz" in nodes
    assert "subdag" in nodes


def test_reuse_subdag_end_to_end():
    fg = graph.FunctionGraph(tests.resources.reuse_subdag, config={"op": "subtract"})
    prefixless_nodes = []
    prefixed_nodes = collections.defaultdict(list)
    for name, node in fg.nodes.items():
        name_split = name.split(".")
        if len(name_split) == 1:
            prefixless_nodes.append(node)
        else:
            namespace, name = name_split
            prefixed_nodes[namespace].append(node)
    node_set = set(fg.nodes)
    assert {
        "e_1",
        "e_2",
        "e_3",
        "e_4",
        "f_1",
        "f_2",
        "f_3",
        "f_4",
    } - node_set == set()  # All the nodes outputted by our subdags
    assert {
        "v1.d",
        "v2.d",
        "v3.d",
        "v4.d",
        "v1.e",
        "v2.e",
        "v3.e",
        "v4.e",
        "v1.f",
        "v2.f",
        "v3.f",
        "v4.f",
    } - node_set == set()  # All these nodes must be in the DAG -- they're all the namespaced nodes
    assert {
        "v1.e",
        "v2.e",
        "v3.e",
        "v4.e",
        "v1.f",
        "v2.f",
        "v3.f",
        "v4.f",
    } - node_set == set()  # All these nodes must be in the DAG
    assert {"a", "b"} - node_set == set()  # common nodes shared by the DAG, not as part of subdags
    assert {"e", "f", "d"} - node_set == {
        "e",
        "f",
        "d",
    }  # We've defined a node e, f, and d, but they're only namespaced in subdags
    assert {
        "v1.c",
        "v2.c",
        "v3.c",
    } - node_set == set()  # These are all static values that are namespaced in the subDAGs
    # The following are all static values
    assert fg.nodes["v1.c"].callable() == 10
    assert fg.nodes["v2.c"].callable() == 20
    assert fg.nodes["v3.c"].callable() == 30

    # Source assigned this
    assert list(fg.nodes["v4.e"].input_types)[0] == "v4.c"
    assert list(fg.nodes["v4.c"].input_types)[0] == "b"
    assert fg.nodes["v4.c"].callable(b=1234) == 1234
    # # Check that the config is assigned and overwritten correctly
    assert fg.nodes["v1.d"].callable(**{"v1.c": 10, "a": 100}) == 100 - 10
    assert fg.nodes["v3.d"].callable(**{"v3.c": 10, "a": 100}) == 100 + 10
    res = fg.execute(nodes=[fg.nodes["sum_everything"]])
    assert res["sum_everything"] == 318


def test_parameterized_subdag():
    def bar(input_1: int) -> int:
        return input_1 + 1

    def foo(input_2: int) -> int:
        return input_2 + 1

    @config.when(baz_version="standard")
    def baz__standard(foo: int, bar: int) -> int:
        return foo + bar

    @config.when(baz_version="alternate")
    def baz__alternate(foo: int, bar: int) -> int:
        return foo * bar

    def subdag_processor(foo: int, bar: int, baz: int) -> Tuple[int, int, int]:
        return foo, bar, baz

    fns = [bar, foo, baz__standard, baz__alternate]

    decorator = parameterized_subdag(
        *fns,
        inputs={"input_1": source("external_input_1"), "input_2": source("external_input_2")},
        config={"baz_version": "standard"},
        v0={},
        v1={"config": {"baz_version": "standard"}, "inputs": {"input_1": value(100)}},
        v2={
            "config": {"baz_version": "alternate"},
            "inputs": {"input_2": value(200)},
        },
    )
    dag_generated = decorator.generate_nodes(subdag_processor, {})
    assert len(dag_generated) == len(set([item.name for item in dag_generated]))
    nodes_by_name = {node.name: node for node in dag_generated}
    assert "v0" in nodes_by_name
    assert "v1" in nodes_by_name
    assert "v2" in nodes_by_name
    assert "v0.baz" in nodes_by_name
    assert "v1.baz" in nodes_by_name
    assert "v2.baz" in nodes_by_name
    assert nodes_by_name["v0.baz"].callable(**{"v0.foo": 1, "v0.bar": 2}) == 3
    assert nodes_by_name["v1.baz"].callable(**{"v1.foo": 1, "v1.bar": 2}) == 3
    assert nodes_by_name["v2.baz"].callable(**{"v2.foo": 1, "v2.bar": 2}) == 2


def test_nested_subdag():
    def bar(input_1: int) -> int:
        return input_1 + 1

    def foo(input_2: int) -> int:
        return input_2 + 1

    @subdag(
        foo,
        bar,
    )
    def inner_subdag(foo: int, bar: int) -> Tuple[int, int]:
        return foo, bar

    @subdag(inner_subdag, inputs={"input_2": value(10)}, config={"plus_one": True})
    def outer_subdag_1(inner_subdag: Tuple[int, int]) -> int:
        return sum(inner_subdag)

    @subdag(inner_subdag, inputs={"input_2": value(3)}, config={"plus_one": False})
    def outer_subdag_2(inner_subdag: Tuple[int, int]) -> int:
        return sum(inner_subdag)

    def sum_all(outer_subdag_1: int, outer_subdag_2: int) -> int:
        return outer_subdag_1 + outer_subdag_2

    # we only need to generate from the outer subdag
    # as it refers to the inner one
    full_module = ad_hoc_utils.create_temporary_module(outer_subdag_1, outer_subdag_2, sum_all)
    fg = graph.FunctionGraph(full_module, config={})
    assert "outer_subdag_1" in fg.nodes
    assert "outer_subdag_2" in fg.nodes
    res = fg.execute(nodes=[fg.nodes["sum_all"]], inputs={"input_1": 2})
    # This is effectively the function graph
    assert res["sum_all"] == sum_all(
        outer_subdag_1(inner_subdag(bar(2), foo(10))), outer_subdag_2(inner_subdag(bar(2), foo(3)))
    )


def test_nested_subdag_with_config():
    def bar(input_1: int) -> int:
        return input_1 + 1

    @config.when(broken=False)
    def foo(input_2: int) -> int:
        return input_2 + 1

    @subdag(
        foo,
        bar,
    )
    def inner_subdag(foo: int, bar: int) -> Tuple[int, int]:
        return foo, bar

    @subdag(inner_subdag, inputs={"input_2": value(10)})
    def outer_subdag_1(inner_subdag: Tuple[int, int]) -> int:
        return sum(inner_subdag)

    @subdag(inner_subdag, inputs={"input_2": value(3)})
    def outer_subdag_2(inner_subdag: Tuple[int, int]) -> int:
        return sum(inner_subdag)

    def sum_all(outer_subdag_1: int, outer_subdag_2: int) -> int:
        return outer_subdag_1 + outer_subdag_2

    # we only need to generate from the outer subdag
    # as it refers to the inner one
    full_module = ad_hoc_utils.create_temporary_module(outer_subdag_1, outer_subdag_2, sum_all)
    fg = graph.FunctionGraph(full_module, config={"broken": False})
    assert "outer_subdag_1" in fg.nodes
    assert "outer_subdag_2" in fg.nodes
    res = fg.execute(nodes=[fg.nodes["sum_all"]], inputs={"input_1": 2})
    # This is effectively the function graph
    assert res["sum_all"] == sum_all(
        outer_subdag_1(inner_subdag(bar(2), foo(10))), outer_subdag_2(inner_subdag(bar(2), foo(3)))
    )


def test_subdag_with_external_nodes_input():
    def bar(input_1: int) -> int:
        return input_1 + 1

    def foo(input_2: int) -> int:
        return input_2 + 1

    @subdag(foo, bar, external_inputs=["baz"])
    def foo_bar_baz(foo: int, bar: int, baz: int) -> int:
        return foo + bar + baz

    full_module = ad_hoc_utils.create_temporary_module(foo_bar_baz)
    fg = graph.FunctionGraph(full_module, config={})
    # since we've provided it above,
    assert "baz" in fg.nodes
    assert fg.nodes["baz"].user_defined
    res = fg.execute(nodes=[fg.nodes["foo_bar_baz"]], inputs={"input_1": 2, "input_2": 3, "baz": 4})
    assert res["foo_bar_baz"] == foo_bar_baz(foo(3), bar(2), 4)


def test_parameterized_subdag_with_external_inputs_global():
    def bar(input_1: int) -> int:
        return input_1 + 1

    def foo(input_2: int) -> int:
        return input_2 + 1

    @parameterized_subdag(
        foo,
        bar,
        external_inputs=["baz"],
        foo_bar_baz_input_1={"inputs": {"input_1": value(10), "input_2": value(20)}},
        foo_bar_baz_input_2={"inputs": {"input_1": value(30), "input_2": value(40)}},
    )
    def foo_bar_baz(foo: int, bar: int, baz: int) -> int:
        return foo + bar + baz

    def foo_bar_baz_summed(foo_bar_baz_input_1: int, foo_bar_baz_input_2: int) -> int:
        return foo_bar_baz_input_1 + foo_bar_baz_input_2

    full_module = ad_hoc_utils.create_temporary_module(foo_bar_baz, foo_bar_baz_summed)
    fg = graph.FunctionGraph(full_module, config={})
    # since we've provided it above,
    assert "baz" in fg.nodes
    assert fg.nodes["baz"].user_defined
    res = fg.execute(nodes=[fg.nodes["foo_bar_baz_summed"]], inputs={"baz": 100})
    assert res["foo_bar_baz_summed"] == foo_bar_baz(foo(10), foo(20), 100) + foo_bar_baz(
        bar(30), foo(40), 100
    )


def test_parameterized_subdag_with_config():
    def bar(input_1: int) -> int:
        return input_1 + 1

    @config.when(broken=False)
    def foo(input_2: int) -> int:
        return input_2 + 1

    @subdag(
        foo,
        bar,
    )
    def inner_subdag(foo: int, bar: int) -> Tuple[int, int]:
        return foo, bar

    @subdag(inner_subdag, inputs={"input_2": value(10)})
    def outer_subdag_1(inner_subdag: Tuple[int, int]) -> int:
        return sum(inner_subdag)

    @subdag(inner_subdag, inputs={"input_2": value(3)})
    def outer_subdag_2(inner_subdag: Tuple[int, int]) -> int:
        return sum(inner_subdag)

    def sum_all(outer_subdag_1: int, outer_subdag_2: int) -> int:
        return outer_subdag_1 + outer_subdag_2

    # we only need to generate from the outer subdag
    # as it refers to the inner one
    full_module = ad_hoc_utils.create_temporary_module(outer_subdag_1, outer_subdag_2, sum_all)
    fg = graph.FunctionGraph(full_module, config={"broken": False})
    assert "outer_subdag_1" in fg.nodes
    assert "outer_subdag_2" in fg.nodes
    res = fg.execute(nodes=[fg.nodes["sum_all"]], inputs={"input_1": 2})
    # This is effectively the function graph
    assert res["sum_all"] == sum_all(
        outer_subdag_1(inner_subdag(bar(2), foo(10))), outer_subdag_2(inner_subdag(bar(2), foo(3)))
    )


def test_nested_parameterized_subdag_with_config():
    def bar(input_1: int) -> int:
        return input_1 + 1

    @config.when(broken=False)
    def foo(input_2: int) -> int:
        return input_2 + 1

    @parameterized_subdag(foo, bar, inner_subdag={})
    def inner_subdag(foo: int, bar: int) -> Tuple[int, int]:
        return foo, bar

    @parameterized_subdag(
        inner_subdag,
        outer_subdag_1={"inputs": {"input_2": value(10)}},
        outer_subdag_2={"inputs": {"input_2": value(3)}},
    )
    def outer_subdag_n(inner_subdag: Tuple[int, int]) -> int:
        return sum(inner_subdag)

    def sum_all(outer_subdag_1: int, outer_subdag_2: int) -> int:
        return outer_subdag_1 + outer_subdag_2

    # we only need to generate from the outer subdag
    # as it refers to the inner one
    full_module = ad_hoc_utils.create_temporary_module(outer_subdag_n, sum_all)
    fg = graph.FunctionGraph(full_module, config={"broken": False})
    assert "outer_subdag_1" in fg.nodes
    assert "outer_subdag_2" in fg.nodes
    res = fg.execute(nodes=[fg.nodes["sum_all"]], inputs={"input_1": 2})
    # This is effectively the function graph
    assert res["sum_all"] == sum_all(
        outer_subdag_n(inner_subdag(bar(2), foo(10))), outer_subdag_n(inner_subdag(bar(2), foo(3)))
    )