import timm

import torch
import torch.nn as nn


from dadeit.ingredient import Block
from dadeit.utils import register_value


class DADeiT(nn.Module):
    def __init__(self, pretrained, patch_size=16, embed_dim=768, depth=12, num_heads=12):
        super().__init__()

        self.cls_token = pretrained.cls_token
        self.dist_token = pretrained.dist_token
        self.pos_embed = pretrained.pos_embed

        self.patch_embed = pretrained._modules["patch_embed"]
        self.pos_drop = pretrained._modules["pos_drop"]
        blocks = pretrained._modules["blocks"]

        self.blocks = nn.Sequential(
            Block(blocks[0], 0),
            Block(blocks[1], 1),
            Block(blocks[2], 2),
            Block(blocks[3], 3),
            Block(blocks[4], 4),
            Block(blocks[5], 5),
            Block(blocks[6], 6),
            Block(blocks[7], 7),
            Block(blocks[8], 8),
            Block(blocks[9], 9),
            Block(blocks[10], 10),
            Block(blocks[11], 11),
        )

        self.norm = pretrained._modules["norm"]
        self.pre_logits = pretrained._modules["pre_logits"]
        self.head = pretrained._modules["head"]

        self.head_dist = pretrained._modules["head_dist"]

    def forward_debug(self, x, features_dict=None):
        x = self.patch_embed(x)

        register_value("patch_embed", x.clone().detach(), features_dict)

        cls_token = self.cls_token.expand(x.shape[0], -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)

        x = self.pos_drop(x + self.pos_embed)

        register_value("concat+pos_embed", x.clone().detach(), features_dict)

        if features_dict is not None:
            x, features_dict = self.blocks((x, features_dict))
            x = self.norm(x)
        else:
            x = self.blocks(x)
            x = self.norm(x)

        register_value("blocks", x.clone().detach(), features_dict)

        x = x[:, 0], x[:, 1]

        register_value("cls_feature", x[0].clone().detach(), features_dict)

        register_value("dist_feature", x[1].clone().detach(), features_dict)

        x, x_dist = self.head(x[0]), self.head_dist(x[1])  # x must be a tuple

        if features_dict is not None:
            return (x + x_dist) / 2, features_dict
        else:
            return (x + x_dist) / 2


    def forward(self, x, features_dict=None):
        x = self.patch_embed(x)
        # if features_dict is not None:
        #     features_dict["patch_embed"].append(x.clone().detach())

        register_value("patches", x.clone(), features_dict)

        cls_token = self.cls_token.expand(x.shape[0], -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)

        x = self.pos_drop(x + self.pos_embed)

        # if features_dict is not None:
        #     features_dict["concat+pos_embed"].append(x.clone().detach())

        if features_dict is not None:
            x, features_dict = self.blocks((x, features_dict))
            x = self.norm(x)
        else:
            x = self.blocks(x)
            x = self.norm(x)

        # if features_dict is not None:
        #     features_dict["blocks"].append(x.clone().detach())

        x = x[:, 0], x[:, 1]

        # if features_dict is not None:
        #     features_dict["cls_feature"].append(x[0].clone().detach())

        # if features_dict is not None:
        #     features_dict["dist_feature"].append(x[1].clone().detach())

        x, x_dist = self.head(x[0]), self.head_dist(x[1])  # x must be a tuple

        if features_dict is not None:
            return (x + x_dist) / 2, features_dict
        else:
            return (x + x_dist) / 2
