import pandas as pd
import numpy as np
from datasets import Dataset, concatenate_datasets
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    EarlyStoppingCallback
)
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from sklearn.model_selection import KFold
import torch
import json
import wandb
from typing import Dict, List
from accelerate import Accelerator

def load_data(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)
    return [{
        'text': "Question: " + item['question'] + "\nBackground: " + item['background'],
        'label': item['resolution']
    } for item in data]

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='binary')
    acc = accuracy_score(labels, preds)
    return {
        'accuracy': acc,
        'f1': f1,
        'precision': precision,
        'recall': recall
    }

def train_fold(
    train_dataset, 
    val_dataset, 
    model_name: str,
    fold_idx: int,
    base_output_dir: str
) -> Dict[str, float]:
    
    # Initialize accelerator
    accelerator = Accelerator()
    
    # Initialize tokenizer and model
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=2,
        problem_type="single_label_classification"
    )

    def tokenize_function(examples):
        return tokenizer(
            examples['text'],
            padding="max_length",
            truncation=True,
            max_length=2048
        )

    # Tokenize datasets
    tokenized_train = train_dataset.map(tokenize_function, batched=True)
    tokenized_val = val_dataset.map(tokenize_function, batched=True)

    # Define training arguments
    # training_args = TrainingArguments(
    #     output_dir=f"{base_output_dir}/fold_{fold_idx}",
    #     learning_rate=1e-5,
    #     per_device_train_batch_size=4,
    #     per_device_eval_batch_size=4,
    #     num_train_epochs=3,
    #     weight_decay=0.01,
    #     evaluation_strategy="steps",
    #     eval_steps=500,
    #     load_best_model_at_end=True,
    #     metric_for_best_model="f1",
    #     save_total_limit=1,
    #     report_to="wandb",  # Enable wandb logging
    #     run_name=f"fold_{fold_idx}"
    # )
    # Define training arguments with adjusted parameters
    training_args = TrainingArguments(
        output_dir=f"{base_output_dir}/{fold_idx}",
        learning_rate=2e-6,  # Reduced learning rate by 10x
        per_device_train_batch_size=2,
        per_device_eval_batch_size=2,
        gradient_accumulation_steps=4,  # Accumulate gradients over 4 steps (effective batch size = 8)
        num_train_epochs=8,
        weight_decay=0.01,
        evaluation_strategy="steps",
        eval_steps=50,
        # load_best_model_at_end=True,
        metric_for_best_model="f1",
        # save_total_limit=1,
        report_to="wandb",
        save_strategy="no",
        run_name=f"{fold_idx}",
        warmup_ratio=0.1,  # Changed from warmup_steps to warmup_ratio
        logging_steps=2,
        lr_scheduler_type="cosine",  # Added cosine learning rate scheduler
        # Multi-GPU settings
        # fp16=True,                    # Enable mixed precision training
        # dataloader_num_workers=4,     # Parallel data loading
        # ddp_find_unused_parameters=False,  # Optimize DDP
    )


    # Initialize trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_train,
        eval_dataset=tokenized_val,
        compute_metrics=compute_metrics,
        # callbacks=[EarlyStoppingCallback(early_stopping_patience=5)]
    )

    # Prepare everything with accelerator
    trainer = accelerator.prepare(trainer)

    # Train the model
    trainer.train()
    
    # Evaluate on validation set
    eval_results = trainer.evaluate()
    
    return eval_results

