import torch
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_
from just_time_windows.Actor.actor import Actor


def train_batch(actor, baseline, batch, optimizer, gradient_clipping=True, comparison_model=None, compute_cost_ratio=True):
    device = actor.device

    actor.train_mode()
    actor.train()
    actor_output = actor(batch)
    actor_cost, log_probs = actor_output['total_time'], actor_output['log_probs']

    with torch.no_grad():
        baseline.greedy_search()
        baseline_output = baseline(batch)
        baseline_cost = baseline_output['total_time']

    loss = ((actor_cost - baseline_cost).detach() * log_probs).mean()

    optimizer.zero_grad()
    loss.backward()

    if gradient_clipping:
        for group in optimizer.param_groups:
            params = [p for p in group['params'] if p.grad is not None]
            if params:
                clip_grad_norm_(params, max_norm=1, norm_type=2)

    optimizer.step()

    if compute_cost_ratio:
        if comparison_model is None:
            normalize = actor.apply_normalization
            comparison_model = Actor(model=None, num_neighbors_action=1, normalize=normalize, device=device)

        with torch.no_grad():
            comp_output = comparison_model(batch)
            comp_cost = comp_output['total_time']

        a = comp_cost.sum().item()
        b = actor_cost.sum().item()
        return b / a

    return None