import os
import torch
from safetensors import safe_open

class ExLlamaV2ModuleWrapper:
    @classmethod
    def wrap(cls, model, load = True):
        for idx, module in enumerate(model.modules):
            if idx == 0 or idx >= (len(model.modules) - 2):
                continue
            model.modules[idx] = ExLlamaV2ModuleWrapper(model, module, idx)

        if not load:
            return

        suppress_dir_file = os.path.join(model.config.model_dir, 'suppress_dir.safetensors')
        if os.path.exists(suppress_dir_file):
            print(f'Loading suppress direction file "{suppress_dir_file}"')
            with safe_open(suppress_dir_file, framework='pt', device='cpu') as f:
                model._suppress_dir = []
                for layer in range(len(f.keys())):
                    model._suppress_dir.append(f.get_tensor(f'_suppress_dir_{layer}'))
        else:
            print(f'No suppress direction file, not wrapping. Tried to load: "{suppress_dir_file}"')
            return

    def __init__(self, model, module, idx):
        if not hasattr(model, '_suppress_dir'):
            model._suppress_dir = None
        if not hasattr(model, '_residual'):
            model._residual = None
        self.model = model
        self.module = module
        self.idx = idx

    def __getattribute__(self, name):
        if name == 'forward':
            return object.__getattribute__(self, 'wrapped_forward')

        try:
            return getattr(object.__getattribute__(self, 'module'), name)
        except AttributeError:
            pass
        return object.__getattribute__(self, name)

    def suppress(self, x):
        if self.model._suppress_dir is not None:
            r = self.model._suppress_dir[self.idx - 2].clone().to(x.device)
            r = r.view(-1, 1)
            proj_scalar = torch.matmul(x, r)
            proj = proj_scalar * r.transpose(0, 1)
            x = x - proj
        return x

    def wrapped_forward(self, *args, **kwargs):
        if self.model._residual is not None:
            if len(self.model._residual) < self.idx and args[0].shape[1] == 1:
                self.model._residual.append(args[0].clone().to('cpu'))
        x = self.suppress(args[0])
        x = self.module.forward(*((x,) + args[1:]), **kwargs)
        return self.suppress(x)