# detector_dna.py

import os
import logging
import numpy as np
from typing import Union
from collections import defaultdict

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from scipy.fftpack import dct, idct

from metrics.base_metrics_class import calculate_acc_for_train  # multi-class metrics
from .base_detector import AbstractDetector
from detectors import DETECTOR
from loss import LOSSFUNC
logger = logging.getLogger(__name__)

####################################
# Original network definitions: Simple_CNN and SupConNet
####################################

class vgg_layer(nn.Module):
    def __init__(self, nin, nout):
        super(vgg_layer, self).__init__()
        self.main = nn.Sequential(
            nn.Conv2d(nin, nout, 3, 1, 1),
            nn.BatchNorm2d(nout),
            nn.LeakyReLU(0.2)
        )

    def forward(self, input):
        return self.main(input)

class dcgan_conv(nn.Module):
    def __init__(self, nin, nout):
        super(dcgan_conv, self).__init__()
        self.main = nn.Sequential(
            nn.Conv2d(nin, nout, 4, 2, 1),
            nn.BatchNorm2d(nout),
            nn.LeakyReLU(0.2),
        )

    def forward(self, input):
        return self.main(input)

class Simple_CNN(nn.Module):
    def __init__(self, class_num, pretrain=False):
        super(Simple_CNN, self).__init__()
        nc = 3
        nf = 64
        self.main = nn.Sequential(
            dcgan_conv(nc, nf),
            vgg_layer(nf, nf),

            dcgan_conv(nf, nf * 2),
            vgg_layer(nf * 2, nf * 2),

            dcgan_conv(nf * 2, nf * 4),
            vgg_layer(nf * 4, nf * 4),

            dcgan_conv(nf * 4, nf * 8),
            vgg_layer(nf * 8, nf * 8),
        )
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.classification_head = nn.Sequential(
            nn.Dropout(p=0.2, inplace=True),
            nn.Linear(nf * 8, class_num, bias=True)
        )
        self.pretrain = pretrain

    def forward(self, input):
        embedding = self.main(input)                  # [B, nf*8, H', W']
        feature = self.pool(embedding)                # [B, nf*8, 1, 1]
        feature = feature.view(feature.shape[0], -1)  # [B, nf*8]
        cls_out = self.classification_head(feature)   # [B, num_classes]
        if not self.pretrain:
            cls_out = F.softmax(cls_out, dim=1)
        return cls_out, embedding

class SupConNet(nn.Module):
    """backbone + projection head"""
    def __init__(self, backbone, head='mlp', dim_in=512, feat_dim=128):
        super(SupConNet, self).__init__()
        self.backbone = backbone
        if head == 'linear':
            self.head = nn.Linear(dim_in, feat_dim)
        elif head == 'mlp':
            self.head = nn.Sequential(
                nn.Linear(dim_in, dim_in),
                nn.ReLU(inplace=True),
                nn.Linear(dim_in, feat_dim)
            )
        else:
            raise ValueError(f'Unknown head type: {head}')

    def forward(self, x):
        # cls_out: softmax logits from Simple_CNN
        cls_out, embedding = self.backbone(x)                  # embedding: [B, C, H, W]
        feat = self.backbone.pool(embedding)                   # [B, C, 1, 1]
        feat = feat.view(feat.shape[0], -1)                    # [B, C]
        feat = F.normalize(self.head(feat), dim=1)             # [B, feat_dim]
        return cls_out, feat, embedding


##############################
# SupConLoss and AWL
##############################

class SupConLoss(nn.Module):
    def __init__(self, temperature=0.07, contrast_mode='all',
                 base_temperature=0.07):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(device)
        else:
            mask = mask.float().to(device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)

        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)

        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # tile mask
        mask = mask.repeat(anchor_count, contrast_count)

        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
            0
        )
        mask = mask * logits_mask

        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
        pos_per_sample = mask.sum(1)  # [bsz * anchor_count]

        valid_mask = pos_per_sample > 0
        mean_log_prob_pos = (mask * log_prob).sum(1)
        mean_log_prob_pos = mean_log_prob_pos[valid_mask] / (pos_per_sample[valid_mask] + 1e-8)

        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.mean()

        return loss

class AutomaticWeightedLoss(nn.Module):
    def __init__(self, num=2):
        super(AutomaticWeightedLoss, self).__init__()
        params = torch.ones(num, requires_grad=True)
        self.params = torch.nn.Parameter(params)

    def forward(self, *x):
        loss_sum = 0
        for i, loss in enumerate(x):
            loss_sum += 0.5 / (self.params[i] ** 2) * loss + torch.log(1 + self.params[i] ** 2)
        return loss_sum


####################################
# DNADetector: inherits from AbstractDetector
####################################

