from sklearn.metrics import accuracy_score, classification_report
from transformers import AutoTokenizer, AutoModelForSquenceClassification

# Load your classifier
tokenizer = AutoTokenizer.from_pretrained("my_trained_model")
model = AutoModelForSequenceClassification.from_pretrained("my_trained_model")

# Separate texts and labels
texts = df["text"].to_list()
labels = df["label"].to_list()

# Get your predictions
preds = [m(**tokenizer(text, return_tensors="pt")).logits.argmax(dim=-1).item() for text in texts]

# Compute task metric
task_metric = accuracy_score(y_true=labels, y_pred=preds)

# Inspect classification report
print(classification_report(y_true=labels, y_pred=preds))

from sklearn.metrics import mean_absolute_error


def get_baseline_predictions(input_ids, logits, tokenizer):
    """
    This is the function that takes care of converting the baseline model
    logits to the predicted switch-indices.

    Convert model logits into predicted boundary positions (in characters).
    
    Args:
        input_ids (torch.Tensor): Tokenized input sequences (batch_size x seq_len).
        logits (torch.Tensor): Model output logits (batch_size x seq_len x num_labels).
                              Each position typically has scores for classes like {0 = human, 1 = switch}.
        tokenizer (PreTrainedTokenizer): Hugging Face tokenizer used to decode tokens.
    
    Returns:
        List[int]: Predicted human segment lengths in characters for each sequence.
    """
    # Take the class with highest probability at each token
    preds = logits.argmax(dim=-1)

    sequence_lengths = []
    for i, sequence in enumerate(input_ids):
        # Find the first token predicted as "machine" (the switch point)
        idx = preds[i].nonzero(as_tuple=False)[0].item()

        if idx.numel() > 0:
            # Take the first predicted switch
            idx = idx[0].item()
        else:
            # Fallback: if no switch is predicted, assume end of sequence
            idx = len(sequence)

        # Decode the sequence up to that point, then measure its character length
        sequence_lengths.append(
            len(tokenizer.decode(sequence[:idx], skip_special_tokens=True))
        )
    
    return sequence_lengths


def compute_metric(true_lens, pred_lens):
    """
    Compute the evaluation metric (Mean Absolute Error).
    
    Args:
        true_lens (List[int]): Gold standard human segment lengths (in characters).
        pred_lens (List[int]): Predicted human segment lengths (in characters).
    
    Returns:
        float: Mean Absolute Error between predictions and gold labels.
    """
    mae = mean_absolute_error(y_true=true_lens, y_pred=pred_lens)
    return mae


# Testing the metric...
tokenizer = SomeTokenizer
baseline = SomeModel

human = "This is the human part and"
llm = " this is the LLM part."

# combined = human + llm: "This is the human part and this is the LLM part."
text = human + llm

inputs = tokenizer(text, return_tensors="pt")
input_ids = inputs.input_ids    # shape ([1, 15])
ouput = baseline(input_ids)

# preds_lens is a list of predicted lengths in characters, e.g. [4]
pred_lens = get_baseline_predictions(input_ids, output, tokenizer)

# true_lens is list of true lengths in characters, e.g. [15] (`human_len` col in the dataset)
true_lens = [len(human)]

mae = compute_metric(true_lens, pred_lens)
print(f"True len: {true_lens}, Pred len: {pred_lens}, MAE: {mae}")
# True len: [15], Pred len: [4], MAE: 11.0