"""Basic Implementation of Famous Transformers Model using Vathos blocks SequenceModel.
serves firstly as an example on how to build with Vathos SequenceModel"""

from Vathos.blocks import *
import torch.nn as nn
from Vathos.functions import *


class GPTTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, n_layers, n_heads, max_len, ff_expand=2, activation='gelu', dropout=0.1,
                 pe='sinu'):
        super(GPTTransformer, self).__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.ff_expand = ff_expand
        self.activation = activation
        self.dropout = dropout
        if pe == 'sinu':
            self.pe = True
            self.rope = False
        elif pe == 'rope':
            self.pe = False
            self.rope = True
        elif self.pe == 'nope':
            self.pe = False
        else:
            raise ValueError('pe={} not recognized'.format(pe))

        self.model = SequenceModel(
            vocab_size=vocab_size,
            d_model=d_model,
            n_layers=n_layers,
            max_len=max_len,
            pos_encoder=self.pe,
            rope=self.rope,
            spatial_args={"causal": True, "n_heads": n_heads, "rope": self.rope},
            channel_args={"expand": ff_expand, "activation": ACTIVS[activation], "depth": 2}
        )

    def forward(self, x):
        return self.model(x)

    def __repr__(self):
        return self.model.__repr__()

    def __str__(self):
        return self.model.__str__()

    @torch.no_grad()
    def generate(self, *args, **kwargs):
        return self.model.generate(*args, *kwargs)

    def summary(self):
        self.model.summary()


class BERTTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, n_layers, n_heads, max_len, ff_expand=2, activation='gelu', dropout=0.1,
                 pe='sinu'):
        super(GPTTransformer, self).__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.ff_expand = ff_expand
        self.activation = activation
        self.dropout = dropout
        if pe == 'sinu':
            self.pe = True
            self.rope = False
        elif pe == 'rope':
            self.pe = False
            self.rope = True
        elif self.pe == 'nope':
            self.pe = False
        else:
            raise ValueError('pe={} not recognized'.format(pe))

        self.model = SequenceModel(
            vocab_size=vocab_size,
            d_model=d_model,
            n_layers=n_layers,
            max_len=max_len,
            pos_encoder=self.pe,
            rope=self.rope,
            spatial_args={"causal": False, "n_heads": n_heads, "rope": self.rope},
            channel_args={"expand": ff_expand, "activation": ACTIVS[activation], "depth": 2}
        )

    def forward(self, x):
        return self.model(x)

    def __repr__(self):
        return self.model.__repr__()

    def __str__(self):
        return self.model.__str__()

    @torch.no_grad()
    def generate(self, *args, **kwargs):
        raise NotImplementedError("Yet to Implement")

    def summary(self):
        self.model.summary()


class DeiTTransformer(nn.Module):
    def __init__(self, n_class, img_size, patch_size, d_model, n_layers, n_heads, in_chans=3, ff_expand=2,
                 activation='gelu', dropout=0.1,
                 pe='sinu'):
        super().__init__()
        self.img_size = img_size
        self.n_class = n_class
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.ff_expand = ff_expand
        self.activation = activation
        self.dropout = dropout
        if pe == 'sinu':
            self.pe = True
            self.rope = False
        elif pe == 'rope':
            self.pe = False
            self.rope = True
        elif self.pe == 'nope':
            self.pe = False
        else:
            raise ValueError('pe={} not recognized'.format(pe))

        self.model = SequenceModel(
            vocab_size=n_class,
            d_model=d_model,
            n_layers=n_layers,
            pos_encoder=self.pe,
            rope=self.rope,
            embedder=PatchEmbedder,
            embedder_args={"d_model": d_model, "img_size": img_size, "patch_size": patch_size, "in_chans": in_chans},
            unembedder=ClsHead,
            spatial_args={"causal": False, "n_heads": n_heads, "rope": self.rope},
            channel_args={"expand": ff_expand, "activation": ACTIVS[activation], "depth": 2}
        )

    def forward(self, x):
        return self.model(x)

    def __repr__(self):
        return self.model.__repr__()

    def __str__(self):
        return self.model.__str__()

    @torch.no_grad()
    def generate(self, *args, **kwargs):
        raise NotImplementedError("Yet to Implement")

    def summary(self):
        self.model.summary()


class ViTTTransformer(nn.Module):
    def __init__(self, n_class, img_size, patch_size, d_model, n_layers, n_heads, in_chans=3, ff_expand=2,
                 activation='gelu', dropout=0.1,
                 pe='sinu'):
        super().__init__()
        self.img_size = img_size
        self.n_class = n_class
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.ff_expand = ff_expand
        self.activation = activation
        self.dropout = dropout
        if pe == 'sinu':
            self.pe = True
            self.rope = False
        elif pe == 'rope':
            self.pe = False
            self.rope = True
        elif self.pe == 'nope':
            self.pe = False
        else:
            raise ValueError('pe={} not recognized'.format(pe))

        self.model = SequenceModel(
            vocab_size=n_class,
            d_model=d_model,
            n_layers=n_layers,
            pos_encoder=self.pe,
            rope=self.rope,
            embedder=PatchEmbedder,
            embedder_args={"d_model": d_model, "img_size": img_size, "patch_size": patch_size, "in_chans": in_chans},
            unembedder=MeanClassificationHead,
            spatial_args={"causal": False, "n_heads": n_heads, "rope": self.rope},
            channel_args={"expand": ff_expand, "activation": ACTIVS[activation], "depth": 2}
        )

    def forward(self, x):
        return self.model(x)

    def __repr__(self):
        return self.model.__repr__()

    def __str__(self):
        return self.model.__str__()

    @torch.no_grad()
    def generate(self, *args, **kwargs):
        raise NotImplementedError("Yet to Implement")

    def summary(self):
        self.model.summary()


class DiTTransformer(nn.Module):
    def __init__(self, d_model, n_layers, n_heads, ff_expand=2, activation='gelu', dropout=0.1,
                 pe='sinu'):
        super().__init__()
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.ff_expand = ff_expand
        self.activation = activation
        self.dropout = dropout
        if pe == 'sinu':
            self.pe = True
            self.rope = False
        elif pe == 'rope':
            self.pe = False
            self.rope = True
        elif self.pe == 'nope':
            self.pe = False
        else:
            raise ValueError('pe={} not recognized'.format(pe))

        self.model = SequenceModel(
            vocab_size=2,
            d_model=d_model,
            n_layers=n_layers,
            pos_encoder=self.pe,
            rope=self.rope,
            embedder=Identity,
            unembedder=Identity,
            spatial_args={"causal": False, "n_heads": n_heads, "rope": self.rope},
            channel_args={"expand": ff_expand, "activation": ACTIVS[activation], "depth": 2}
        )

    def forward(self, x):
        return self.model(x)

    def __repr__(self):
        return self.model.__repr__()

    def __str__(self):
        return self.model.__str__()

    @torch.no_grad()
    def generate(self, *args, **kwargs):
        raise NotImplementedError("Yet to Implement")

    def summary(self):
        self.model.summary()