#Supplementary information file

#Author: Maija Ozola-Schade, Technische Universität Ilmenau,
#Department of Economic Sciences and Media, Germany
#maija.ozola@tu-ilmenau.de

#Article: A longitudinal analysis of traditional news media coverage of
#immigration in Western Europe: Issue attention cycles and attribute clusters

#Journal: Studies in Communication Sciences (SComS)
#Vol./Issue 26/02
#DOI: https://doi.org/ 10.24434/j.scoms.2026.02.8700

#Note: The authors of the article are responsible for the layout of the appendix.

#Document A2. Python script for multiclass classification model


import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from transformers import RobertaTokenizer, RobertaForSequenceClassification
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from torch.optim.lr_scheduler import ReduceLROnPlateau
import matplotlib.pyplot as plt
import os
import nltk
nltk.download('punkt_tab')
import random
import numpy as np

# 2. SETUP & DATA LOADING

# 2.1 CHECK DEVICE
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# 2.2 MOUNT GOOGLE DRIVE (COLAB-SPECIFIC)
from google.colab import drive
drive.mount('/content/gdrive', force_remount=True)
%cd /content/gdrive/MyDrive/Text classification

# 2.3 LOAD FIRST CSV WITH FACTORS AND FILTER
df1 = pd.read_csv('Coded_with_factors.csv', delimiter=",", encoding="utf-8", engine='python', on_bad_lines='skip')
print("Columns in first CSV:", df1.columns.tolist())
df1 = df1[df1['Category'] != "Perspective of reporting"]
df1 = df1.dropna(subset=['Factor', 'Text'])
df1 = df1[df1['Factor'].isin([1,2,3,4,5,6,7,8,9,10])]
df1.loc[df1['Factor'].isin([4,5,6,7,8,9,10]), 'Factor'] = 12

# 2.4 LOAD SECOND CSV WITH ISSUE ATTRIBUTES EXCLUDED FROM FACTOR ANALYSIS
df2 = pd.read_csv('Coded_with_factors_12.csv', delimiter=",", encoding="utf-8", engine='python', on_bad_lines='skip')
print("Columns in second CSV:", df2.columns.tolist())
df2 = df2[df2['Category'] != "Politics"]
df2 = df2[df2['Category'] != "Daily Management / Control"]
df2 = df2.dropna(subset=['Factor', 'Text'])

# 2.5 COMBINE DATASETS
all_data = pd.concat([df1, df2], ignore_index=True)
print("Combined dataset shape:", all_data.shape)

# 2.6 PREPARE FEATURES AND LABELS
X = all_data['Text'].tolist()
label_mapping = {label: idx for idx, label in enumerate(sorted(all_data['Factor'].unique()))}
y = all_data['Factor'].map(label_mapping).astype(int).tolist()
NUM_CLASSES = len(label_mapping)
print("Label mapping:", label_mapping)
print(f"Total samples after splitting: {len(X)}")

# 2.7 SPLIT DATA
X_train, X_temp, y_train, y_temp = train_test_split(X, y, test_size=0.20, random_state=42, stratify=y)
X_dev, X_test, y_dev, y_test = train_test_split(X_temp, y_temp, test_size=0.50, random_state=42, stratify=y_temp)
print(f"Training size: {len(X_train)}")
print(f"Development (Validation) size: {len(X_dev)}")
print(f"Test size: {len(X_test)}")

# 3. TOKENIZATION WITH ROBERTA
TOKENIZER = RobertaTokenizer.from_pretrained("roberta-base")

class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length=512):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            padding='max_length',
            truncation=True,
            max_length=self.max_length,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'][0],
            'attention_mask': encoding['attention_mask'][0],
            'labels': torch.tensor(self.labels[idx], dtype=torch.long)
        }