def cross_validate(n_folds: int = 5):
    # Initialize accelerator
    accelerator = Accelerator()
    
    # Load all datasets
    train_data = load_data("/fast/XXXX-3/forecasting/datasets/manifold/manifold_binary_train_with_r1queries.json")
    # test_data = load_data("/fast/XXXX-3/forecasting/datasets/menge/binary_test.json")
    test_data = load_data("/fast/XXXX-3/forecasting/datasets/manifold/manifold_binary_validation_set.json")
    
    # Convert to HuggingFace datasets
    train_dataset = Dataset.from_list(train_data)
    test_dataset = Dataset.from_list(test_data)
    
    
    # shuffle both datasets
    train_dataset = train_dataset.shuffle(seed=42)
    test_dataset = test_dataset.shuffle(seed=42)
    
    # Balance the training dataset
    num1s = np.sum(train_dataset['label'])
    num0s = len(train_dataset['label']) - num1s
    
    if num1s > num0s:
        # keep num0s number of 1s
        indices_1s = [i for i, label in enumerate(train_dataset['label']) if label == 1]
        indices_0s = [i for i, label in enumerate(train_dataset['label']) if label == 0]
        
        # Take only num0s samples from the 1s
        sampled_1s_indices = indices_1s[:num0s]
        only1s = train_dataset.select(sampled_1s_indices)
        only0s = train_dataset.select(indices_0s)
        train_dataset = concatenate_datasets([only1s, only0s])
    else:
        # only keep num1s number of 0s
        indices_1s = [i for i, label in enumerate(train_dataset['label']) if label == 1]
        indices_0s = [i for i, label in enumerate(train_dataset['label']) if label == 0]
        
        # Take only num1s samples from the 0s
        sampled_0s_indices = indices_0s[:num1s]
        only0s = train_dataset.select(sampled_0s_indices)
        only1s = train_dataset.select(indices_1s)
        train_dataset = concatenate_datasets([only0s, only1s])
    
    # Balance the test dataset
    test_num1s = np.sum(test_dataset['label'])
    test_num0s = len(test_dataset['label']) - test_num1s
    
    if test_num1s > test_num0s:
        # keep test_num0s number of 1s
        test_indices_1s = [i for i, label in enumerate(test_dataset['label']) if label == 1]
        test_indices_0s = [i for i, label in enumerate(test_dataset['label']) if label == 0]
        
        # Take only test_num0s samples from the 1s
        test_sampled_1s_indices = test_indices_1s[:test_num0s]
        test_only1s = test_dataset.select(test_sampled_1s_indices)
        test_only0s = test_dataset.select(test_indices_0s)
        test_dataset = concatenate_datasets([test_only1s, test_only0s])
    else:
        # only keep test_num1s number of 0s
        test_indices_1s = [i for i, label in enumerate(test_dataset['label']) if label == 1]
        test_indices_0s = [i for i, label in enumerate(test_dataset['label']) if label == 0]
        
        # Take only test_num1s samples from the 0s
        test_sampled_0s_indices = test_indices_0s[:test_num1s]
        test_only0s = test_dataset.select(test_sampled_0s_indices)
        test_only1s = test_dataset.select(test_indices_1s)
        test_dataset = concatenate_datasets([test_only0s, test_only1s])
    
    train_dataset = train_dataset.shuffle(seed=42)
    test_dataset = test_dataset.shuffle(seed=42)
    
    # Add this diagnostic print
    train_label_counts = np.bincount(train_dataset['label'])
    test_label_counts = np.bincount(test_dataset['label'])
    print(f"Train label counts: {train_label_counts}")
    print(f"Test label counts: {test_label_counts}")

    # Initialize k-fold
    kf = KFold(n_splits=n_folds, shuffle=True, random_state=42)
    
    # Initialize wandb only on the main process
    if accelerator.is_main_process:
        wandb.init(project="forecasting-classifier", name="manifold-without-validation")
    
    # Store results for each fold
    fold_results = []
    model_name = "microsoft/deberta-v3-large"
    
    # Combine train and validation datasets for cross-validation
    combined_train = concatenate_datasets([train_dataset])
    
    
    # Train final model on all training data and evaluate on validation set
    if accelerator.is_main_process:
        print("Training final model on all training data...")
    train_all = concatenate_datasets([train_dataset])
    final_results = train_fold(
        train_all,
        test_dataset,
        model_name,
        fold_idx="manifold",
        base_output_dir="./results"
    )
    
    # Log results only on the main process
    if accelerator.is_main_process:
        wandb.log({
            "final_accuracy": final_results['eval_accuracy'],
            "final_f1": final_results['eval_f1'],
            "final_precision": final_results['eval_precision'],
            "final_recall": final_results['eval_recall']
        })
        
        wandb.finish()
    

if __name__ == "__main__":
    # Enable deterministic behavior for reproducibility across GPUs
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    
    cross_validate(n_folds=5)
    # avg_metrics, final_results = cross_validate(n_folds=5)
    # print("\nAverage Cross-Validation Metrics:")
    # print(avg_metrics)
    # print("\nFinal Model Performance on Validation Set:")
    # print(final_results)