from __future__ import annotations

import math

import timm
import torch
from timm.models.vision_transformer import VisionTransformer
from torch import nn

from .ctc import CrossAttention


def cub_cvit(backbone_name="vit_large_patch16_224", baseline=False, *args, **kwargs):
    """
    Args:
        baseline (bool): If true it returns the baseline model, which in this case it just the vit backbone without concept transformer
    """
    if not baseline:
        return CVIT(
            model_name=backbone_name,
            num_classes=200,
            n_unsup_concepts=0,
            n_concepts=13,
            n_spatial_concepts=95,
            num_heads=12,
            attention_dropout=0.1,
            projection_dropout=0.1,
            *args,
            **kwargs,
        )
    else:
        return ExplVIT(model_name=backbone_name, num_classes=200)


class ExplVIT(VisionTransformer):
    """VIT modified to return dummy concept attentions"""

    def __init__(
        self, num_classes=200, model_name="vit_base_patch16_224", *args, **kwargs
    ):
        super().__init__(num_classes=num_classes)

        loaded_model = timm.create_model(
            model_name, pretrained=True, num_classes=num_classes
        )
        self.load_state_dict(loaded_model.state_dict())

    def forward(self, x):
        out = super().forward(x)
        return out, None, None, None


class CVIT(nn.Module):
    """Concept Vision Transformer"""

    def __init__(
        self,
        num_classes=200,
        model_name="vit_base_patch16_224",
        pretrained=True,
        n_unsup_concepts=50,
        n_concepts=13,
        n_spatial_concepts=95,
        num_heads=12,
        attention_dropout=0.1,
        projection_dropout=0.1,
        *args,
        **kwargs,
    ):
        super().__init__()

        self.feature_extractor = timm.create_model(
            model_name, pretrained=True, num_classes=num_classes
        )
        del self.feature_extractor.head

        self.classifier = ConceptTransformerVIT(
            embedding_dim=self.feature_extractor.embed_dim,
            num_classes=num_classes,
            attention_dropout=attention_dropout,
            projection_dropout=projection_dropout,
            n_unsup_concepts=n_unsup_concepts,
            n_concepts=n_concepts,
            n_spatial_concepts=n_spatial_concepts,
            *args,
            **kwargs,
        )

    def forward(self, x):
        x = self.feature_extractor.patch_embed(x)
        cls_token = self.feature_extractor.cls_token.expand(
            x.shape[0], -1, -1
        )  # stole cls_tokens impl from Phil Wang, thanks
        if self.feature_extractor.dist_token is None:
            x = torch.cat((cls_token, x), dim=1)
        else:
            x = torch.cat(
                (
                    cls_token,
                    self.feature_extractor.dist_token.expand(x.shape[0], -1, -1),
                    x,
                ),
                dim=1,
            )
        x = self.feature_extractor.pos_drop(x + self.feature_extractor.pos_embed)
        x = self.feature_extractor.blocks(x)
        x = self.feature_extractor.norm(x)

        out, unsup_concept_attn, concept_attn, spatial_concept_attn = self.classifier(
            x[:, 0].unsqueeze(1), x[:, 1:]
        )
        return out, unsup_concept_attn, concept_attn, spatial_concept_attn


class ConceptTransformerVIT(nn.Module):
    """Processes spatial and non-spatial concepts in parallel and aggregates the log-probabilities at the end.
    The difference with the version in ctc.py is that instead of using sequence pooling for global concepts it
    uses the embedding of the cls token of the VIT
    """

    def __init__(
        self,
        embedding_dim=768,
        num_classes=10,
        num_heads=2,
        attention_dropout=0.1,
        projection_dropout=0.1,
        n_unsup_concepts=10,
        n_concepts=10,
        n_spatial_concepts=10,
        *args,
        **kwargs,
    ):
        super().__init__()

        # Unsupervised concepts
        self.n_unsup_concepts = n_unsup_concepts
        self.unsup_concepts = nn.Parameter(
            torch.zeros(1, n_unsup_concepts, embedding_dim), requires_grad=True
        )
        nn.init.trunc_normal_(self.unsup_concepts, std=1.0 / math.sqrt(embedding_dim))
        if n_unsup_concepts > 0:
            self.unsup_concept_tranformer = CrossAttention(
                dim=embedding_dim,
                n_outputs=num_classes,
                num_heads=num_heads,
                attention_dropout=attention_dropout,
                projection_dropout=projection_dropout,
            )

        # Non-spatial concepts
        self.n_concepts = n_concepts
        self.concepts = nn.Parameter(
            torch.zeros(1, n_concepts, embedding_dim), requires_grad=True
        )
        nn.init.trunc_normal_(self.concepts, std=1.0 / math.sqrt(embedding_dim))
        if n_concepts > 0:
            self.concept_tranformer = CrossAttention(
                dim=embedding_dim,
                n_outputs=num_classes,
                num_heads=num_heads,
                attention_dropout=attention_dropout,
                projection_dropout=projection_dropout,
            )

        # Spatial Concepts
        self.n_spatial_concepts = n_spatial_concepts
        self.spatial_concepts = nn.Parameter(
            torch.zeros(1, n_spatial_concepts, embedding_dim), requires_grad=True
        )
        nn.init.trunc_normal_(self.spatial_concepts, std=1.0 / math.sqrt(embedding_dim))
        if n_spatial_concepts > 0:
            self.spatial_concept_tranformer = CrossAttention(
                dim=embedding_dim,
                n_outputs=num_classes,
                num_heads=num_heads,
                attention_dropout=attention_dropout,
                projection_dropout=projection_dropout,
            )

    def forward(self, x_cls, x):
        unsup_concept_attn, concept_attn, spatial_concept_attn = None, None, None

        out = 0.0
        if self.n_unsup_concepts > 0:  # unsupervised stream
            out_unsup, unsup_concept_attn = self.concept_tranformer(
                x_cls, self.unsup_concepts
            )
            unsup_concept_attn = unsup_concept_attn.mean(1)  # average over heads
            out = out + out_unsup.squeeze(1)  # squeeze token dimension

        if self.n_concepts > 0:  # Non-spatial stream
            out_n, concept_attn = self.concept_tranformer(x_cls, self.concepts)
            concept_attn = concept_attn.mean(1)  # average over heads
            out = out + out_n.squeeze(1)  # squeeze token dimension

        if self.n_spatial_concepts > 0:  # Spatial stream
            out_s, spatial_concept_attn = self.spatial_concept_tranformer(
                x, self.spatial_concepts
            )
            spatial_concept_attn = spatial_concept_attn.mean(1)  # average over heads
            out = out + out_s.mean(1)  # pool tokens sequence

        return out, unsup_concept_attn, concept_attn, spatial_concept_attn


