from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import math
import torch
import torch.nn as nn
import alphavideo.action.AlphAction.utils.registry as registry
from .utils.IA_helper import has_memory, has_person, has_object


class InteractionBlock(nn.Module):
    def __init__(self, dim_person, dim_other, dim_out, dim_inner, structure_config,
                 max_others=20,
                 temp_pos_len=-1,  # configs for memory feature
                 dropout=0.):
        super(InteractionBlock, self).__init__()
        self.dim_person = dim_person
        self.dim_other = dim_other
        self.dim_out = dim_out
        self.dim_inner = dim_inner
        self.max_others = max_others
        self.scale_value = dim_inner ** (-0.5)
        # config for temporal position, only used for memory feature,
        self.temp_pos_len = temp_pos_len

        bias = not structure_config.NO_BIAS
        init_std = structure_config.CONV_INIT_STD

        self.query = nn.Conv3d(dim_person, dim_inner, 1, bias)
        init_layer(self.query, init_std, bias)

        self.key = nn.Conv3d(dim_other, dim_inner, 1, bias)
        init_layer(self.key, init_std, bias)

        self.value = nn.Conv3d(dim_other, dim_inner, 1, bias)
        init_layer(self.value, init_std, bias)

        self.out = nn.Conv3d(dim_inner, dim_out, 1, bias)
        if structure_config.USE_ZERO_INIT_CONV:
            out_init = 0
        else:
            out_init = init_std
        init_layer(self.out, out_init, bias)

        self.softmax = nn.Softmax(dim=2)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

        self.use_ln = structure_config.LAYER_NORM

        if dim_person != dim_out:
            self.shortcut = nn.Conv3d(dim_person, dim_out, 1, bias)
            init_layer(self.shortcut, init_std, bias)
        else:
            self.shortcut = None

        if self.temp_pos_len > 0:
            self.temporal_position_k = nn.Parameter(torch.zeros(temp_pos_len, 1, self.dim_inner, 1, 1, 1))
            self.temporal_position_v = nn.Parameter(torch.zeros(temp_pos_len, 1, self.dim_inner, 1, 1, 1))

    def forward(self, person, others):
        """
        :param person: [n, channels, t, h, w]
        :param others: [n, num_other, channels, t, h, w]
        """
        device = person.device
        n, dim_person, t, h, w = person.size()
        _, max_others, dim_others, t_others, h_others, w_others = others.size()

        query_batch = person
        key = fuse_batch_num(others)  # [n*num_other, channels, t, h, w]

        query_batch = self.query(query_batch)
        key_batch = self.key(key).contiguous().view(n, self.max_others, self.dim_inner, t_others, h_others, w_others)
        value_batch = self.value(key).contiguous().view(n, self.max_others, self.dim_inner, t_others, h_others,
                                                        w_others)

        if self.temp_pos_len > 0:
            max_person_per_sec = max_others // self.temp_pos_len
            key_batch = key_batch.contiguous().view(n, self.temp_pos_len, max_person_per_sec, self.dim_inner, t_others,
                                                    h_others,
                                                    w_others)
            key_batch = key_batch + self.temporal_position_k
            key_batch = key_batch.contiguous().view(n, self.max_others, self.dim_inner, t_others,
                                                                    h_others, w_others)

            value_batch = value_batch.contiguous().view(n, self.temp_pos_len, max_person_per_sec, self.dim_inner,
                                                        t_others,
                                                        h_others, w_others)
            value_batch = value_batch + self.temporal_position_v
            value_batch = value_batch.contiguous().view(n, self.max_others, self.dim_inner, t_others,
                                                                        h_others, w_others)

        query_batch = query_batch.contiguous().view(n, self.dim_inner, -1).transpose(1, 2)  # [n, thw, dim_inner]
        key_batch = key_batch.contiguous().view(n, self.max_others, self.dim_inner, -1).transpose(2, 3)
        key_batch = key_batch.contiguous().view(n, self.max_others * t_others * h_others * w_others, -1).transpose(
            1, 2)

        qk = torch.bmm(query_batch, key_batch)  # n, thw, max_other * thw

        qk_sc = qk * self.scale_value

        weight = self.softmax(qk_sc)

        value_batch = value_batch.contiguous().view(n, self.max_others, self.dim_inner, -1).transpose(2, 3)
        value_batch = value_batch.contiguous().view(n, self.max_others * t_others * h_others * w_others, -1)
        out = torch.bmm(weight, value_batch)  # n, thw, dim_inner

        out = out.contiguous().view(n, t * h * w, -1)
        out = out.transpose(1, 2)
        out = out.contiguous().view(n, self.dim_inner, t, h, w)

        if self.use_ln:
            if not hasattr(self, "layer_norm"):
                self.layer_norm = nn.LayerNorm([self.dim_inner, t, h, w], elementwise_affine=False).to(
                    device)
            out = self.layer_norm(out)

        out = self.relu(out)

        out = self.out(out)
        out = self.dropout(out)

        if self.shortcut:
            person = self.shortcut(person)

        out = out + person
        return out


