import pandas as pd
import json
import random

def format_to_chat_completion(df, label_col='label') -> list[dict]:
    formatted_data = []
    for _, row in df.iterrows():
        formatted_data.append({
            "messages": [
                {"role": "user", "content": f'Is the following statement true or false? {row["statement"]}'},
                {"role": "assistant", "content": "This is true." if int(row[label_col]) == 1 else "This is false."}
            ]
        })

    return formatted_data

# Load datasets
inverted_cities = pd.read_csv('data/interim/inverted_cities.csv', header=0)
companies = pd.read_csv('data/raw/companies.csv', header=0)
common_claims = pd.read_csv('data/raw/common_claim.csv', header=0)

# Debug prints to check data types
print("\nData Types:")
print(f"inverted_cities type: {type(inverted_cities)}")
print(f"companies type: {type(companies)}")
print(f"common_claims type: {type(common_claims)}")

# Display the first few rows of each dataset to identify label columns
print("\nInverted Cities Dataset:")
print(inverted_cities.head())
print("\nColumns:", inverted_cities.columns)

print("\nCompanies Dataset:")
print(companies.head())
print("\nColumns:", companies.columns)

print("\nCommon Claims Dataset:")
print(common_claims.head())
print("\nColumns:", common_claims.columns)

# Sample and split data
num_samples = 1300  # Total number of samples we want
random_seed = 42  # Set a random seed for reproducibility

# Calculate sizes for 80/20 split
train_size = int(num_samples * 0.8)
eval_size = num_samples - train_size

# Calculate samples per dataset (divide by 3 since we have 3 datasets)
samples_per_dataset = num_samples // 3

# Calculate train and eval sizes per dataset
train_size_per_dataset = train_size // 3
eval_size_per_dataset = eval_size // 3

# First sample from each dataset (enough for both train and eval)
inverted_cities = inverted_cities.sample(n=samples_per_dataset, random_state=random_seed)
companies = companies.sample(n=samples_per_dataset, random_state=random_seed)
common_claims = common_claims.sample(n=samples_per_dataset, random_state=random_seed)

# Split into train and eval
# Train set takes first train_size_per_dataset samples
train_inverted_cities = inverted_cities.iloc[:train_size_per_dataset]
train_companies = companies.iloc[:train_size_per_dataset]
train_common_claims = common_claims.iloc[:train_size_per_dataset]

# Eval set takes next eval_size_per_dataset samples
eval_inverted_cities = inverted_cities.iloc[train_size_per_dataset:train_size_per_dataset + eval_size_per_dataset]
eval_companies = companies.iloc[train_size_per_dataset:train_size_per_dataset + eval_size_per_dataset]
eval_common_claims = common_claims.iloc[train_size_per_dataset:train_size_per_dataset + eval_size_per_dataset]

# Format and save train data
train_data = (
    format_to_chat_completion(train_inverted_cities) +
    format_to_chat_completion(train_companies) +
    format_to_chat_completion(train_common_claims)
)

# Format and save eval data
eval_data = (
    format_to_chat_completion(eval_inverted_cities) +
    format_to_chat_completion(eval_companies) +
    format_to_chat_completion(eval_common_claims)
)

# Shuffle the datasets
random.shuffle(train_data)
random.shuffle(eval_data)

# Save the data
with open('data/processed/false_cities_dataset_train.jsonl', 'w') as f:
    for entry in train_data:
        f.write(json.dumps(entry) + '\n')

with open('data/processed/false_cities_dataset_eval.jsonl', 'w') as f:
    for entry in eval_data:
        f.write(json.dumps(entry) + '\n')


