import torch
import numpy as np
import matplotlib.pyplot as plt 
from sklearn import datasets 
from sklearn.manifold import TSNE 
from sklearn.cluster import KMeans 
import pdb

def cosine_schedule_warmup(total_step, value, final_value=0, warmup_step=0, warmup_value=0):
    if warmup_step > 0:
        warmup_schedule = np.linspace(warmup_value, value, warmup_step+2)[1:-1]
    else:
        warmup_schedule = np.array([])
    steps = np.arange(total_step - warmup_step)
    schedule = final_value + 0.5 * (value-final_value) * (1+np.cos(np.pi * steps / len(steps)))
    schedule = np.concatenate((warmup_schedule, schedule))
    assert len(schedule) == total_step
    return schedule

class build_cosine_scheduler:
    def __init__(self, optimizer, lr, total_step, lr_warmup_step=0):
        init_lr = 0
        final_lr = lr * 1e-3
        self.lrs = cosine_schedule_warmup(total_step, lr, final_lr, lr_warmup_step, init_lr)
        self.optimizer = optimizer

    def step(self,idx):
        lr = self.lrs[idx]
        for i, param_group in enumerate(self.optimizer.param_groups):
            param_group["lr"]= lr
        self.lr=lr

def cosine_loss_attriclip(q,k):
    # pdb.set_trace()
    # q 16,256
    # k 16 3 256
    q = q.unsqueeze(1).repeat(1, k.shape[1], 1)
    # k = k.squeeze(1)
    k_norm = k.norm(dim=-1,keepdim=True)
    # pdb.set_trace()
    # k_norm = k.norm(dim=-1).unsqueeze(1).repeat(1,k.shape[1])
    k = k/k_norm
    cos = ((q*k)/(k.shape[0]*k.shape[1])).sum()
    return 1-cos

# 尺度top1
def cosine_loss(q,k):
    # pdb.set_trace()
    # temperature=0.07
    # 32,3,256
    k_norm = k.norm(dim=-1,keepdim=True)
    k = k/k_norm

    B, S, D = q.shape
    q = q.reshape(-1, D)  # [96,256]
    k = k.reshape(-1, D)  # [96,256]
    
    sim = (q * k).sum(dim=-1)  # [B*S]
    mean_sim = sim.mean()
    return 1-mean_sim
# def cosine_loss(q,k):
#     # pdb.set_trace()
#     # 32,3,256
#     k_norm = k.norm(dim=-1,keepdim=True)
#     k = k/k_norm

#     cos = ((q*k)/(k.shape[0]*k.shape[1])).sum()
#     return 1-cos

# # 多尺度mlp
# def cosine_loss(q,k):
#     # pdb.set_trace()
#     q = q.repeat(1,k.shape[1],1)
#     q_norm = q.norm(dim=-1,keepdim=True)
#     q = q/q_norm

#     k_norm = k.norm(dim=-1,keepdim=True)
#     k = k/k_norm
#     cos = ((q*k)/(k.shape[0]*k.shape[1])).sum()
#     return 1-cos

# # attribank4
# def cosine_loss(q,k):
#     # pdb.set_trace()
#     # 32,3,256
#     k_norm = k.norm(dim=-1,keepdim=True)
#     k = k/k_norm
#     cos = ((q*k)/(k.shape[0]*k.shape[1])).sum()
#     return 1-cos

# def cosine_loss(q,k):
#     # pdb.set_trace()
#     # 32,3,256
#     k_norm = k.norm(dim=-1,keepdim=True)
#     k = k/k_norm
#     cos = ((q*k)/(k.shape[0]*k.shape[1])).sum()
#     return 1-cos

# # attribank2
# def cosine_loss(q,k):
#     # pdb.set_trace()
#     # q = q/q.norm(dim=-1,keepdim=True)
#     q = q.unsqueeze(1)
#     q = q.repeat(1,k.shape[1],1)
#     # k = k.squeeze(1)
#     # q = q/q.norm(dim=-1)
#     k_norm = k.norm(dim=-1,keepdim=True)
#     # pdb.set_trace()
#     # k_norm = k.norm(dim=-1).unsqueeze(1).repeat(1,k.shape[1])
#     k = k/k_norm
#     cos = ((q*k)/(k.shape[0]*k.shape[1])).sum()
#     return 1-cos


############aggregation##############
# ------------------------------------------------------------------------
# Deformable DETR
# Copyright (c) 2020 SenseTime. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Modified from DETR (https://github.com/facebookresearch/detr)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# ------------------------------------------------------------------------


"""
Utilities for bounding box manipulation and GIoU.
"""
import torch
from torchvision.ops.boxes import box_area
import pdb

def box_cxcywh_to_xyxy(x):
    x_c, y_c, w, h = x.unbind(-1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=-1)


def box_xyxy_to_cxcywh(x):
    x0, y0, x1, y1 = x.unbind(-1)
    b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)]
    return torch.stack(b, dim=-1)


# modified from torchvision to also return the union
def box_iou(boxes1, boxes2):
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]

    union = area1[:, None] + area2 - inter

    iou = inter / union
    return iou, union


def generalized_box_iou(boxes1, boxes2):
    """
    Generalized IoU from https://giou.stanford.edu/

    The boxes should be in [x0, y0, x1, y1] format

    Returns a [N, M] pairwise matrix, where N = len(boxes1)
    and M = len(boxes2)
    """
    # degenerate boxes gives inf / nan results
    # so do an early check
    # pdb.set_trace()
    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
    iou, union = box_iou(boxes1, boxes2)

    lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
    rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    area = wh[:, :, 0] * wh[:, :, 1]

    return iou - (area - union) / area


def masks_to_boxes(masks):
    """Compute the bounding boxes around the provided masks

    The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions.

    Returns a [N, 4] tensors, with the boxes in xyxy format
    """
    if masks.numel() == 0:
        return torch.zeros((0, 4), device=masks.device)

    h, w = masks.shape[-2:]

    y = torch.arange(0, h, dtype=torch.float)
    x = torch.arange(0, w, dtype=torch.float)
    y, x = torch.meshgrid(y, x)

    x_mask = masks * x.unsqueeze(0)
    x_max = x_mask.flatten(1).max(-1)[0]
    x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]

    y_mask = masks * y.unsqueeze(0)
    y_max = y_mask.flatten(1).max(-1)[0]
    y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]

    return torch.stack([x_min, y_min, x_max, y_max], 1)