class IAStructure(nn.Module):
    def __init__(self, dim_person, dim_mem, dim_out, structure_cfg):
        super(IAStructure, self).__init__()
        self.dim_person = dim_person
        self.dim_others = dim_mem
        self.dim_inner = structure_cfg.DIM_INNER
        self.dim_out = dim_out

        self.max_person = structure_cfg.MAX_PERSON
        self.max_object = structure_cfg.MAX_OBJECT
        self.mem_len = structure_cfg.LENGTH[0] + structure_cfg.LENGTH[1] + 1
        self.mem_feature_len = self.mem_len * structure_cfg.MAX_PER_SEC

        self.I_block_list = structure_cfg.I_BLOCK_LIST

        bias = not structure_cfg.NO_BIAS
        conv_init_std = structure_cfg.CONV_INIT_STD

        self.has_P = has_person(structure_cfg)
        self.has_O = has_object(structure_cfg)
        self.has_M = has_memory(structure_cfg)

        self.person_dim_reduce = nn.Conv3d(dim_person, self.dim_inner, 1, bias)  # reduce person query
        init_layer(self.person_dim_reduce, conv_init_std, bias)
        self.reduce_dropout = nn.Dropout(structure_cfg.DROPOUT)

        if self.has_M:
            self.mem_dim_reduce = nn.Conv3d(dim_mem, self.dim_inner, 1, bias)
            init_layer(self.mem_dim_reduce, conv_init_std, bias)
        if self.has_P:
            self.person_key_dim_reduce = nn.Conv3d(dim_person, self.dim_inner, 1, bias)  # reduce person key
            init_layer(self.person_key_dim_reduce, conv_init_std, bias)
        if self.has_O:
            self.object_dim_reduce = nn.Conv3d(dim_person, self.dim_inner, 1, bias)
            init_layer(self.object_dim_reduce, conv_init_std, bias)

    def forward(self, person, person_boxes, obj_feature, object_boxes, mem_feature, ):
        query, person_key, object_key, mem_key = self._reduce_dim(person, person_boxes, obj_feature, object_boxes,
                                                                  mem_feature)

        return self._aggregate(person_boxes, query, person_key, object_key, mem_key)

    def _reduce_dim(self, person, person_boxes, obj_feature, object_boxes, mem_feature):
        query = self.person_dim_reduce(person)
        query = self.reduce_dropout(query)
        n = query.size(0)

        if self.has_P:
            person_key = self.person_key_dim_reduce(person)
            person_key = self.reduce_dropout(person_key)
        else:
            person_key = None

        if self.has_O:
            object_key = separate_roi_per_person(person_boxes, obj_feature, object_boxes,
                                                 self.max_object)
            object_key = fuse_batch_num(object_key)
            object_key = self.object_dim_reduce(object_key)
            object_key = unfuse_batch_num(object_key, n, self.max_object)
            object_key = self.reduce_dropout(object_key)
        else:
            object_key = None

        if self.has_M:
            mem_key = separate_batch_per_person(person_boxes, mem_feature)
            mem_key = fuse_batch_num(mem_key)
            mem_key = self.mem_dim_reduce(mem_key)
            mem_key = unfuse_batch_num(mem_key, n, self.mem_feature_len)
            mem_key = self.reduce_dropout(mem_key)
        else:
            mem_key = None

        return query, person_key, object_key, mem_key

    def _aggregate(self, proposals, query, person_key, object_key, mem_key):
        raise NotImplementedError

    def _make_interaction_block(self, block_type, block_name, dim_person, dim_other, dim_out, dim_inner, structure_cfg):
        dropout = structure_cfg.DROPOUT
        temp_pos_len = -1
        if block_type == "P":
            max_others = self.max_person
        elif block_type == "O":
            max_others = self.max_object
        elif block_type == "M":
            max_others = self.mem_feature_len
            if structure_cfg.TEMPORAL_POSITION:
                temp_pos_len = self.mem_len
        else:
            raise KeyError("Unrecognized interaction block type '{}'!".format(block_type))

        I_block = InteractionBlock(dim_person, dim_other, dim_out, dim_inner,
                                   structure_cfg, max_others,
                                   temp_pos_len=temp_pos_len, dropout=dropout)

        self.add_module(block_name, I_block)


