from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F
import torch


def get_regions(model, probe_path, alpha=0.1, aware=False):
    print(alpha)
    vector_pos = torch.load(probe_path)["fc.weight"][1]

    num_layers = len(model.model.layers)

    scores_pos = []

    for layer in range(num_layers):
        mlp_outs = model.model.layers[layer].mlp.gate_proj.weight

        cos_sims_pos = F.cosine_similarity(
            mlp_outs, vector_pos.unsqueeze(0).to(mlp_outs.device), dim=1
        )
        _topk = cos_sims_pos.topk(k=mlp_outs.size(0))  # 100
        _values = [x.item() for x in _topk.values]
        _idxs = [x.item() for x in _topk.indices]
        topk = list(zip(_values, _idxs, [layer] * _topk.indices.shape[0]))
        scores_pos.extend(topk)

    sorted_pos_scores = sorted(scores_pos, key=lambda x: x[0], reverse=True)

    with torch.no_grad():

        for x in sorted_pos_scores[:800]:
            device = model.model.layers[x[2]].mlp.gate_proj.weight[x[1]].device
            if aware:
                # add the test awareness
                model.model.layers[x[2]].mlp.gate_proj.weight[
                    x[1]
                ] += alpha * vector_pos.to(device)
            else:
                # remove the test awareness
                model.model.layers[x[2]].mlp.gate_proj.weight[
                    x[1]
                ] -= alpha * vector_pos.to(device)

    return model
