from typing import Callable

import torch
import torch.nn as nn

from torch.nn import functional as F

from deepstruct.graph import LayeredGraph
from deepstruct.sparse import MaskedLinearLayer


class LinearBlockUnit(nn.Module):
    def __init__(self, input_size, output_size, stride=1):
        super(LinearBlockUnit, self).__init__()
        self.act = nn.ReLU()
        # self.conv = nn.Conv2d(1, 10, kernel_size=3, stride=stride, padding=1)
        self.linear = nn.Linear(input_size, output_size)
        self.bn = nn.BatchNorm1d(output_size)

    def forward(self, x):
        out = self.act(x)
        out = self.linear(out)
        out = self.bn(out)
        return out


class FiveWaySparseBlockNet(nn.Module):
    def __init__(self, input_size, output_classes, hidden_block_size=100):
        super(FiveWaySparseBlockNet, self).__init__()

        self.input_blocks = nn.ModuleList()
        for _ in range(5):
            self.input_blocks.append(LinearBlockUnit(input_size, hidden_block_size))

        self.hidden_blocks = nn.ModuleList()
        for _ in range(5):
            self.hidden_blocks.append(
                LinearBlockUnit(hidden_block_size, hidden_block_size)
            )

        self.output_blocks = nn.ModuleList()
        for _ in range(5):
            self.output_blocks.append(
                LinearBlockUnit(hidden_block_size, output_classes)
            )

    def forward(self, x):
        input_to_blocks = {}
        for idx, block_unit in enumerate(self.input_blocks):
            input_to_blocks[idx] = block_unit(x)

        hidden_block_results = {}
        for idx, block_unit in enumerate(self.hidden_blocks):
            hidden_block_results[idx] = block_unit(input_to_blocks[idx])

        output_blocks = []
        for idx, block_unit in enumerate(self.output_blocks):
            output_blocks.append(block_unit(hidden_block_results[idx]))

        return torch.mean(torch.stack(output_blocks), dim=0)


cell_constructor_func: Callable[
    [bool, bool, int, int, int], nn.Module
] = lambda is_input, is_output, in_degree, out_degree, layer: nn.Conv2d(10, 10, 1)
skip_map_fn: Callable[[int, int], nn.Module] = lambda layer: nn.Conv2d(10, 10, 1)
layer_channel_size_func: Callable[[int], int] = lambda layer: 10


