import torch

# hyperparams
epochs = 0
finetune = 2


def train_model(model, data_loader, optimizer, criterion, device, epochs):

    model.train()

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

        total_loss = 0.0
        last_subepoch = -1

        for batch_idx, batch in enumerate(data_loader):

            x_batch, y_batch = batch

            x_batch = x_batch.to(device)
            y_batch = y_batch.to(device)

            optimizer.zero_grad()

            output = model(x_batch)

            output = output.view(-1, output.size(-1))
            y_batch = y_batch.view(-1)

            current_loss = criterion(output, y_batch)
            current_loss.backward()

            if current_loss.item() > 3:
                torch.nn.utils.clip_grad_norm_(
                    model.parameters(),
                    max_norm=0.5
                )
            else:
                torch.nn.utils.clip_grad_norm_(
                    model.parameters(),
                    max_norm=1.0
                )

            torch.cuda.empty_cache()

            optimizer.step()

            total_loss += current_loss.item()

            # Calculate progress through this epoch
            progress = (batch_idx + 1) / len(data_loader)

            # Only display every 0.05 epoch
            subepoch = int(progress / 0.05) * 0.05

            if subepoch > last_subepoch and subepoch < 1.0:
                print(
                    f"\rEpoch {epoch - 1 + subepoch:.2f}",
                    end="",
                    flush=True
                )
                last_subepoch = subepoch

        avg_loss = total_loss / len(data_loader)

        # Clear the subepoch progress line
        print("\r" + " " * 40 + "\r", end="")

        # Normal epoch result
        print(
            f"Epoch {epoch}/{epochs}, "
            f"Average Loss: {avg_loss:.4f}"
        )

    print("The weights have been updated.")