import logging
import time

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import roc_auc_score

from base.base_dataset import BaseADDataset
from base.base_net import BaseNet
from base.base_trainer import BaseTrainer


class AETrainer(BaseTrainer):
    def __init__(
        self,
        optimizer_name: str = "adam",
        lr: float = 0.001,
        n_epochs: int = 150,
        lr_milestones: tuple = (),
        batch_size: int = 128,
        weight_decay: float = 1e-6,
        device: str = "cuda",
        n_jobs_dataloader: int = 0,
    ):
        super().__init__(
            optimizer_name,
            lr,
            n_epochs,
            lr_milestones,
            batch_size,
            weight_decay,
            device,
            n_jobs_dataloader,
        )

        # Results
        self.train_time = None
        self.test_auc = None
        self.test_time = None

    def train(self, dataset: BaseADDataset, ae_net: BaseNet, k_fold_idx: int = None):
        logger = logging.getLogger()

        # Get train data loader
        if k_fold_idx is not None:
            train_loader, _ = dataset.loaders_k_fold(
                fold_idx=k_fold_idx,
                batch_size=self.batch_size,
                num_workers=self.n_jobs_dataloader,
            )
        else:
            train_loader, _, _ = dataset.loaders(
                batch_size=self.batch_size, num_workers=self.n_jobs_dataloader
            )

        # Set loss
        criterion = nn.MSELoss(reduction="none")

        # Set device
        ae_net = ae_net.to(self.device)
        criterion = criterion.to(self.device)

        # Set optimizer (Adam optimizer for now)
        optimizer = optim.Adam(
            ae_net.parameters(), lr=self.lr, weight_decay=self.weight_decay
        )

        # Set learning rate scheduler
        scheduler = optim.lr_scheduler.MultiStepLR(
            optimizer, milestones=self.lr_milestones, gamma=0.1
        )

        # Prepare containers for results (similar to test)
        all_indices = []
        all_labels_exp_based = []
        all_labels_manual_based = []
        all_semi_targets = []
        all_file_ids = []
        all_frame_ids = []
        all_scores = []

        # Training
        logger.info("Starting pretraining...")
        start_time = time.time()
        ae_net.train()
        ae_net.summary(receptive_field=True)  # Add network summary before training

        for epoch in range(self.n_epochs):
            epoch_loss = 0.0
            n_batches = 0
            epoch_start_time = time.time()
            for data in train_loader:
                (
                    inputs,
                    labels_exp_based,
                    labels_manual_based,
                    semi_targets,
                    idx,
                    (file_id, frame_id),
                ) = data
                inputs, idx = (
                    inputs.to(self.device),
                    idx.to(self.device),
                )

                # Zero the network parameter gradients
                optimizer.zero_grad()

                # Forward + backward + optimize
                rec = ae_net(inputs)
                rec_loss = criterion(rec, inputs)
                scores = torch.mean(rec_loss, dim=tuple(range(1, rec.dim())))
                loss = torch.mean(rec_loss)
                loss.backward()
                optimizer.step()

                epoch_loss += loss.item()
                n_batches += 1

                # Save all relevant information from dataloader
                all_indices.extend(idx.detach().cpu().numpy())
                all_labels_exp_based.extend(labels_exp_based.detach().cpu().numpy())
                all_labels_manual_based.extend(
                    labels_manual_based.detach().cpu().numpy()
                )
                all_semi_targets.extend(semi_targets.detach().cpu().numpy())
                all_file_ids.extend(file_id.detach().cpu().numpy())
                all_frame_ids.extend(frame_id.detach().cpu().numpy())
                all_scores.extend(scores.detach().cpu().numpy())

            scheduler.step()
            if epoch in self.lr_milestones:
                logger.info(
                    "  LR scheduler: new learning rate is %g"
                    % float(scheduler.get_last_lr()[0])
                )

            # log epoch statistics
            epoch_train_time = time.time() - epoch_start_time
            logger.info(
                f"| Epoch: {epoch + 1:03}/{self.n_epochs:03} | Train Time: {epoch_train_time:.3f}s "
                f"| Train Loss: {epoch_loss / n_batches:.6f} |"
            )

        self.train_time = time.time() - start_time

        # Save all results as member variables (like in test)
        self.train_indices = np.array(all_indices)
        self.train_labels_exp_based = np.array(all_labels_exp_based)
        self.train_labels_manual_based = np.array(all_labels_manual_based)
        self.train_semi_targets = np.array(all_semi_targets)
        self.train_file_ids = np.array(all_file_ids)
        self.train_frame_ids = np.array(all_frame_ids)
        self.train_scores = np.array(all_scores)
        self.train_loss = epoch_loss / n_batches if n_batches > 0 else float("nan")

        logger.info("Pretraining Time: {:.3f}s".format(self.train_time))
        logger.info("Finished pretraining.")

        return ae_net

    def test(self, dataset: BaseADDataset, ae_net: BaseNet, k_fold_idx: int = None):
        logger = logging.getLogger()

        # Get test data loader
        if k_fold_idx is not None:
            _, test_loader = dataset.loaders_k_fold(
                fold_idx=k_fold_idx,
                batch_size=self.batch_size,
                num_workers=self.n_jobs_dataloader,
            )
        else:
            _, test_loader, _ = dataset.loaders(
                batch_size=self.batch_size, num_workers=self.n_jobs_dataloader
            )

        # Set loss
        criterion = nn.MSELoss(reduction="none")

        # Set device for network
        ae_net = ae_net.to(self.device)
        criterion = criterion.to(self.device)

        # Prepare containers for results
        all_indices = []
        all_labels_exp_based = []
        all_labels_manual_based = []
        all_semi_targets = []
        all_file_ids = []
        all_frame_ids = []
        all_scores = []

        # Testing
        logger.info("Testing autoencoder...")
        epoch_loss = 0.0
        n_batches = 0
        start_time = time.time()
        ae_net.eval()
        ae_net.summary(receptive_field=True)  # Add network summary before testing

        with torch.no_grad():
            for data in test_loader:
                (
                    inputs,
                    labels_exp_based,
                    labels_manual_based,
                    semi_targets,
                    idx,
                    (file_id, frame_id),
                ) = data
                inputs, idx = (
                    inputs.to(self.device),
                    idx.to(self.device),
                )

                rec = ae_net(inputs)
                rec_loss = criterion(rec, inputs)
                scores = torch.mean(rec_loss, dim=tuple(range(1, rec.dim())))

                loss = torch.mean(rec_loss)
                epoch_loss += loss.item()
                n_batches += 1

                # Save all relevant information from dataloader
                all_indices.extend(idx.detach().cpu().numpy())
                all_labels_exp_based.extend(labels_exp_based.detach().cpu().numpy())
                all_labels_manual_based.extend(
                    labels_manual_based.detach().cpu().numpy()
                )
                all_semi_targets.extend(semi_targets.detach().cpu().numpy())
                all_file_ids.extend(file_id.detach().cpu().numpy())
                all_frame_ids.extend(frame_id.detach().cpu().numpy())
                all_scores.extend(scores.detach().cpu().numpy())

        self.test_time = time.time() - start_time

        # Save all results as member variables
        self.test_indices = np.array(all_indices)
        self.test_labels_exp_based = np.array(all_labels_exp_based)
        self.test_labels_manual_based = np.array(all_labels_manual_based)
        self.test_semi_targets = np.array(all_semi_targets)
        self.test_file_ids = np.array(all_file_ids)
        self.test_frame_ids = np.array(all_frame_ids)
        self.test_scores = np.array(all_scores)

        # No performance metric is calculated using labels, as this is pre-training
        self.test_loss = epoch_loss / n_batches if n_batches > 0 else float("nan")

        # Log results
        logger.info("Test Loss: {:.6f}".format(self.test_loss))
        logger.info("Test Time: {:.3f}s".format(self.test_time))
        logger.info("Finished testing autoencoder.")
