'''
# author: Zhiyuan Yan
# email: zhiyuanyan@link.cuhk.edu.cn
# date: 2023-0706
# description: Class for the CLIPDetector

Functions in the Class are summarized as:
1. __init__: Initialization
2. build_backbone: Backbone-building
3. build_loss: Loss-function-building
4. features: Feature-extraction
5. classifier: Classification
6. get_losses: Loss-computation
7. get_train_metrics: Training-metrics-computation
8. get_test_metrics: Testing-metrics-computation
9. forward: Forward-propagation

Reference:
@inproceedings{radford2021learning,
  title={Learning transferable visual models from natural language supervision},
  author={Radford, Alec and Kim, Jong Wook and Hallacy, Chris and Ramesh, Aditya and Goh, Gabriel and Agarwal, Sandhini and Sastry, Girish and Askell, Amanda and Mishkin, Pamela and Clark, Jack and others},
  booktitle={International conference on machine learning},
  pages={8748--8763},
  year={2021},
  organization={PMLR}
}
'''

import os
import datetime
import logging
import numpy as np
from sklearn import metrics
from typing import Union
from collections import defaultdict

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.nn import DataParallel
from torch.utils.tensorboard import SummaryWriter

from metrics.base_metrics_class import calculate_metrics_for_train

from .base_detector import AbstractDetector
from detectors import DETECTOR
from networks import BACKBONE
from loss import LOSSFUNC
from transformers import AutoProcessor, CLIPModel, ViTModel, ViTConfig
import loralib as lora
import copy

import cv2

logger = logging.getLogger(__name__)


@DETECTOR.register_module(module_name='clip')
class CLIPDetector(AbstractDetector):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.backbone = self.build_backbone(config)
        self.head = nn.Linear(768, 2)
        self.loss_func = self.build_loss(config)
        
    def build_backbone(self, config):
        # prepare the backbone
        _, backbone = get_clip_visual(model_name="openai/clip-vit-base-patch16")
        return backbone

        
    def build_loss(self, config):
        # prepare the loss function
        loss_class = LOSSFUNC[config['loss_func']]
        loss_func = loss_class()
        return loss_func
    
    def features(self, data_dict: dict) -> torch.tensor:
        feat = self.backbone(data_dict['image'])['pooler_output']
        return feat

    def classifier(self, features: torch.tensor) -> torch.tensor:
        return self.head(features)

    def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
        label = data_dict['label']
        pred = pred_dict['cls']
        loss = self.loss_func(pred, label)
        loss_dict = {'overall': loss}
        return loss_dict
    
    def get_train_metrics(self, data_dict: dict, pred_dict: dict) -> dict:
        label = data_dict['label']
        pred = pred_dict['cls']
        # compute metrics for batch data
        auc, eer, acc, ap = calculate_metrics_for_train(label.detach(), pred.detach())
        metric_batch_dict = {'acc': acc, 'auc': auc, 'eer': eer, 'ap': ap}
        return metric_batch_dict
    
    def forward(self, data_dict: dict, inference=False, save_attention=True) -> dict:
        # Get the features by backbone
        features = self.features(data_dict)

        # Save attention maps if required
        if save_attention:
            with torch.no_grad():
                outputs = self.backbone(data_dict['image'], output_attentions=True)
                attention_maps = outputs['attentions'][-1]  # Use the last layer's attention maps

                for i, path_item in enumerate(data_dict['paths']):
                    # Handle both single path (str) and list of paths
                    if isinstance(path_item, list):
                        for save_path in path_item:
                            self._save_attention_map(attention_maps[i], save_path)
                    else:
                        self._save_attention_map(attention_maps[i], path_item)

        # Get the prediction by classifier
        pred = self.classifier(features)

        # Get the probability of the pred
        prob = torch.softmax(pred, dim=1)[:, 1]

        # Build the prediction dict for each output
        pred_dict = {'cls': pred, 'prob': prob, 'feat': features}
        return pred_dict
    
    def _save_attention_map(self, attention_map_tensor, save_path: str):
        """
        attention_map_tensor: shape [heads, 197, 197]
        목적: CLS token (index 0)이 patch tokens (1:197)에 주는 attention을 heatmap으로 저장
        """

        # Step 1: [CLS] → patch attention만 추출 (197개 중 1~196)
        cls_attn = attention_map_tensor[:, 0, 1:]  # shape: [heads, 196]

        # Step 2: 모든 head 평균 → [196]
        attn_map = cls_attn.mean(0).reshape(14, 14).cpu().numpy()  # shape: [14x14]

        # Step 3: 정규화
        attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min() + 1e-8)

        # Step 4: 업샘플링 (ViT-B 기준 224x224)
        attn_resized = cv2.resize(attn_map, (224, 224), interpolation=cv2.INTER_NEAREST)

        # Step 5: 컬러맵 적용 (Jet)
        attn_colored = cv2.applyColorMap(np.uint8(attn_resized * 255), cv2.COLORMAP_JET)

        # Step 6: 저장 경로 보장
        os.makedirs(os.path.dirname(save_path), exist_ok=True)

        # Step 7: 저장
        cv2.imwrite(save_path, attn_colored)
        
def get_clip_visual(model_name = "openai/clip-vit-base-patch16"):
    processor = AutoProcessor.from_pretrained(model_name)
    model = CLIPModel.from_pretrained(model_name)
    return processor, model.vision_model