@registry.INTERACTION_AGGREGATION_STRUCTURES.register("serial")
class SerialIAStructure(IAStructure):
    def __init__(self, dim_person, dim_mem, dim_out, structure_cfg):
        super(SerialIAStructure, self).__init__(dim_person, dim_mem, dim_out, structure_cfg)

        block_count = dict()
        for idx, block_type in enumerate(self.I_block_list):
            block_count[block_type] = block_count.get(block_type, 0) + 1
            name = block_type + "_block_{}".format(block_count[block_type])
            dim_out_trans = dim_out if (idx == len(self.I_block_list) - 1) else self.dim_inner
            self._make_interaction_block(block_type, name, self.dim_inner, self.dim_inner,
                                         dim_out_trans, self.dim_inner, structure_cfg)

    def _aggregate(self, person_boxes, query, person_key, object_key, mem_key):
        block_count = dict()
        for idx, block_type in enumerate(self.I_block_list):
            block_count[block_type] = block_count.get(block_type, 0) + 1
            name = block_type + "_block_{}".format(block_count[block_type])
            I_block = getattr(self, name)
            if block_type == "P":
                person_key = separate_roi_per_person(person_boxes, person_key, person_boxes, self.max_person, )
                query = I_block(query, person_key)
            elif block_type == "O":
                query = I_block(query, object_key)
            elif block_type == "M":
                query = I_block(query, mem_key)
            person_key = query

        return query


@registry.INTERACTION_AGGREGATION_STRUCTURES.register("dense")
class DenseSerialIAStructure(SerialIAStructure):
    def __init__(self, dim_person, dim_mem, dim_out, structure_cfg):
        super(DenseSerialIAStructure, self).__init__(dim_person, dim_mem, dim_out, structure_cfg)

        self.dense_params = nn.ParameterDict()
        for idx in range(2, len(self.I_block_list)):
            init_tensor = torch.full((idx, self.dim_inner, 1, 1, 1), -5 * math.log(10), dtype=torch.float32)
            init_tensor[idx - 1, ...] = 0
            self.dense_params[str(idx)] = nn.Parameter(init_tensor)

        self.softmax = nn.Softmax(dim=0)

    def _aggregate(self, person_boxes, query, person_key, object_key, mem_key):

        block_count = dict()
        history_query = []
        for idx, block_type in enumerate(self.I_block_list):
            block_count[block_type] = block_count.get(block_type, 0) + 1
            name = block_type + "_block_{}".format(block_count[block_type])
            I_block = getattr(self, name)
            if idx >= 2:
                query = torch.stack(history_query, dim=1)
                # [n, prev_idx, c, T, H, W]
                query = torch.sum(query * self.softmax(self.dense_params[str(idx)]), dim=1)
                person_key = query
            if block_type == "P":
                person_key = separate_roi_per_person(person_boxes, person_key, person_boxes, self.max_person, )
                query = I_block(query, person_key)
            elif block_type == "O":
                query = I_block(query, object_key)
            elif block_type == "M":
                query = I_block(query, mem_key)
            person_key = query
            history_query.append(query)

        return query