class DeepDACellNetwork(nn.Module):
    def __init__(
        self,
        num_classes,
        cell_constructor: cell_constructor_func,
        layer_channel_size: layer_channel_size_func,
        fn_skip_map: skip_map_fn,
        structure: LayeredGraph,
    ):
        super(DeepDACellNetwork, self).__init__()

        self._structure = structure
        assert structure.num_layers > 0
        self._layer_channel_size = layer_channel_size

        self._input_nodes = [n for n in structure.nodes if structure.in_degree(n) == 0]
        self._output_nodes = [
            n for n in structure.nodes if structure.out_degree(n) == 0
        ]
        assert len(self._input_nodes) > 0, "No input nodes in structure: len=%d" % len(
            self.input_nodes
        )
        assert (
            len(self._output_nodes) > 0
        ), "No output nodes in structure: len=%d" % len(self.output_nodes)

        # For each node in our structure, we create a cell based on the given cell constructor
        # The cell constructor has to return a nn.Module which maps [B, C_s, N, M, D] -> [B, C_t, N, M]
        # in which B stands for batch size, C_s for the channel size of the source layer defined by
        # layer_channel_size_func, N and M for width and height of the features, C_t for the channel size of the target
        # layer defined by layer_channel_size_func and D for the outgoing degree of the particular node (which needs to
        # be condensed, e.g. by aggregating on the last dimension)
        self._node_cells = {
            node: cell_constructor(
                is_input=node in self._input_nodes,
                is_output=node in self._output_nodes,
                in_degree=structure.in_degree(node),
                out_degree=structure.out_degree(node),
                layer=structure.get_layer(node),
            )
            for node in self._structure.nodes
        }
        self._nodes = nn.ModuleList(list(self._node_cells.values()))

        # Additionaly maps for skip-connections are created. These map channels from one
        self._node_maps = {}
        self._node_maps_list = []
        for source_layer in self._structure.layers[:-1]:
            self._node_maps[source_layer] = {}
            source_output_channel_size = self._layer_channel_size(source_layer + 1)
            for target_layer in self._structure.layers[source_layer + 1 :]:
                target_channel_size = self._layer_channel_size(target_layer + 1)
                if source_output_channel_size is not target_channel_size:
                    # self._node_maps[source_layer][target_layer] = nn.Conv2d(source_output_channel_size, target_channel_size, kernel_size=1)
                    self._node_maps[source_layer][target_layer] = fn_skip_map(
                        source_output_channel_size, target_channel_size
                    )
                    # print(self._node_maps[source_layer][target_layer])
                    self._node_maps_list.append(
                        self._node_maps[source_layer][target_layer]
                    )
        self._maps = nn.ModuleList(self._node_maps_list)

        # Extract number of out channels from node of --last-- layer (not any output node)
        # A node in the last layer has actually the number of out channels everything is mapped to
        any_output_node = self._structure.get_vertices(self._structure.last_layer)[0]
        num_out_channel = self._node_cells[any_output_node].num_out_channel

        self.last_representation_size = 1000
        self.convlast = nn.Conv2d(
            num_out_channel, self.last_representation_size, kernel_size=1
        )
        self.bnlast = nn.BatchNorm2d(self.last_representation_size)
        self.fc = MaskedLinearLayer(self.last_representation_size, num_classes)

    def _apply_skip_layer_adaption(self, source_node, target_node, input):
        source_layer = self._structure.get_layer(source_node)
        target_layer = self._structure.get_layer(target_node)
        self._layer_channel_size(source_layer)
        self._layer_channel_size(target_layer)

    def forward(self, input):
        # input: [batch_size, channels, N, M] = [B, C_in, N, M]

        outputs = {node: None for node in self._structure.nodes}

        for layer in self._structure.layers:
            for current_node in self._structure.get_vertices(layer):

                if current_node in self._input_nodes:
                    current_input = [input]
                else:
                    # current_input = [ self._node_maps[self._structure.get_layer(u)][layer](outputs[u])
                    #                  if layer in self._node_maps[self._structure.get_layer(u)] else outputs[u] for (u, _) in self._structure.in_edges(current_node)]
                    current_input = []
                    for (source, _) in self._structure.in_edges(current_node):
                        source_layer = self._structure.get_layer(source)
                        if layer in self._node_maps[source_layer]:
                            """print('From source layer', source_layer)
                            print('To target layer', layer)
                            print(self._node_maps[source_layer][layer].in_channels)
                            print(self._node_maps[source_layer][layer].out_channels)
                            print('Source output shape:', outputs[source].shape)
                            print('Mapped source output to new channel', self._node_maps[source_layer][layer].out_channels)"""
                            current_input.append(
                                self._node_maps[source_layer][layer](outputs[source])
                            )
                        else:
                            current_input.append(outputs[source])

                feed = torch.stack(current_input, dim=-1)

                outputs[current_node] = self._node_cells[current_node](feed)

            # Force GPU to possibly free memory as stacks of collected inputs/outputs can grow rapidly
            torch.cuda.empty_cache()

        # Collect all outputs of the output nodes into a list and stack it
        # Use the channel output size of the last layer as targeted outgoing channel size
        # Every output node from previous layers are mapped through a convolution to the next layer
        out_list = []
        for node in self._output_nodes:
            source_layer = self._structure.get_layer(node)
            if (
                source_layer is self._structure.last_layer
                or self._structure.last_layer not in self._node_maps[source_layer]
            ):
                if source_layer is not self._structure.last_layer:
                    print(
                        "Channel size source layer:",
                        self._layer_channel_size(source_layer),
                    )
                    print(
                        "Channel size last layer:",
                        self._layer_channel_size(self._structure.last_layer),
                    )
                    print("Output shape:", outputs[node].shape)
                out_list.append(outputs[node])
            else:
                map_to_last_layer = self._node_maps[source_layer][
                    self._structure.last_layer
                ]
                out_list.append(map_to_last_layer(outputs[node]))
        print([o.shape for o in out_list])
        result = torch.mean(torch.stack(out_list), dim=0)  # [B, C_out, N, M]

        # Classifier
        y = F.relu(result)  # [B, C_out, N, M]
        y = self.convlast(y)  # [B, 1000, N, M]
        y = self.bnlast(y)  # [B, 1000, N, M]
        y = F.adaptive_avg_pool2d(y, (1, 1))  # [B, 1000, 1, 1]
        y = y.view(y.size(0), -1)  # [B, 1000]
        return self.fc(y)  # [B, num_classes]
        # return F.log_softmax(y, dim=1) # [B, num_classes]