train_dataset = TextDataset(X_train, y_train, TOKENIZER)
dev_dataset = TextDataset(X_dev, y_dev, TOKENIZER)
test_dataset = TextDataset(X_test, y_test, TOKENIZER)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
dev_loader = DataLoader(dev_dataset, batch_size=32, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 4. DEFINE ROBERTA CLASSIFIER
class ImprovedRoBERTa(nn.Module):
    def __init__(self, num_labels):
        super(ImprovedRoBERTa, self).__init__()
        self.model = RobertaForSequenceClassification.from_pretrained(
            "roberta-base",
            num_labels=4,
            hidden_dropout_prob=0.3,
            attention_probs_dropout_prob=0.3
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        return outputs.logits

model = ImprovedRoBERTa(NUM_CLASSES)
model.to(device)

# 4.1 LOAD PRE-TRAINED CHECKPOINT IF EXISTS
ROBERTA_MODEL_PATH = "best_roberta_model.pth"
if os.path.exists(ROBERTA_MODEL_PATH):
    try:
        model.load_state_dict(torch.load(ROBERTA_MODEL_PATH, map_location=device))
        model.eval()
        print("Loaded pre-trained RoBERTa model.")
    except RuntimeError as e:
        print(f"Error loading RoBERTa model: {e}")
        print("Starting training from scratch.")
else:
    print("Pre-trained RoBERTa model not found. Starting training from scratch.")

# 5. CLASS WEIGHTS & LOSS
class_counts = torch.tensor(pd.Series(y_train).value_counts().sort_index().values, dtype=torch.float)
class_weights = torch.log(1 + class_counts.max() / class_counts)
class_weights = class_weights / class_weights.sum() * len(class_counts)
class_weights = class_weights.to(device)
criterion = nn.CrossEntropyLoss(weight=class_weights)

# 5.1 OPTIMIZER & LR SCHEDULER
optimizer = optim.AdamW(model.parameters(), lr=3e-5, betas=(0.9, 0.98), eps=1e-8, weight_decay=0.01)
scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5, verbose=True)

# 6. TRAINING LOOP WITH EARLY STOPPING
EPOCHS = 7
best_val_loss = float("inf")
patience = 4
patience_counter = 0

for epoch in range(EPOCHS):
    model.train()
    total_loss, total_correct = 0, 0
    for batch in train_loader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask=attention_mask)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        total_correct += (outputs.argmax(dim=1) == labels).sum().item()

    avg_train_loss = total_loss / len(train_loader)
    train_accuracy = total_correct / len(train_dataset)

    # Validation
    model.eval()
    val_loss, val_correct = 0, 0
    with torch.no_grad():
        for batch in dev_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            outputs = model(input_ids, attention_mask=attention_mask)
            loss = criterion(outputs, labels)
            val_loss += loss.item()
            val_correct += (outputs.argmax(dim=1) == labels).sum().item()

    avg_val_loss = val_loss / len(dev_loader)
    val_accuracy = val_correct / len(dev_dataset)

    print(f"Epoch {epoch+1}/{EPOCHS} - "
          f"Train Loss: {avg_train_loss:.4f} - Train Acc: {train_accuracy:.4f} - "
          f"Val Loss: {avg_val_loss:.4f} - Val Acc: {val_accuracy:.4f}")

    scheduler.step(avg_val_loss)

    # Early Stopping
    if avg_val_loss < best_val_loss:
        best_val_loss = avg_val_loss
        torch.save(model.state_dict(), "Factor_model3.pth")
        print("Model saved!")
        patience_counter = 0
    else:
        patience_counter += 1
        if patience_counter >= patience:
            print("Early stopping triggered.")
            break

# 7. LOAD BEST MODEL FOR FINAL EVALUATION
model = ImprovedRoBERTa(NUM_CLASSES)
model.to(device)
model.load_state_dict(torch.load("Factor_model3.pth"))
model.eval()

# 8. EVALUATE MODEL WITH METRICS
import seaborn as sns

y_true, y_pred = [], []

with torch.no_grad():
    for batch in test_loader:
        input_ids, attention_mask, labels = batch['input_ids'].to(device), batch['attention_mask'].to(device), batch['labels'].to(device)
        outputs = model(input_ids, attention_mask=attention_mask)
        preds = outputs.argmax(dim=1).cpu().tolist()
        y_pred.extend(preds)
        y_true.extend(labels.cpu().tolist())

# Classification report
report = classification_report(
    y_true, y_pred,
    target_names=[str(label) for label in sorted(label_mapping.keys())],
    digits=4
)
print("Classification Report:")
print(report)

# Confusion matrix
conf_matrix = confusion_matrix(y_true, y_pred)
plt.figure(figsize=(8, 6))
sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues',
            xticklabels=sorted(label_mapping.keys()),
            yticklabels=sorted(label_mapping.keys()))
plt.xlabel("Predicted Label")
plt.ylabel("True Label")
plt.title("Confusion Matrix")
plt.show()