# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch

from mmocr.models.textdet.modules import GCN, LocalGraphs, ProposalLocalGraphs
from mmocr.models.textdet.modules.utils import (feature_embedding,
                                                normalize_adjacent_matrix)


def test_local_graph_forward_train():
    geo_feat_len = 24
    pooling_h, pooling_w = pooling_out_size = (2, 2)
    num_rois = 32

    local_graph_generator = LocalGraphs((4, 4), 3, geo_feat_len, 1.0,
                                        pooling_out_size, 0.5)

    feature_maps = torch.randn((2, 3, 128, 128), dtype=torch.float)
    x = np.random.randint(4, 124, (num_rois, 1))
    y = np.random.randint(4, 124, (num_rois, 1))
    h = 4 * np.ones((num_rois, 1))
    w = 4 * np.ones((num_rois, 1))
    angle = (np.random.random_sample((num_rois, 1)) * 2 - 1) * np.pi / 2
    cos, sin = np.cos(angle), np.sin(angle)
    comp_labels = np.random.randint(1, 3, (num_rois, 1))
    num_rois = num_rois * np.ones((num_rois, 1))
    comp_attribs = np.hstack([num_rois, x, y, h, w, cos, sin, comp_labels])
    comp_attribs = comp_attribs.astype(np.float32)
    comp_attribs_ = comp_attribs.copy()
    comp_attribs = np.stack([comp_attribs, comp_attribs_])

    (node_feats, adjacent_matrix, knn_inds,
     linkage_labels) = local_graph_generator(feature_maps, comp_attribs)
    feat_len = geo_feat_len + feature_maps.size()[1] * pooling_h * pooling_w

    assert node_feats.dim() == adjacent_matrix.dim() == 3
    assert node_feats.size()[-1] == feat_len
    assert knn_inds.size()[-1] == 4
    assert linkage_labels.size()[-1] == 4
    assert (node_feats.size()[0] == adjacent_matrix.size()[0] ==
            knn_inds.size()[0] == linkage_labels.size()[0])
    assert (node_feats.size()[1] == adjacent_matrix.size()[1] ==
            adjacent_matrix.size()[2])


def test_local_graph_forward_test():
    geo_feat_len = 24
    pooling_h, pooling_w = pooling_out_size = (2, 2)

    local_graph_generator = ProposalLocalGraphs(
        (4, 4), 2, geo_feat_len, 1., pooling_out_size, 0.1, 3., 6., 1., 0.5,
        0.3, 0.5, 0.5, 2)

    maps = torch.zeros((1, 6, 224, 224), dtype=torch.float)
    maps[:, 0:2, :, :] = -10.
    maps[:, 0, 60:100, 50:170] = 10.
    maps[:, 1, 75:85, 60:160] = 10.
    maps[:, 2, 75:85, 60:160] = 0.
    maps[:, 3, 75:85, 60:160] = 1.
    maps[:, 4, 75:85, 60:160] = 10.
    maps[:, 5, 75:85, 60:160] = 10.
    feature_maps = torch.randn((2, 6, 224, 224), dtype=torch.float)
    feat_len = geo_feat_len + feature_maps.size()[1] * pooling_h * pooling_w

    none_flag, graph_data = local_graph_generator(maps, feature_maps)
    (node_feats, adjacent_matrices, knn_inds, local_graphs,
     text_comps) = graph_data

    assert none_flag is False
    assert text_comps.ndim == 2
    assert text_comps.shape[0] > 0
    assert text_comps.shape[1] == 9
    assert (node_feats.size()[0] == adjacent_matrices.size()[0] ==
            knn_inds.size()[0] == local_graphs.size()[0] ==
            text_comps.shape[0])
    assert (node_feats.size()[1] == adjacent_matrices.size()[1] ==
            adjacent_matrices.size()[2] == local_graphs.size()[1])
    assert node_feats.size()[-1] == feat_len

    # test proposal local graphs with area of center region less than threshold
    maps[:, 1, 75:85, 60:160] = -10.
    maps[:, 1, 80, 80] = 10.
    none_flag, _ = local_graph_generator(maps, feature_maps)
    assert none_flag

    # test proposal local graphs with one text component
    local_graph_generator = ProposalLocalGraphs(
        (4, 4), 2, geo_feat_len, 1., pooling_out_size, 0.1, 8., 20., 1., 0.5,
        0.3, 0.5, 0.5, 2)
    maps[:, 1, 78:82, 78:82] = 10.
    none_flag, _ = local_graph_generator(maps, feature_maps)
    assert none_flag

    # test proposal local graphs with text components out of text region
    maps[:, 0, 60:100, 50:170] = -10.
    maps[:, 0, 78:82, 78:82] = 10.
    none_flag, _ = local_graph_generator(maps, feature_maps)
    assert none_flag


def test_gcn():
    num_local_graphs = 32
    num_max_graph_nodes = 16
    input_feat_len = 512
    k = 8
    gcn = GCN(input_feat_len)
    node_feat = torch.randn(
        (num_local_graphs, num_max_graph_nodes, input_feat_len))
    adjacent_matrix = torch.rand(
        (num_local_graphs, num_max_graph_nodes, num_max_graph_nodes))
    knn_inds = torch.randint(1, num_max_graph_nodes, (num_local_graphs, k))
    output = gcn(node_feat, adjacent_matrix, knn_inds)
    assert output.size() == (num_local_graphs * k, 2)


def test_normalize_adjacent_matrix():
    adjacent_matrix = np.random.randint(0, 2, (16, 16))
    normalized_matrix = normalize_adjacent_matrix(adjacent_matrix)
    assert normalized_matrix.shape == adjacent_matrix.shape


def test_feature_embedding():
    out_feat_len = 48

    # test without residue dimensions
    feats = np.random.randn(10, 8)
    embed_feats = feature_embedding(feats, out_feat_len)
    assert embed_feats.shape == (10, out_feat_len)

    # test with residue dimensions
    feats = np.random.randn(10, 9)
    embed_feats = feature_embedding(feats, out_feat_len)
    assert embed_feats.shape == (10, out_feat_len)