"""
Implements the Trainer class to handle training loops, evaluation, and model
checkpointing for classification tasks. Supports mixed-precision, TorchScript
compilation, and periodic metrics logging.
"""

import os
import time
from datetime import datetime
from tqdm import tqdm
import torch
from utils import model_size


class Trainer:
    def __init__(
        self,
        model,
        train_loader,
        test_loader,
        optimizer,
        scheduler,
        loss_fn,
        epochs,
        freq_print,
        device,
        state_dict_name,
        mixed_precision,
        torch_dot_compile,
        before_training,
    ):
        """
        Trainer for classification models with logging and checkpointing.
        """
        self.model = model
        self.train_loader = train_loader
        self.test_loader = test_loader
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.epochs = epochs
        self.freq_print = freq_print
        self.device = device
        self.state_dict_name = state_dict_name
        self.mixed_precision = mixed_precision
        self.before_training = before_training

        if loss_fn.lower() == "ce":
            self.loss_fn = torch.nn.CrossEntropyLoss()
        self.accuracy_fn = (
            lambda logits, labels: (logits.argmax(dim=1) == labels).float().mean()
        )

        self.step_count = 0

        if torch_dot_compile:
            # Compile the model in C++
            self.model = torch.compile(self.model, mode="default", fullgraph=False)

        if mixed_precision:
            self.scaler = torch.cuda.amp.GradScaler()

    def train(self):
        """
        Run the full training loop, including optional pre-training evaluation,
        per-epoch training, validation, logging, and model checkpointing.
        """
        self.train_losses, self.test_metrics = [], []

        if self.before_training:
            # Initial evaluation before training
            initial_train = self._evaluate_metrics(self.train_loader)
            initial_test = self._evaluate_metrics(self.test_loader)

            n_params, n_bytes = model_size(self.model)
            print(
                f"Before Training | "
                f"Train Loss: {initial_train['loss']:.4f} | "
                f"Train Acc:  {initial_train['acc']*100:.2f}% | "
                f"Test Loss:  {initial_test['loss']:.4f} | "
                f"Test Acc:   {initial_test['acc']*100:.2f}% | "
                f"Params: {n_params} | Size: {n_bytes/1e9:.3f} GB"
            )
        # Main training loop
        for epoch in range(1, self.epochs + 1):
            start_time = time.time()

            train_loss, train_acc = self._train_one_epoch()

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

            test_metrics = self._evaluate_metrics(self.test_loader)
            current_lr = self.optimizer.param_groups[0]["lr"]
            epoch_time = time.time() - start_time

            # GPU memory usage
            free, total = torch.cuda.mem_get_info(self.device)
            mem_used = (total - free) / 1024**3

            # Print progress
            if self.freq_print is not None:
                print(
                    f"Epoch {epoch}/{self.epochs} | "
                    f"Train Loss: {train_loss:.4f} | "
                    f"Train Acc:  {train_acc*100:.2f}% | "
                    f"Test Loss:  {test_metrics['loss']:.4f} | "
                    f"Test Acc:   {test_metrics['acc']*100:.2f}% | "
                    f"GPU: {mem_used:.2f} GB | "
                    f"Time: {epoch_time:.1f}s"
                )

        # Save final state_dict
        if self.state_dict_name is not None:
            state_dict_path = f"model_weights/{self.state_dict_name}_time_{datetime.now():%d_%m_%Y_%H_%M}"
            os.makedirs(os.path.dirname(state_dict_path), exist_ok=True)
            torch.save(self.model.state_dict(), state_dict_path)

            print(f"Saved model state_dict at {state_dict_path}")

    def _train_one_epoch(self):
        """
        Train the model for one epoch over the training DataLoader.
        """
        self.model.train()
        running_loss = 0.0
        running_corrects = 0
        total_samples = 0

        for x, y in tqdm(self.train_loader, leave=False):
            x, y = (
                x.to(self.device, non_blocking=True),
                y.to(self.device, non_blocking=True).long(),
            )
            self.optimizer.zero_grad()

            with torch.amp.autocast(device_type="cuda", enabled=self.mixed_precision):
                logits = self.model(x)
                loss = self.loss_fn(logits, y)

            if self.mixed_precision:
                self.scaler.scale(loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                loss.backward()
                self.optimizer.step()

            running_loss += loss.item() * x.size(0)
            preds = logits.argmax(dim=1)
            running_corrects += (preds == y).sum().item()
            total_samples += x.size(0)
            self.step_count += 1

        epoch_loss = running_loss / total_samples
        epoch_acc = running_corrects / total_samples
        return epoch_loss, epoch_acc

    def _evaluate(self, data_loader):
        """
        Compute average loss over a DataLoader.
        """
        self.model.eval()
        total_loss = 0.0

        with torch.no_grad():
            for x, y in data_loader:
                x, y = (
                    x.to(self.device, non_blocking=True),
                    y.to(self.device, non_blocking=True).long(),
                )
                logits = self.model(x)
                loss = self.loss_fn(logits, y)
                total_loss += loss.item()

        return total_loss / len(data_loader)

    def _evaluate_metrics(self, data_loader):
        """
        Compute average loss and accuracy over a DataLoader.
        """
        self.model.eval()
        tot_loss, tot_acc = 0.0, 0.0

        with torch.no_grad():
            for x, y in tqdm(data_loader, leave=False):
                x, y = (
                    x.to(self.device, non_blocking=True),
                    y.to(self.device, non_blocking=True).long(),
                )
                logits = self.model(x)
                loss = self.loss_fn(logits, y)
                acc = (logits.argmax(dim=1) == y).float().mean()

                tot_loss += loss.item()
                tot_acc += acc.item()

        n = len(data_loader)
        return {"loss": tot_loss / n, "acc": tot_acc / n}
