from .imports import *
from .data import preload_and_process_data, get_data_loader
from .model import GeneformerMultiTask
from .utils import calculate_task_specific_metrics
from torch.utils.tensorboard import SummaryWriter
import pandas as pd
import os
from tqdm import tqdm
import random
import numpy as np
import torch


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def initialize_wandb(config):
    if config.get("use_wandb", False):
        import wandb
        wandb.init(project=config["wandb_project"], config=config)
        print("Weights & Biases (wandb) initialized and will be used for logging.")
    else:
        print("Weights & Biases (wandb) is not enabled. Logging will use other methods.")

def create_model(config, num_labels_list, device):
    model = GeneformerMultiTask(
        config["pretrained_path"],
        num_labels_list,
        dropout_rate=config["dropout_rate"],
        use_task_weights=config["use_task_weights"],
        task_weights=config["task_weights"],
        max_layers_to_freeze=config["max_layers_to_freeze"],
        use_attention_pooling=config["use_attention_pooling"]
    )
    if config["use_data_parallel"]:
        model = nn.DataParallel(model)
    return model.to(device)

def setup_optimizer_and_scheduler(model, config, total_steps):
    optimizer = AdamW(model.parameters(), lr=config["learning_rate"], weight_decay=config["weight_decay"])
    warmup_steps = int(config["warmup_ratio"] * total_steps)

    if config["lr_scheduler_type"] == "linear":
        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps)
    elif config["lr_scheduler_type"] == "cosine":
        scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, num_cycles=0.5)

    return optimizer, scheduler

def train_epoch(model, train_loader, optimizer, scheduler, device, config, writer, epoch):
    model.train()
    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{config['epochs']}")
    for batch_idx, batch in enumerate(progress_bar):
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = [batch['labels'][task_name].to(device) for task_name in config["task_names"]]

        loss, _, _ = model(input_ids, attention_mask, labels)
        loss.backward()

        if config["gradient_clipping"]:
            torch.nn.utils.clip_grad_norm_(model.parameters(), config["max_grad_norm"])

        optimizer.step()
        scheduler.step()

        writer.add_scalar('Training Loss', loss.item(), epoch * len(train_loader) + batch_idx)
        if config.get("use_wandb", False):
            wandb.log({'Training Loss': loss.item()})

        # Update progress bar
        progress_bar.set_postfix({'loss': f"{loss.item():.4f}"})

    return loss.item()  # Return the last batch loss

def validate_model(model, val_loader, device, config):
    model.eval()
    val_loss = 0.0
    task_true_labels = {task_name: [] for task_name in config["task_names"]}
    task_pred_labels = {task_name: [] for task_name in config["task_names"]}
    task_pred_probs = {task_name: [] for task_name in config["task_names"]}

    with torch.no_grad():
        for batch in val_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = [batch['labels'][task_name].to(device) for task_name in config["task_names"]]
            loss, logits, _ = model(input_ids, attention_mask, labels)
            val_loss += loss.item()

            for sample_idx in range(len(batch['input_ids'])):
                for i, task_name in enumerate(config["task_names"]):
                    true_label = batch['labels'][task_name][sample_idx].item()
                    pred_label = torch.argmax(logits[i][sample_idx], dim=-1).item()
                    pred_prob = torch.softmax(logits[i][sample_idx], dim=-1).cpu().numpy()
                    task_true_labels[task_name].append(true_label)
                    task_pred_labels[task_name].append(pred_label)
                    task_pred_probs[task_name].append(pred_prob)

    val_loss /= len(val_loader)
    return val_loss, task_true_labels, task_pred_labels, task_pred_probs

def log_metrics(task_metrics, val_loss, config, writer, epochs):
    for task_name, metrics in task_metrics.items():
        print(f"{task_name} - Validation F1 Macro: {metrics['f1']:.4f}, Validation Accuracy: {metrics['accuracy']:.4f}")
        if config.get("use_wandb", False):
            import wandb
            wandb.log({
                f'{task_name} Validation F1 Macro': metrics['f1'],
                f'{task_name} Validation Accuracy': metrics['accuracy']
            })

    writer.add_scalar('Validation Loss', val_loss, epochs)
    for task_name, metrics in task_metrics.items():
        writer.add_scalar(f'{task_name} - Validation F1 Macro', metrics['f1'], epochs)
        writer.add_scalar(f'{task_name} - Validation Accuracy', metrics['accuracy'], epochs)

def save_validation_predictions(val_cell_id_mapping, task_true_labels, task_pred_labels, task_pred_probs, config, trial_number=None):
    if trial_number is not None:
        trial_results_dir = os.path.join(config["results_dir"], f"trial_{trial_number}")
        os.makedirs(trial_results_dir, exist_ok=True)
        val_preds_file = os.path.join(trial_results_dir, "val_preds.csv")
    else:
        val_preds_file = os.path.join(config["results_dir"], "manual_run_val_preds.csv")
    
    rows = []
    for sample_idx in range(len(val_cell_id_mapping)):
        row = {'Cell ID': val_cell_id_mapping[sample_idx]}
        for task_name in config["task_names"]:
            row[f'{task_name} True'] = task_true_labels[task_name][sample_idx]
            row[f'{task_name} Pred'] = task_pred_labels[task_name][sample_idx]
            row[f'{task_name} Probabilities'] = ','.join(map(str, task_pred_probs[task_name][sample_idx]))
        rows.append(row)
    
    df = pd.DataFrame(rows)
    df.to_csv(val_preds_file, index=False)
    print(f"Validation predictions saved to {val_preds_file}")


