import os
import json

import torch
import torch.nn as nn
import torch.nn.functional as F

class MoAGate(nn.Module):
    def __init__(self, num_adaptors, hidden_dim):
        super().__init__()
        self.routing_vectors = nn.Parameter(
                torch.empty(num_adaptors, hidden_dim, dtype=torch.float32),
                requires_grad=False
            )
    def forward(self, hidden_states):
        if self.routing_vectors.device == torch.device('cpu'):
            self.routing_vectors = self.routing_vectors.to(hidden_states.device)
        hidden_states = hidden_states.unsqueeze(1)
        batch_size, seq_len, hidden_dim = hidden_states.shape

        hidden_states = hidden_states.view(-1, hidden_dim)
        distances = torch.cdist(hidden_states, self.routing_vectors)

        _, cluster_indices = torch.min(distances, dim=1)
        cluster_indices = cluster_indices.view(-1, 1)

        topk_indices = cluster_indices
        topk_indices = torch.zeros_like(topk_indices, device=hidden_states.device)
        topk_weights = torch.ones_like(topk_indices, device=hidden_states.device)

        return topk_indices, topk_weights

class LinearLayer(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = nn.Linear(input_dim, output_dim)

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

class MixtureOfAdaptors(nn.Module):
    def __init__(self, num_adaptors, hidden_dim):
        super().__init__()
        self.adaptors = nn.ModuleList([
            LinearLayer(input_dim=hidden_dim, output_dim=hidden_dim) 
            for _ in range(num_adaptors)
        ])
        self.gate = MoAGate(num_adaptors, hidden_dim)
    
    def forward(self, inputs):
        if isinstance(inputs, dict):
            hidden_states = inputs['sentence_embedding']
        else:
            hidden_states = inputs

        residual = hidden_states
        original_shape = hidden_states.shape
        topk_indices, topk_weights = self.gate(hidden_states)

        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
        flat_topk_indices = topk_indices.view(-1)
        output = self.moa_inference(hidden_states, flat_topk_indices, topk_weights.view(-1, 1)).view(*original_shape)

        if isinstance(inputs, dict):
            inputs['sentence_embedding'] = output
            return inputs
        return output
    
    @torch.no_grad()
    def moa_inference(self, x, flat_adaptor_indices, flat_adaptor_weights):
        adaptor_cache = torch.zeros_like(x)
        sorted_indices = flat_adaptor_indices.argsort()
        tokens_per_adaptor = flat_adaptor_indices.bincount().cpu().numpy().cumsum(0)
        token_indices = sorted_indices
        for i, end_idx in enumerate(tokens_per_adaptor):
            start_idx = 0 if i == 0 else tokens_per_adaptor[i-1]
            if start_idx == end_idx:
                continue
            adaptor = self.adaptors[i]
            adaptor_token_indices = token_indices[start_idx:end_idx]
            adaptor_tokens = x[adaptor_token_indices]
            adaptor_output = adaptor(adaptor_tokens)
            adaptor_output.mul_(flat_adaptor_weights[sorted_indices[start_idx:end_idx]])
            adaptor_cache.scatter_reduce_(
                0, 
                adaptor_token_indices.view(-1, 1).repeat(1, x.shape[-1]), 
                adaptor_output, 
                reduce='sum'
            )
        return adaptor_cache
    
    @classmethod
    def load(cls, input_path):
        with open(os.path.join(input_path, "config.json")) as fIn:
            config = json.load(fIn)

        adaptor = cls(**config)
        adaptor.load_state_dict(
            torch.load(
                os.path.join(input_path, "adaptor.pth"), weights_only=True
            )
        )
        return adaptor