from tqdm import tqdm
import pandas as pd
import torch

from attention_dynamic_model import set_decode_type
from reinforce_baseline import validate

from utils import generate_data_onfly, get_cur_time
from time import gmtime, strftime
from utils import FastTensorDataLoader
from uuid import uuid4
import os


class IterativeMean:
    def __init__(self):
        self.sum = 0
        self.n = 0

    def update_state(self, val):
        self.sum += val
        self.n += 1

    def result(self):
        return self.sum / self.n


def train_model(
    optimizer,
    model_torch,
    baseline,
    validation_dataset,
    samples=1280000,
    batch=128,
    val_batch_size=1000,
    start_epoch=0,
    end_epoch=5,
    from_checkpoint=False,
    grad_norm_clipping=1.0,
    batch_verbose=1000,
    graph_size=20,
    mod_dir=None,
    mem_efficient=False,
    log_results=False,
):
    if mod_dir is None:
        dt = strftime("%Y-%m-%d", gmtime())
        mod_dir = f"./model_{graph_size}_{dt}_{uuid4()[:4]}"
        os.makedirs(mod_dir)

    def rein_loss(model, inputs, baseline, num_batch):
        """Calculate loss for REINFORCE algorithm"""

        # Evaluate model, get costs and log probabilities
        cost, log_likelihood = model(inputs)

        # Evaluate baseline
        # For first wp_n_epochs we take the combination of baseline and ema for previous batches
        # after that we take a slice of precomputed baseline values
        bl_val = (
            bl_vals[num_batch] if bl_vals is not None else baseline.eval(inputs, cost)
        )

        # Calculate loss
        reinforce_loss = torch.mean((cost - bl_val.detach()) * log_likelihood)

        return reinforce_loss, torch.mean(cost)

    def mem_efficient_rein_loss(model, inputs, baseline, num_batch):
        rein_detached_loss, (cost, log_likelihood) = model.fwd_rein_loss(
            inputs, baseline, bl_vals, num_batch
        )
        return rein_detached_loss, torch.mean(cost)

    def grad(model, inputs, baseline, num_batch):
        """Calculate gradients"""
        if mem_efficient:
            loss, cost = mem_efficient_rein_loss(model, inputs, baseline, num_batch)
        else:
            loss, cost = rein_loss(model, inputs, baseline, num_batch)
            loss.backward()
        grads = [param.grad.view(-1) for param in model.parameters()]
        grads = torch.cat(grads)
        # we can return detached loss since it's backwarded already above
        return loss.detach(), cost, grads

    # For plotting
    train_loss_results = []
    train_cost_results = []
    val_cost_avg = []

    min_val_cost_avg = 1e9
    # Training loop
    for epoch in range(start_epoch, end_epoch):

        # Create dataset on current epoch
        data = generate_data_onfly(num_samples=samples, graph_size=graph_size)

        epoch_loss_avg = IterativeMean()
        epoch_cost_avg = IterativeMean()

        # Skip warm-up stage when we continue training from checkpoint
        if from_checkpoint and baseline.alpha != 1.0:
            print("Skipping warm-up mode")
            baseline.alpha = 1.0

        # If epoch > wp_n_epochs then precompute baseline values for the whole dataset else None
        bl_vals = baseline.eval_all(data)  # (samples, ) or None
        bl_vals = (
            torch.reshape(bl_vals, (-1, batch)) if bl_vals is not None else None
        )  # (n_batches, batch) or None

        print("Current decode type: {}".format(model_torch.decode_type))

        train_batches = FastTensorDataLoader(
            data[0], data[1], data[2], batch_size=batch, shuffle=False
        )

        for num_batch, x_batch in tqdm(
            enumerate(train_batches), desc="batch calculation at epoch {}".format(epoch)
        ):
            optimizer.zero_grad()
            # Optimize the model
            loss_value, cost_val, grads = grad(
                model_torch, x_batch, baseline, num_batch
            )

            # Clip gradients by grad_norm_clipping
            init_global_norm = torch.linalg.norm(grads)
            torch.nn.utils.clip_grad_norm_(model_torch.parameters(), grad_norm_clipping)
            grads = [param.grad.view(-1) for param in model_torch.parameters()]
            grads = torch.cat(grads)
            global_norm = torch.linalg.norm(grads)

            if num_batch % batch_verbose == 0:
                print(
                    "grad_global_norm = {}, clipped_norm = {}".format(
                        init_global_norm.cpu().numpy(), global_norm.cpu().numpy()
                    )
                )

            optimizer.step()

            # Track progress
            epoch_loss_avg.update_state(loss_value)
            epoch_cost_avg.update_state(cost_val)

            if num_batch % batch_verbose == 0:
                print(
                    "Epoch {} (batch = {}): Loss: {}: Cost: {}".format(
                        epoch,
                        num_batch,
                        epoch_loss_avg.result(),
                        epoch_cost_avg.result(),
                    )
                )

        # Update baseline if the candidate model is good enough. In this case also create new baseline dataset
        baseline.epoch_callback(model_torch, epoch)
        set_decode_type(model_torch, "sampling")

        # Validate current model
        val_cost = validate(validation_dataset, model_torch, val_batch_size)
        val_cost_avg.append(val_cost)

        if val_cost < min_val_cost_avg:
            torch.save(model_torch.state_dict(), f"{mod_dir}/adm.pt")
            min_val_cost_avg = val_cost_avg

        train_loss_results.append(epoch_loss_avg.result())
        train_cost_results.append(epoch_cost_avg.result())

        if log_results:
            pd.DataFrame(
                data={
                    "epochs": list(range(start_epoch, epoch + 1)),
                    "train_loss": [x.cpu().numpy() for x in train_loss_results],
                    "train_cost": [x.cpu().numpy() for x in train_cost_results],
                    "val_cost": [x.cpu().numpy() for x in val_cost_avg],
                }
            ).to_csv(f"{mod_dir}/adm_training_results.csv", index=False)

        print(
            get_cur_time(),
            "Epoch {}: Loss: {}: Cost: {}".format(
                epoch, epoch_loss_avg.result(), epoch_cost_avg.result()
            ),
        )