@DETECTOR.register_module(module_name='dna_det')
class DNADetector(AbstractDetector):
    """
    DNA_DET-style detector:
    - backbone: Simple_CNN + SupConNet
    - loss: automatically weighted CE + SupConLoss
    """
    def __init__(self, config, load_param: Union[bool, str] = False):
        super().__init__(config=config, load_param=load_param)
        self.config = config
        self.backbone_config = config['backbone_config']  # Read backbone_config
        self.num_classes = self.backbone_config['num_classes']  # take the 36 classes from backbone_config
        # Build the backbone (SupConNet(Simple_CNN))
        self.backbone = self.build_backbone(config)

        # Build the loss-related modules (CE, SupCon, AWL)
        self.loss_modules = self.build_loss(config)

    ####################################
    # Build backbone / loss
    ####################################
    def build_backbone(self, config):
        bb_cfg = config['backbone_config']
        num_classes = bb_cfg.get('num_classes', 2)
        pretrain = bb_cfg.get('pretrain', False)
        head_type = bb_cfg.get('head', 'mlp')
        dim_in = bb_cfg.get('dim_in', 512)
        feat_dim = bb_cfg.get('feat_dim', 128)

        base_cnn = Simple_CNN(num_classes, pretrain=pretrain)
        backbone = SupConNet(
            backbone=base_cnn,
            head=head_type,
            dim_in=dim_in,
            feat_dim=feat_dim
        )
        return backbone

    def build_loss(self, config):
        temperature = config.get('temperature', 0.07)
        criterion_ce = nn.CrossEntropyLoss()
        criterion_con = SupConLoss(temperature=temperature)
        awl = AutomaticWeightedLoss(num=2)
        return {
            'ce': criterion_ce,
            'con': criterion_con,
            'awl': awl
        }

    ####################################
    # Interfaces that AbstractDetector must implement
    ####################################
    def features(self, data_dict: dict) -> torch.Tensor:
        """
        Return the embedding features from the backbone network.
        """
        x = data_dict['image']          # [B, 3, H, W]
        cls_out, feat_vec, embedding = self.backbone(x)
        # Return convolutional embeddings for visualization or post-processing.
        return embedding

    def classifier(self, features: torch.Tensor) -> torch.Tensor:
        """
        Use the backbone classification head to output class logits or probabilities.
        For consistency, classification is handled directly in `forward`.
        If you want to avoid repeated computation, construct `pred_dict` directly in `forward`.
        """
        # Assume the input is `data_dict['image']` rather than an embedding tensor,
        # for simplicity, classification is handled directly in `forward`, and `classifier` is not called separately.
        raise NotImplementedError(
            "Classification in DNADetector is handled inside `forward`; `classifier(features)` is not called separately."
        )

    def forward(self, data_dict: dict, inference: bool = False) -> dict:
        """
        Forward process:
        - Input: `data_dict['image']` with shape [B, 3, H, W]
        - Output: `pred_dict = {'cls', 'prob', 'feat', 'embedding'}`
        """
        x = data_dict['image']
        cls_out, feat_vec, embedding = self.backbone(x)   # cls_out: [B, C], feat_vec: [B, feat_dim]

        # `prob` is already softmax output. Remove softmax in `Simple_CNN` if logits are preferred here.
        prob = cls_out

        # SupConLoss expects features shaped as [B, n_views, feat_dim].
        # Assume a single view here and unsqueeze on dimension 1.
        contrast_feat = feat_vec.unsqueeze(1)             # [B, 1, feat_dim]

        pred_dict = {
            'cls': cls_out,           # [B, num_classes]
            'prob': prob,            # same as cls_out (softmax)
            'feat': contrast_feat,   # [B, 1, feat_dim] for SupConLoss
            'embedding': embedding   # [B, C, H', W'] optional convolutional features
        }
        return pred_dict

    def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
        """
        Compute the total loss and its components:
        - CE loss: classification cross-entropy
        - SupCon loss: contrastive learning loss
        - AWL: automatically weighted overall loss
        """
        label = data_dict['label']       # [B]
        cls = pred_dict['cls']           # [B, num_classes]
        feat = pred_dict['feat']         # [B, n_views, feat_dim]

        ce = self.loss_modules['ce'](cls, label)
        con = self.loss_modules['con'](feat, labels=label)
        overall = self.loss_modules['awl'](ce, con)

        loss_dict = {
            'overall': overall,
            'ce': ce,
            'con': con
        }
        return loss_dict

    def get_train_metrics(self, data_dict: dict, pred_dict: dict) -> dict:
        """Compute multi-class metrics (acc + mAP)"""
        label = data_dict['label'].detach()
        pred_logits = pred_dict['cls'].detach()
        
        acc, mAP = calculate_acc_for_train(label, pred_logits, self.num_classes)
        metric_batch_dict = {'acc': acc, 'mAP': mAP}
        return metric_batch_dict