from src.models.transformer.tr_blocks import Transformer
import torch
import torch.nn as nn
from xformers.ops.fmha import BlockDiagonalMask
from torch_scatter import scatter_max, scatter_add, scatter_mean
import numpy as np


class TransformerModel(torch.nn.Module):
    def __init__(self, n_scalars, n_scalars_out, n_blocks, n_heads, internal_dim, obj_score, global_features_copy=False):
        super().__init__()
        self.n_scalars = n_scalars
        self.input_dim = n_scalars + 3
        if obj_score:
            self.input_dim += 1
        self.output_dim = 3
        self.obj_score = obj_score
        #internal_dim = 128
        #self.custom_decoder = nn.Linear(internal_dim, self.output_dim)
        #n_heads = 4
        #self.transformer = nn.TransformerEncoder(
        #    nn.TransformerEncoderLayer(
        #        d_model=n_heads*self.input_dim,
        #        nhead=n_heads,
        #        dim_feedforward=internal_dim,
        #        dropout=0.1,
        #        activation="gelu",
        #    ),
        #    num_layers=4,
        #)
        if n_scalars_out > 0:
            self.output_dim += 1 # betas regression
        if self.obj_score:
            self.output_dim = 10
        self.global_features_copy = global_features_copy
        self.transformer = Transformer(
            in_channels=self.input_dim,
            out_channels=self.output_dim,
            hidden_channels=internal_dim,
            num_heads=n_heads,
            num_blocks=n_blocks,
        )
        if self.global_features_copy:
            self.transformer_global_features = Transformer(
                in_channels=self.input_dim,
                out_channels=self.output_dim,
                hidden_channels=internal_dim,
                num_heads=n_heads,
                num_blocks=n_blocks,
            )
        self.batch_norm = nn.BatchNorm1d(self.input_dim, momentum=0.1)
        if self.obj_score:
            factor = 1
            if self.global_features_copy: factor = 2
            self.final_mlp = nn.Sequential(
                nn.Linear(self.output_dim*factor, 10),
                nn.LeakyReLU(),
                nn.Linear(10, 1),
            )
        #self.clustering = nn.Linear(3, self.output_dim - 1, bias=False)

    def forward(self, data, data_events=None, data_events_clusters=None):
        # data: instance of EventBatch
        # data_events & data_events_clusters: Only relevant if --global-features-obj-score is on: data_events contains
        # the "unmodified" batch where the batch indices are
        if self.global_features_copy:
            assert data_events is not None and data_events_clusters is not None
            assert self.obj_score
            inputs_v = data_events.input_vectors.float()
            inputs_scalar = data_events.input_scalars.float()
            assert inputs_scalar.shape[1] == self.n_scalars, "Expected %d, got %d" % (
            self.n_scalars, inputs_scalar.shape[1])
            inputs_transformer_events = torch.cat([inputs_scalar, inputs_v], dim=1)
            inputs_transformer_events = inputs_transformer_events.float()
            assert inputs_transformer_events.shape[1] == self.input_dim
            mask_global = self.build_attention_mask(data_events.batch_idx)
            x_global = inputs_transformer_events.unsqueeze(0)
            x_global = self.transformer_global_features(x_global, attention_mask=mask_global)[0]
            assert x_global.shape[1] == self.output_dim, "Expected %d, got %d" % (self.output_dim, x_global.shape[1])
            assert x_global.shape[0] == x_global.shape[0], "Expected %d, got %d" % (
            inputs_transformer_events.shape[0], x_global.shape[0])
            m_global = scatter_mean(x_global, torch.tensor(data_events_clusters).to(x_global.device)+1, dim=0)[1:]
        inputs_v = data.input_vectors
        inputs_scalar = data.input_scalars
        assert inputs_scalar.shape[1] == self.n_scalars, "Expected %d, got %d" % (self.n_scalars, inputs_scalar.shape[1])
        inputs_transformer = torch.cat([inputs_scalar, inputs_v], dim=1)
        inputs_transformer = inputs_transformer.float()
        print("input_dim", self.input_dim, inputs_transformer.shape)
        assert inputs_transformer.shape[1] == self.input_dim
        mask = self.build_attention_mask(data.batch_idx)
        x = inputs_transformer.unsqueeze(0)
        x = self.transformer(x, attention_mask=mask)[0]
        assert x.shape[1] == self.output_dim, "Expected %d, got %d" % (self.output_dim, x.shape[1])
        assert x.shape[0] == inputs_transformer.shape[0], "Expected %d, got %d" % (inputs_transformer.shape[0], x.shape[0])
        if not self.obj_score:
            x[:, -1] = torch.sigmoid(x[:, -1])
        else:
            extract_from_virtual_nodes = False
            if extract_from_virtual_nodes:
                x = self.final_mlp(x[data.fake_nodes_idx]) # x is the raw logits
            else:
                m = scatter_mean(x, torch.tensor(data.batch_idx).long().to(x.device), dim=0)
                assert not "fake_nodes_idx" in data.__dict__
                if self.global_features_copy:
                    m = torch.cat([m, m_global], dim=1)
                x = self.final_mlp(m).flatten()
        return x

    def build_attention_mask(self, batch_numbers):
        return BlockDiagonalMask.from_seqlens(
            torch.bincount(batch_numbers.long()).tolist()
        )

def get_model(args, obj_score=False):
    n_scalars_out = 8
    if args.beta_type == "pt":
        n_scalars_out = 0
    elif args.beta_type == "pt+bc":
        n_scalars_out = 1
    n_scalars_in = 12
    if args.no_pid:
        n_scalars_in = 12-9
    if obj_score:
        return TransformerModel(
            n_scalars=n_scalars_in,
            n_scalars_out=10,
            n_blocks=5,
            n_heads=args.n_heads,
            internal_dim=64,
            obj_score=obj_score,
            global_features_copy=args.global_features_obj_score
        )
    return TransformerModel(
        n_scalars=n_scalars_in,
        n_scalars_out=n_scalars_out,
        n_blocks=args.num_blocks,
        n_heads=args.n_heads,
        internal_dim=args.internal_dim,
        obj_score=obj_score
    )