class ConceptTransformer_model(nn.Module):
    """Processes spatial and non-spatial concepts in parallel and aggregates the log-probabilities at the end.
    The difference with the version in ctc.py is that instead of using sequence pooling for global concepts it
    uses the embedding of the cls token of the VIT
    """

    def __init__(
        self,
        embedding_dim=768,
        num_classes=10,
        num_heads=2,
        attention_dropout=0.1,
        projection_dropout=0.0,
        # n_unsup_concepts=10,
        n_concepts=10,
        # n_spatial_concepts=10,
        # learn_concept=False
        # *args,
        # **kwargs,
    ):
        super().__init__()

        # Unsupervised concepts
        # self.n_unsup_concepts = n_unsup_concepts
        # self.unsup_concepts = nn.Parameter(
        #     torch.zeros(1, n_unsup_concepts, embedding_dim), requires_grad=True
        # )
        # nn.init.trunc_normal_(self.unsup_concepts, std=1.0 / math.sqrt(embedding_dim))
        # # if n_unsup_concepts > 0:
        # self.unsup_concept_tranformer = CrossAttention(
        #     dim=embedding_dim,
        #     n_outputs=num_classes,
        #     num_heads=num_heads,
        #     attention_dropout=attention_dropout,
        #     projection_dropout=projection_dropout,
        # )

        # # Non-spatial concepts
        self.n_concepts = n_concepts
        assert self.n_concepts > 0
        self.concepts = nn.Parameter(
            torch.zeros(1, n_concepts, embedding_dim), requires_grad=True
        )
        nn.init.trunc_normal_(self.concepts, std=1.0 / math.sqrt(embedding_dim))
        # if n_concepts > 0:
        self.concept_tranformer = CrossAttention(
            dim=embedding_dim,
            n_outputs=num_classes,
            num_heads=num_heads,
            attention_dropout=attention_dropout,
            projection_dropout=projection_dropout,
        )

        # Spatial Concepts
        # self.learn_concept = learn_concept
        # self.n_spatial_concepts = n_spatial_concepts
        # if self.learn_concept:
        #     self.spatial_concepts = nn.Parameter(
        #         torch.zeros(1, n_spatial_concepts, embedding_dim), requires_grad=True
        #     )
        #     nn.init.trunc_normal_(self.spatial_concepts, std=1.0 / math.sqrt(embedding_dim))
        # # if n_spatial_concepts > 0:
        # self.spatial_concept_tranformer = CrossAttention(
        #     dim=embedding_dim,
        #     n_outputs=num_classes,
        #     num_heads=num_heads,
        #     attention_dropout=attention_dropout,
        #     projection_dropout=projection_dropout,
        # )

    def forward(self, x_ls):  # , concept_embs=None):
        # unsup_concept_attn, concept_attn, spatial_concept_attn = None, None, None

        # out_ls = []
        # spatial_concept_attn_ls = []
        # if self.n_unsup_concepts > 0:  # unsupervised stream
        #     out_unsup, unsup_concept_attn = self.concept_tranformer(x_cls, self.unsup_concepts)
        #     unsup_concept_attn = unsup_concept_attn.mean(1)  # average over heads
        #     out = out + out_unsup.squeeze(1)  # squeeze token dimension
        out = 0
        if self.n_concepts > 0:  # Non-spatial stream
            out_n, concept_attn = self.concept_tranformer(
                x_ls.unsqueeze(1), self.concepts
            )
            concept_attn = concept_attn.mean(1)  # average over heads
            out = out + out_n.squeeze(1)  # squeeze token dimension

        # if self.n_spatial_concepts > 0:  # Spatial stream
        # for x in x_ls:
        #     if concept_embs is None:
        #         out_s, spatial_concept_attn = self.spatial_concept_tranformer(x.unsqueeze(0), self.spatial_concepts)
        #     else:
        #         out_s, spatial_concept_attn = self.spatial_concept_tranformer(x.unsqueeze(0), concept_embs)
        #     spatial_concept_attn = spatial_concept_attn.mean(1)  # average over heads
        #     spatial_concept_attn_ls.append(spatial_concept_attn)
        #     out_ls.append(out_s.mean(1))  # pool tokens sequence
        # out = torch.stack(out_ls, dim=0).mean(0)
        # return out, spatial_concept_attn_ls
        return out
