import torch
import wandb
from torch.cuda.amp import GradScaler

from datasets.stas import crop_around_receptive_field
from training.measures import correlation


scaler = GradScaler()


def model_step(
        images, model, max_coordinate=None, rf_size=None, h=150, w=200, data_key=None
):
    if isinstance(rf_size, int):
        cur_rf_size = (rf_size, rf_size)
    elif len(rf_size) == 3:
        cur_rf_size = rf_size[1:]
    else:
        cur_rf_size = rf_size
    if cur_rf_size != (h, w):
        images = crop_around_receptive_field(
            max_coordinate=max_coordinate, images=images, rf_size=cur_rf_size, h=h, w=w
        )

    output = model(images.double(), data_key=data_key)
    if len(output.shape) == 1:
        output = torch.unsqueeze(output, 1)
    elif (len(output.shape) >= 3) and (output.shape[2] == 1):
        output = torch.squeeze(output, 2)

    return output

def train_step(
        images,
        responses,
        optimizer,
        model,
        loss_function,
        rf_size=None,
        max_coordinate=None,
        h=150,
        w=200,
        regularization=True,
        penalty_coef=1,
        data_key=None,
        regularizer_start=0,
        epoch=None,
        track_gradients=True,
        grad_step=0,
        wandb_log=True

):
    if regularizer_start != 0:
        assert epoch is not None
    optimizer.zero_grad()
    output = model_step(
        images, model, max_coordinate, rf_size, h, w, data_key=data_key
    )

    if output is None:
        return None, None, None, None

    if output.shape != responses.shape:
        responses = responses.squeeze(-1)
    assert output.shape == responses.shape

    corr = correlation(output, responses, 1e-12)

    if regularization and ((epoch is None) or (epoch >= regularizer_start)):
        core_spatial, core_temporal, readout = model.regularizer(data_key=data_key)
        penalty = core_spatial + core_temporal + readout
    else:
        penalty = 0
        core_spatial, core_temporal, readout = (
            torch.Tensor([0]),
            torch.Tensor([0]),
            torch.Tensor([0]),
        )

    loss = loss_function(output.double(), responses.double()) + penalty_coef * penalty

    loss.backward()
    optimizer.step()

    return loss, [core_spatial, core_temporal, readout], corr


def save_checkpoint(epoch, model, path, optimizer, valid_corr):
    torch.save(
        {
            "epoch": epoch,
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict() if optimizer is not None else None,
            "valid_corr": valid_corr,
        },
        path,
    )


