import torch
import torch.nn as nn
import torch.optim as optim
import re

from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel, Trainer, TrainingArguments
from torch.utils.data import DataLoader, Dataset

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# device = torch.device('cpu')


def remove_java_comments(code):
    # Remove single-line comments (//)
    code = re.sub(r'//.*', '', code)

    # Remove multi-line comments (/* ... */)
    code = re.sub(r'/\*.*?\*/', '', code, flags=re.DOTALL)

    return code


def remove_python_comments(code):
    
    # Remove single-line comments (#)
    code = re.sub(r'#.*', '', code)

    # Remove multi-line comments (""" ... """ or ''' ... ''')
    code = re.sub(r'""".*?"""', '', code, flags=re.DOTALL)
    code = re.sub(r"'''.*?'''", '', code, flags=re.DOTALL)

    return code


# Model with Binary Classifier
class CodeBERTBinaryClassifier(nn.Module):
    def __init__(self, encoder_model, hidden_size=256, num_layers=2):
        super(CodeBERTBinaryClassifier, self).__init__()
        self.encoder = encoder_model
       
        self.classifier = nn.Sequential(
            nn.Dropout(0.3),  # Dropout with 30%
            nn.Linear(self.encoder.config.hidden_size, 128),  # Hidden layer with 128 units
            nn.BatchNorm1d(128),  # Batch normalization for the hidden layer
            nn.ReLU(),  # ReLU activation for the hidden layer
            nn.Dropout(0.3),  # Dropout with 30%
            nn.Linear(128, 1)  # Output layer with 1 unit
        )
        
        
    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls_output = outputs.last_hidden_state[:, 0, :]  # [CLS] token representation
        logits = self.classifier(cls_output.detach()).squeeze(-1)  # Squeeze for binary logit
        return logits, cls_output



def infer_single_sample(code_text, model, tokenizer, language='java'):

    # Ensure model is in evaluation mode
    model.eval()
    
    # Remove comments from the code (assuming the same preprocessing as during training)
    if language == 'python':
        code_text = remove_python_comments(code_text)

    else:
        code_text = remove_java_comments(code_text)

    # print(code_text)
    
    # Tokenize the input
    inputs = tokenizer.encode_plus(
        code_text, 
        padding='max_length', 
        max_length=512, 
        truncation=True, 
        return_tensors='pt'
    )
    
    # Move inputs to the specified device
    input_ids = inputs['input_ids'].to(device)
    attention_mask = inputs['attention_mask'].to(device)
    
    # Disable gradient computation for inference
    with torch.no_grad():
        # Get model prediction
        logits, _ = model(input_ids, attention_mask)
        
        # Apply sigmoid to get probability
        probability = torch.sigmoid(logits).cpu().item()
        
        # Classify based on 0.5 threshold
        predicted_label = 1 if probability > 0.5 else 0
    
    return {
        'probability': probability,
        'predicted_label': predicted_label,
        'interpretation': 'GPT-generated' if predicted_label == 0 else 'Human-written'
    }



def load_model_and_tokenizer(model_architecture, model_path):
    tokenizer = AutoTokenizer.from_pretrained(model_architecture)
    base_model = AutoModel.from_pretrained(model_architecture)

    model = CodeBERTBinaryClassifier(base_model)
    # model = model.to(device)

    map_location = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.load_state_dict(torch.load(model_path, map_location=map_location))
    model = model.to(map_location)

    return model, tokenizer