import torch
import numpy as np
from typing import Dict, Any, Optional
from torch.utils.data import DataLoader
from tqdm import tqdm
from vox_bench.utils import handle_keyboard_interrupt

class Learner:

    def __init__(
        self,
        model: torch.nn.Module,
        optimizer: torch.optim.Optimizer,
        scheduler,
        head,
        datamodule,
        device: torch.device,
        wandb_handler,
        checkpoint_name: str = "best_model",
    ):
        self.model = model.to(device)
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.head = head

        self.train_loader: Optional[DataLoader] = getattr(datamodule, "train_dl", None)
        self.val_loader: Optional[DataLoader] = getattr(datamodule, "val_dl", None)
        self.test_loader: Optional[DataLoader] = getattr(datamodule, "test_dl", None)

        self.device = device
        self.wandb = wandb_handler
        self.checkpoint_name = checkpoint_name

    def _move_to_device(self, data: Dict):
        x_dict = {k: v.to(self.device) for k, v in data.get("X", {}).items()}
        y_dict = {k: v.to(self.device) for k, v in data.get("Y", {}).items()}
        return x_dict, y_dict

    def _step(self, data: Dict, training: bool = True):
        x_dict, y_dict = self._move_to_device(data)

        outputs = self.head.forward(self.model, x_dict)
        loss = self.head.compute_loss(outputs, y_dict)

        if training:
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

        return loss.item(), outputs, y_dict

    def _run_eval(self, loader: DataLoader, prefix: str, epoch: int):

        self.model.eval()
        aggregated = {}
        total_loss = 0.0

        pbar = tqdm(
            loader,
            desc=f"Epoch {epoch:03d} [{prefix}]",
            leave=False,
        )

        with torch.no_grad():
            for data in pbar:
                loss, outputs, y_dict = self._step(data, training=False)

                batch_eval = self.head.evaluate_batch(outputs, y_dict)

                for key, value in batch_eval.items():
                    aggregated.setdefault(key, []).append(value)

                total_loss += loss
                pbar.set_postfix(loss=f"{loss:.4f}")

        avg_loss = total_loss / max(len(loader), 1)

        return aggregated, avg_loss
        
    @handle_keyboard_interrupt
    def fit(self, epochs: int):

        best_metric = -float("inf")
        checkpoint_filename = f"{self.checkpoint_name}.pth"

        for epoch in range(1, epochs + 1):

            train_metrics = {}

            if self.train_loader is not None:

                self.model.train()
                total_loss = 0.0

                pbar = tqdm(
                    self.train_loader,
                    desc=f"Epoch {epoch:03d} [train]",
                    leave=False,
                )

                for data in pbar:
                    loss, _, _ = self._step(data, training=True)
                    total_loss += loss
                    pbar.set_postfix(loss=f"{loss:.4f}")

                train_avg_loss = total_loss / max(len(self.train_loader), 1)
                train_metrics["train/loss"] = train_avg_loss

            val_metrics = {}

            if self.val_loader is not None:

                aggregated, val_avg_loss = self._run_eval(
                    self.val_loader,
                    prefix="val",
                    epoch=epoch,
                )

                val_metrics = self.head.compute_epoch_metrics(
                    aggregated=aggregated,
                    prefix="val",
                    epoch=epoch,
                    wandb=self.wandb,
                )

                val_metrics["val/loss"] = val_avg_loss

                monitor_value = self.head.get_monitor_value(val_metrics)

                if self.scheduler is not None:
                    self.scheduler.step(monitor_value)

                is_best = monitor_value > best_metric
                if is_best:
                    best_metric = monitor_value

                self.wandb.save_checkpoint(
                    model=self.model,
                    filename=checkpoint_filename,
                    epoch=epoch,
                    best_metric=best_metric,
                    metadata={"model_class": self.model.__class__.__name__},
                    is_best=is_best,
                    extra_metrics={"monitor": monitor_value},
                )
            else:
                is_best = False

            log_data = {
                "epoch": epoch,
                "lr": self.optimizer.param_groups[0]["lr"],
                **train_metrics,
                **val_metrics,
            }

            self.wandb.log_metrics(log_data)

            print(
                f"[Epoch {epoch:03d}/{epochs}] "
                f"LR: {self.optimizer.param_groups[0]['lr']:.2e} "
                + (f"| Train loss: {train_metrics.get('train/loss', 0):.4f} " if train_metrics else "")
                + (f"| Val loss: {val_metrics.get('val/loss', 0):.4f} " if val_metrics else "")
                + ("← new best" if is_best else "")
            )

        print("\n[Done] Training complete.")


    def test(self):

        if self.test_loader is None:
            print("No test loader available.")
            return {}

        aggregated, avg_loss = self._run_eval(
            self.test_loader,
            prefix="test",
            epoch=0,
        )

        test_metrics = self.head.compute_epoch_metrics(
            aggregated=aggregated,
            prefix="test",
            epoch=0,
            wandb=self.wandb,
        )

        test_metrics["test/loss"] = avg_loss
        
        self.wandb.log_metrics(test_metrics)
        
        # print("\n[Test Results]")
        # for key, value in test_metrics.items():
        #     print(f"{key}: {value:.4f}")

        return test_metrics

    def finish_run(self):
        self.wandb.finish_run()