"""
Implements the Trainer class for regression-based neural operators.
Includes mixed-precision training, periodic plotting, and dual-metric evaluation
using MSE and Relative MSE. Also provides RelativeMSELoss for normalized error.

Classes:
    RelativeMSELoss: Computes normalized MSE relative to target magnitude.
    Trainer: Executes training loops, evaluation, plotting, and checkpointing for regression.
"""

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


class RelativeMSELoss(torch.nn.Module):
    """
    Loss function that computes mean squared error normalized by target magnitude.
    """

    def __init__(self, eps=1e-8):
        super().__init__()
        self.eps = eps

    def forward(self, pred, target):
        numerator = torch.sum((pred - target) ** 2)
        denominator = torch.sum(target**2) + self.eps
        return numerator / denominator


class Trainer:
    def __init__(
        self,
        model,
        train_loader,
        test_loader,
        optimizer,
        scheduler,
        loss_fn,
        epochs,
        freq_print,
        freq_plot,
        plot_fn,
        plot_path,
        show,
        device,
        state_dict_name,
        mixed_precision,
        torch_dot_compile,
    ):
        """
        Regression Trainer for neural operators with plotting and checkpointing.
        """
        self.train_loader = train_loader
        self.test_loader = test_loader
        self.loss_fn = loss_fn
        self.model = model
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.epochs = epochs
        self.freq_print = freq_print
        self.freq_plot = freq_plot
        self.plot_fn = plot_fn
        self.plot_path = plot_path
        self.show = show
        self.device = device
        self.state_dict_name = state_dict_name
        self.mixed_precision = mixed_precision
        self.torch_dot_compile = torch_dot_compile

        # Losses
        self.MSE_loss_fn = torch.nn.MSELoss()
        self.rel_loss_fn = RelativeMSELoss()

        self.step_count = 0  # step_count counts the number of optimization steps, it's different from number of epochs, more specifically step_count = epoch * num_batches_per_epoch

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

        # Cache a fixed test sample for consistent plotting
        self.eval_input = None
        self.eval_target = None
        self._cache_plot_test_point()

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

    def train(self):
        """
        Execute the training loop with periodic plotting, evaluation, and checkpointing.
        Prints initial metrics, per-epoch metrics, and saves model state if specified.
        """
        self.train_losses, self.test_losses = [], []
        # Evaluate before the first training epoch
        initial_train_loss = self._evaluate(self.train_loader)
        initial_test_loss = self._evaluate_metrics(self.test_loader)

        n_params, n_bytes = model_size(self.model)
        print(
            f"Before Training | Train {self.loss_fn.__class__.__name__}: {initial_train_loss:.6f} | Test MSE: {initial_test_loss['MSE']:.6f} | Test rel_MSE: {initial_test_loss['rel_MSE']:.6f} | Number of parameters: {n_params} | Model size (approx): {n_bytes / 1e9:.3f} GB"
        )

        if self.freq_plot is not None:
            self.plot_fn(
                self.model,
                self.plot_input,
                self.plot_target,
                self.plot_path,
                0,
                self.freq_plot,
                self.show,
            )

        for epoch in range(1, self.epochs + 1):
            start_time = time.time()

            train_loss = self._train_one_epoch()
            self.scheduler.step() if self.scheduler is not None else None
            test_loss = self._evaluate_metrics(self.test_loader)

            current_lr = self.optimizer.param_groups[0]["lr"]
            end_time = time.time()

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

            if self.freq_print is not None:
                print(
                    f"Epoch: {epoch}/{self.epochs} | Train Loss: {train_loss:.6f} | Test MSE: {test_loss['MSE']:.6f} | Test rel_MSE: {test_loss['rel_MSE']:.6f} | GPU Usage: {mem_used:.2f}GB | Epoch duration: {end_time - start_time:.1f}"
                )

        if self.state_dict_name is not None:
            # Save state_dict
            state_dict_path = f"model_weights/{self.state_dict_name}_time_{datetime.now():%d_%m_%Y_%H_%M}"

            # Ensure the directory exists
            os.makedirs(os.path.dirname(state_dict_path), exist_ok=True)

            torch.save(self.model.state_dict(), state_dict_path)

            print(f"save state_dict of the model at {state_dict_path}")

    def _train_one_epoch(self):
        """
        Perform optimization over the training set for one epoch.
        """
        self.model.train()
        train_loss = 0.0
        for x, y in tqdm(self.train_loader, desc=None, leave=False):
            x, y = x.to(self.device), y.to(self.device)
            self.optimizer.zero_grad()

            # Single unified autocast context
            with torch.amp.autocast(device_type="cuda", enabled=self.mixed_precision):
                preds = self.model(x).squeeze(-1)
                loss = self.loss_fn(preds, y)

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

            train_loss += loss.item()

            # Plot prediction
            if self.freq_plot is not None and self.step_count % self.freq_plot == 0:
                self.plot_fn(
                    self.model,
                    self.plot_input,
                    self.plot_target,
                    self.plot_path,
                    self.step_count,
                    self.freq_plot,
                    self.show,
                )
            self.step_count += 1

        return train_loss / len(self.train_loader)

    def _evaluate(self, data_loader):
        """
        Compute average primary loss over a dataset.
        """
        self.model.eval()
        total_loss = 0.0
        with torch.no_grad():
            for x, y in data_loader:
                x, y = x.to(self.device), y.to(self.device)
                output_pred = self.model(x).squeeze(-1)
                loss = self.loss_fn(output_pred, y)
                total_loss += loss.item()
        return total_loss / len(data_loader)

    def _evaluate_metrics(self, data_loader):
        """
        Compute MSE and Relative MSE over a dataset.
        """
        self.model.eval()
        tot_MSE, tot_rel = 0.0, 0.0
        with torch.no_grad():
            for x, y in data_loader:
                x, y = x.to(self.device), y.to(self.device)
                preds = self.model(x).squeeze(-1)

                tot_MSE += self.MSE_loss_fn(preds, y).item()
                tot_rel += self.rel_loss_fn(preds, y).item()

        n = len(data_loader)
        return {"MSE": tot_MSE / n, "rel_MSE": tot_rel / n}

    def _cache_plot_test_point(self):
        """
        Cache a single sample from the test set for consistent plotting during training.
        """
        self.model.eval()
        with torch.no_grad():
            for x, y in self.test_loader:
                self.plot_input = x[0].to(self.device)
                self.plot_target = y[0]
                break