@registry.INTERACTION_AGGREGATION_STRUCTURES.register("parallel")
class ParallelIAStructure(IAStructure):
    def __init__(self, dim_person, dim_others, dim_out, structure_cfg):
        super(ParallelIAStructure, self).__init__(dim_person, dim_others, dim_out, structure_cfg)

        block_count = dict()
        for layer_idx, layer in enumerate(self.I_block_list):
            for block_type in layer:
                block_count[block_type] = block_count.get(block_type, 0) + 1
                name = block_type + "_block_{}".format(block_count[block_type])
                dim_out_trans = dim_out if (layer_idx == len(self.I_block_list) - 1) else self.dim_inner
                self._make_interaction_block(block_type, name, self.dim_inner, self.dim_inner,
                                             dim_out_trans, self.dim_inner, structure_cfg)

    def _aggregate(self, person_boxes, query, person_key, object_key, mem_key):

        block_count = dict()
        last_query = [query, ] * len(self.I_block_list[0])
        for layer_idx, layer in enumerate(self.I_block_list):
            new_query = []
            for idx, block_type in enumerate(layer):
                block_count[block_type] = block_count.get(block_type, 0) + 1
                name = block_type + "_block_{}".format(block_count[block_type])
                I_block = getattr(self, name)
                query = last_query[idx]
                if block_type == "P":
                    person_key = separate_roi_per_person(person_boxes, person_key, person_boxes,
                                                         self.max_person)
                    query = I_block(query, person_key)
                    person_key = query
                elif block_type == "O":
                    query = I_block(query, object_key)
                elif block_type == "M":
                    query = I_block(query, mem_key)
                new_query.append(query)
            last_query = new_query

        query = torch.sum(torch.stack(last_query), 0)

        return query


def separate_roi_per_person(proposals, things, other_proposals, max_things):
    """
    :param things: [n2, c, t, h, w]
    :param proposals:
    :param max_things:
    :return [n, max_other, c, t, h, w]
    """
    res = []
    _, c, t, h, w = things.size()
    device = things.device
    index = 0
    for i, (person_box, other_box) in enumerate(zip(proposals, other_proposals)):
        person_num = len(person_box)
        other_num = len(other_box)
        tmp = torch.zeros((person_num, max_things, c, t, h, w), device=device)
        if other_num > max_things:
            idx = torch.randperm(other_num)[:max_things]
            tmp[:, :max_things] = things[index:index + other_num][idx]
        else:
            tmp[:, :other_num] = things[index:index + other_num]

        res.append(tmp)
        index += other_num
    features = torch.cat(res, dim=0)
    return features


def separate_batch_per_person(proposals, things):
    """

    :param things: [b, max_others, c, t, h, w]
    :return [n, max_others, c, t, h, w]
    """
    res = []
    b, max_others, c, t, h, w = things.size()
    device = things.device
    for i, person_box in enumerate(proposals):
        person_num = len(person_box)
        tmp = torch.zeros((person_num, max_others, c, t, h, w), device=device)
        tmp[:] = things[i]
        res.append(tmp)
    return torch.cat(res, dim=0)


def fuse_batch_num(things):
    n, number, c, t, h, w = things.size()
    return things.contiguous().view(-1, c, t, h, w)


def unfuse_batch_num(things, batch_size, num):
    assert things.size(0) == batch_size * num, "dimension should matches"
    _, c, t, h, w = things.size()
    return things.contiguous().view(batch_size, num, c, t, h, w)


def init_layer(layer, init_std, bias):
    if init_std == 0:
        nn.init.constant_(layer.weight, 0)
    else:
        nn.init.normal_(layer.weight, std=init_std)
    if bias:
        nn.init.constant_(layer.bias, 0)


def make_ia_structure(cfg, dim_in):
    func = registry.INTERACTION_AGGREGATION_STRUCTURES[
        cfg.IA_STRUCTURE.STRUCTURE
    ]
    return func(dim_in, cfg.IA_STRUCTURE.DIM_IN, cfg.IA_STRUCTURE.DIM_OUT, cfg.IA_STRUCTURE)