def train_model(config, device, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list):
    set_seed(config["seed"])
    initialize_wandb(config)

    model = create_model(config, num_labels_list, device)
    total_steps = len(train_loader) * config["epochs"]
    optimizer, scheduler = setup_optimizer_and_scheduler(model, config, total_steps)

    log_dir = os.path.join(config["tensorboard_log_dir"], "manual_run")
    writer = SummaryWriter(log_dir=log_dir)

    epoch_progress = tqdm(range(config["epochs"]), desc="Training Progress")
    for epoch in epoch_progress:
        last_loss = train_epoch(model, train_loader, optimizer, scheduler, device, config, writer, epoch)
        epoch_progress.set_postfix({'last_loss': f"{last_loss:.4f}"})

    val_loss, task_true_labels, task_pred_labels, task_pred_probs = validate_model(model, val_loader, device, config)
    task_metrics = calculate_task_specific_metrics(task_true_labels, task_pred_labels)
    
    log_metrics(task_metrics, val_loss, config, writer, config["epochs"])
    writer.close()

    save_validation_predictions(val_cell_id_mapping, task_true_labels, task_pred_labels, task_pred_probs, config)

    if config.get("use_wandb", False):
        import wandb
        wandb.finish()

    print(f"\nFinal Validation Loss: {val_loss:.4f}")
    return val_loss, model  # Return both the validation loss and the trained model

def objective(trial, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list, config, device):
    set_seed(config["seed"])  # Set the seed before each trial
    initialize_wandb(config)

    # Hyperparameters
    config["learning_rate"] = trial.suggest_float("learning_rate", config["hyperparameters"]["learning_rate"]["low"], config["hyperparameters"]["learning_rate"]["high"], log=config["hyperparameters"]["learning_rate"]["log"])
    config["warmup_ratio"] = trial.suggest_float("warmup_ratio", config["hyperparameters"]["warmup_ratio"]["low"], config["hyperparameters"]["warmup_ratio"]["high"])
    config["weight_decay"] = trial.suggest_float("weight_decay", config["hyperparameters"]["weight_decay"]["low"], config["hyperparameters"]["weight_decay"]["high"])
    config["dropout_rate"] = trial.suggest_float("dropout_rate", config["hyperparameters"]["dropout_rate"]["low"], config["hyperparameters"]["dropout_rate"]["high"])
    config["lr_scheduler_type"] = trial.suggest_categorical("lr_scheduler_type", config["hyperparameters"]["lr_scheduler_type"]["choices"])
    config["use_attention_pooling"] = trial.suggest_categorical("use_attention_pooling", [True, False])

    if config["use_task_weights"]:
        config["task_weights"] = [trial.suggest_float(f"task_weight_{i}", config["hyperparameters"]["task_weights"]["low"], config["hyperparameters"]["task_weights"]["high"]) for i in range(len(num_labels_list))]
        weight_sum = sum(config["task_weights"])
        config["task_weights"] = [weight / weight_sum for weight in config["task_weights"]]
    else:
        config["task_weights"] = None

    # Fix for max_layers_to_freeze
    if isinstance(config["max_layers_to_freeze"], dict):
        config["max_layers_to_freeze"] = trial.suggest_int("max_layers_to_freeze", config["max_layers_to_freeze"]["min"], config["max_layers_to_freeze"]["max"])
    elif isinstance(config["max_layers_to_freeze"], int):
        # If it's already an int, we don't need to suggest it
        pass
    else:
        raise ValueError("Invalid type for max_layers_to_freeze. Expected dict or int.")

    model = create_model(config, num_labels_list, device)
    total_steps = len(train_loader) * config["epochs"]
    optimizer, scheduler = setup_optimizer_and_scheduler(model, config, total_steps)

    log_dir = os.path.join(config["tensorboard_log_dir"], f"trial_{trial.number}")
    writer = SummaryWriter(log_dir=log_dir)

    for epoch in range(config["epochs"]):
        train_epoch(model, train_loader, optimizer, scheduler, device, config, writer, epoch)

    val_loss, task_true_labels, task_pred_labels, task_pred_probs = validate_model(model, val_loader, device, config)
    task_metrics = calculate_task_specific_metrics(task_true_labels, task_pred_labels)
    
    log_metrics(task_metrics, val_loss, config, writer, config["epochs"])
    writer.close()

    save_validation_predictions(val_cell_id_mapping, task_true_labels, task_pred_labels, task_pred_probs, config, trial.number)

    trial.set_user_attr("model_state_dict", model.state_dict())
    trial.set_user_attr("task_weights", config["task_weights"])

    trial.report(val_loss, config["epochs"])

    if trial.should_prune():
        raise optuna.TrialPruned()

    if config.get("use_wandb", False):
        import wandb
        wandb.log({
            "trial_number": trial.number,
            "val_loss": val_loss,
            **{f"{task_name}_f1": metrics['f1'] for task_name, metrics in task_metrics.items()},
            **{f"{task_name}_accuracy": metrics['accuracy'] for task_name, metrics in task_metrics.items()},
            **{k: v for k, v in config.items() if k in ["learning_rate", "warmup_ratio", "weight_decay", "dropout_rate", "lr_scheduler_type", "use_attention_pooling", "max_layers_to_freeze"]}
        })
        wandb.finish()

    return val_loss