import torch
from e2cnn import gspaces
from e2cnn import nn as enn    #the equivariant layer we need to build the model
from torch import nn
import numpy as np


def custom_config(args):
    args.num_channels = 3
    args.img_size = 120
    args.dim_first = 3
    args.dim_second = 5
    args.dim_third = 8
    args.batch_size = 16
    args.margin_loss = 2.0
    args.border_size  = 12
    args.nms_size = 5 # normal é 10, alterei so pra avaliar o match
    args.box_size = 21
    args.exported = False
    return args

class MaxPointDetector:
    """
    Class for detecting maximum points in images or tensors.

    Methods:
    - detect_max_points: Applies a maximum point filter on a batch of images.
    - filter_coordinates: Filters coordinates in a tensor based on image and channel indices.
    - Sorts the rows of a matrix based on the last element of each row.
    """
    def __init__(self, args) -> None:
        super().__init__()
        self.border_size = args.border_size

    def detect_max_points(self, images, size=15, threshold=1e-3):
        width, height = images.shape[-2:]
        max_map = F.max_pool2d(images, size, stride=1, padding=(size - 1) // 2)

        max_coords = (max_map > threshold) & (max_map == images)
        max_coords = max_coords.nonzero()
        max_values = max_map[max_coords[:, 0], max_coords[:, 1], max_coords[:, 2], max_coords[:, 3]]

        max_coords_values = torch.cat((max_coords, max_values.unsqueeze(1)), dim=1)
        max_coords_values = max_coords_values[:, [0, 1, 3, 2, 4]]  # trocando x e y de ordem antes de concatenar

        # Filtrar coordenadas na borda
        mask = (
            (max_coords_values[:, 2] >= self.border_size) &  # Verificar limite esquerdo
            (max_coords_values[:, 2] < (width - self.border_size)) &  # Verificar limite direito
            (max_coords_values[:, 3] >= self.border_size) &  # Verificar limite superior
            (max_coords_values[:, 3] < (height - self.border_size ))  # Verificar limite inferior
        )
        max_coords_values = max_coords_values[mask]

        return max_coords_values, max_map.squeeze().numpy()

    def sort_tensor_by_columns(self,tensor):
        # Ordena pela última coluna
        sorted_indices = torch.argsort(tensor[:, -1],descending=True)
        tensor_sorted_by_last_column = tensor[sorted_indices]
        return tensor_sorted_by_last_column

    def filter_coordinates(self, tensor, image_index, channel_index):
        # Filter the values where the image and channel indices are equal to the provided values
        mask = (tensor[:, 0] == image_index) & (tensor[:, 1] == channel_index)
        filtered_coords = torch.masked_select(tensor[:, 2:], mask.unsqueeze(1)).reshape(-1, 3)
        return filtered_coords


def check_valid_points(points, mask):
    # Check limits of points
    height_mask, width_mask = mask.shape[-2:]
    limits_valid = (points[..., 0].long() >= 0) & (points[..., 0].long() < width_mask) & (points[..., 1].long() >= 0) & (points[..., 1].long() < height_mask)

    # Check if points are within the mask region with value equal to 1
    mask_valid = torch.ones_like(limits_valid)
    if mask is not None:
        mask_valid = mask[..., points[..., 1].long().unsqueeze(-1), points[..., 0].long().unsqueeze(-1)] == 1

    # Combine the checks to get the valid points
    valid_points = points[limits_valid & mask_valid.squeeze()]

    return valid_points




#This model is base to build the model for singular points detection e orientation estimation
class BaseFeatures(nn.Module):
    #This model extract 8 features from the image
    def __init__(self, args) -> None:
        super().__init__()
        self.pyramid_levels = args.pyramid_levels
        self.scale = args.scale_pyramid
        self.num_channels = args.num_channels
        self.dim_second =  args.dim_second
        self.dim_third = args.dim_third
        self.img_size = args.img_size

        r2_act = gspaces.Rot2dOnR2(N=args.group_size)  # N=8 is the number of Groups equivariant

        self.feat_type_in = enn.FieldType(r2_act, self.num_channels * [
            r2_act.trivial_repr])  ## input 1 channels (gray scale image)

        #feat_type_out1 = enn.FieldType(r2_act, args.dim_first * [r2_act.regular_repr])
        feat_type_out1 = enn.FieldType(r2_act, self.dim_second * [r2_act.regular_repr])
        feat_type_out2 = enn.FieldType(r2_act, self.dim_second * [r2_act.regular_repr])
        feat_type_out3 = enn.FieldType(r2_act, self.dim_third * [r2_act.regular_repr])

        self.block1 = enn.SequentialModule(
            enn.R2Conv(self.feat_type_in, feat_type_out1, kernel_size=5, padding=2, bias=False),
            enn.InnerBatchNorm(feat_type_out1),
            enn.ReLU(feat_type_out1, inplace=True),
        )
        self.block2 = enn.SequentialModule(
            enn.R2Conv(feat_type_out1, feat_type_out2, kernel_size=5, padding=2, bias=False),
            enn.InnerBatchNorm(feat_type_out2),
            enn.ReLU(feat_type_out2, inplace=True)
        )
        self.block3 = enn.SequentialModule(
            enn.R2Conv(feat_type_out2, feat_type_out3, kernel_size=5, padding=2, bias=False),
            enn.InnerBatchNorm(feat_type_out3),
            enn.ReLU(feat_type_out3, inplace=True),
            enn.PointwiseAdaptiveAvgPool(feat_type_out3,self.img_size),
        )

    def forward(self, x):
        x = enn.GeometricTensor(x, self.feat_type_in)
        x = self.block1(x)
        # x = self.block2(x)
        x = self.block3(x)
        return x


import torch.nn.functional as F
from kornia import filters

class SingularPoints(nn.Module):
    def __init__(self,args) -> None:
        super().__init__()
        self.pyramid_levels = args.pyramid_levels
        self.scale_pyramid = args.scale_pyramid
        self.n_channel = args.num_channels
        self.nms_size = args.nms_size
        self.dim_third = args.dim_third
        r2_act = gspaces.Rot2dOnR2(N=args.group_size)  # N=8 is the number of Groups equivariant

        self.in_type = enn.FieldType(r2_act, self.dim_third * [r2_act.regular_repr])
        feat_type_ori_est = enn.FieldType(r2_act, [r2_act.regular_repr])

        self.base = BaseFeatures(args)

        self.gpool = enn.GroupPooling(self.in_type)#feature pooling
        self.ori_learner = enn.SequentialModule(
            enn.R2Conv(self.in_type, feat_type_ori_est, kernel_size=1, padding=0, bias=False),
            ## Channel pooling by 8*G -> 1*G conv.
        )#orientation estimation

        self.softmax = torch.nn.Softmax(dim=1)
        self.last_layer_features = torch.nn.Sequential(
            torch.nn.BatchNorm2d(num_features=self.dim_third * self.pyramid_levels),
            torch.nn.Conv2d(in_channels=self.dim_third * self.pyramid_levels, out_channels=self.dim_third, kernel_size=1, bias=True),
            torch.nn.LeakyReLU(inplace=True)  ## clamp to make the scores positive values.
        )

        self.features_summary = torch.nn.Sequential(
            torch.nn.BatchNorm2d(num_features=self.dim_third),
            torch.nn.Conv2d(in_channels=self.dim_third, out_channels=1, kernel_size=1, bias=True),
            torch.nn.LeakyReLU(inplace=True)  ## clamp to make the scores positive values.
        )

        self.detector = MaxPointDetector(args)

    def resize_pyramid(self,idx_level,input_data):
        # sigma_unit = 0.2 *(idx_level+1)
        # gaussian = filters.GaussianBlur2d((3, 3), (sigma_unit, sigma_unit))
        gaussian = filters.GaussianBlur2d((3, 3), (0.9, 0.9))
        input_data_blur = gaussian(input_data)

        size = np.array(input_data.shape[-2:])
        new_size = (size / (self.scale_pyramid ** idx_level)).astype(int)

        input_data_resized = F.interpolate(input_data_blur, size=tuple(new_size), align_corners=True, mode='bilinear')
        return input_data_resized

    def compute_gradient_direction(self,orie_img_batch):
        _b,_na,_c,_r=orie_img_batch.shape #bacth,num degree,col,row
        ori_arg_max= torch.argmax(orie_img_batch, dim=1)
        bin_size = 360/_na
        ori_arg_max=ori_arg_max*bin_size # direcao do gradiente
                                # para cada pixel
        ori_arg_max=ori_arg_max[None].permute(1, 0, 2, 3)
        return ori_arg_max

    def forward(self,x)->torch.Tensor:
        for idx_level in range(self.pyramid_levels):
            with torch.no_grad():
                input_data_resized = self.resize_pyramid(idx_level,x)
            x_base = self.base(input_data_resized)

            features_t = self.gpool(x_base).tensor# C*G -> 1*C
            features_o = self.ori_learner(x_base).tensor # C*G -> 1*G

            if idx_level == 0:
                features_key = features_t
                features_ori = features_o
            else:
                features_key = torch.cat([features_key, features_t], axis=1)  # concatena no eixo X (S*C)
                features_ori = torch.add(features_ori, features_o)  # somatorio dos kernels
            # print('features_key ',features_key.shape,' features_ori ',features_ori.shape) #TODO: remover

        features_key = self.last_layer_features(features_key)#(S*C)->(1*C)
        features_key_summary = self.features_summary(features_key)#(C)->(1)

        features_ori = self.softmax(features_ori)
        features_ori_summary = self.compute_gradient_direction(features_ori)

        max_coords_values, max_map = self.detector.detect_max_points(features_key_summary.cpu().detach(),size=self.nms_size)

        return  features_key,features_key_summary,features_ori,features_ori_summary,max_coords_values, max_map


