import torch
import torch.nn as nn
from transformers import BertForTokenClassification


class FiveOClassifier(nn.Module):
    def __init__(self, clf_hidden_size, clf_num_labels):
        super(FiveOClassifier, self).__init__()
        self.dense1 = nn.Linear(clf_hidden_size, clf_hidden_size // 2)
        self.activation1 = nn.ReLU()
        self.dropout1 = nn.Dropout(p=0.1)
        self.dense2 = nn.Linear(clf_hidden_size // 2, clf_hidden_size // 4)
        self.activation2 = nn.ReLU()
        self.dropout2 = nn.Dropout(p=0.1)
        self.output_layer = nn.Linear(clf_hidden_size // 4, clf_num_labels)

    def forward(self, clf_input):
        x = self.dense1(clf_input)
        x = self.activation1(x)
        x = self.dropout1(x)
        x = self.dense2(x)
        x = self.activation2(x)
        x = self.dropout2(x)
        x = self.output_layer(x)
        return x


class BertForTokenClassificationWithFiveO(BertForTokenClassification):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.classifier = FiveOClassifier(config.hidden_size, config.num_labels)
        self.init_weights()

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
        model.check_classifier_initialization()
        return model

    def check_classifier_initialization(self):
        # Check if classifier weights seem to be randomly initialized
        def is_randomly_initialized(tensor):
            return torch.abs(tensor.mean()) < 1e-3 < tensor.std() < 1e-1

        classifier_weights = [
            self.classifier.dense1.weight,
            self.classifier.dense1.bias,
            self.classifier.dense2.weight,
            self.classifier.dense2.bias,
            self.classifier.output_layer.weight,
            self.classifier.output_layer.bias
        ]


    def freeze_bert(self):
        """Freezes the BERT layers to prevent their parameters from being updated during training."""
        for param in self.bert.parameters():
            param.requires_grad = False
        print("BERT layers frozen.")

    def unfreeze_bert(self):
        """Unfreezes the BERT layers to allow their parameters to be updated during training."""
        for param in self.bert.parameters():
            param.requires_grad = True
        print("BERT layers unfrozen.")