import os
from typing import Callable
import math

import numpy as np
from bson import ObjectId
import torch
from torch import nn

from utils.data_utils import get_collated_batch_for_key


def simple_get_y(collated_batch):
    """
    Get y from collated batch
    Args:
        collated_batch: collated batch to get y from

    Returns:
        np.ndarray: y
    """
    return collated_batch[1]


def simple_model_save(model: nn.Module,
                      training_configuration: dict) -> None:
    """
    Save model to disk
    Args:
        model: model to save
        training_configuration: training configuration
    Returns:
        None
    """
    model_path = os.path.join(training_configuration["training_dir"], "model.pt")
    torch.save(model.state_dict(), model_path)


def simple_model_load(model_configuration: dict,
                      training_configuration: dict,
                      sample_key: str | ObjectId,
                      device: str,
                      model_creation_fn: Callable,
                      *args, **kwargs) -> nn.Module:
    """
    Load model from disk
    Args:
        model_configuration: model configuration
        training_configuration: training configuration
        sample_key: sample key to get sample data batch with
        device: device to load model on
        model_creation_fn: function to create model
    Returns:
        model: loaded model
    """
    training_dir = training_configuration["training_dir"]
    model_state_path = os.path.join(training_dir, "model.pt")
    model = model_creation_fn(model_configuration,
                              sample_key, *args, **kwargs)
    model.to(device)
    model.load_state_dict(torch.load(model_state_path))
    return model


def simple_model_creation(model_configuration: dict,
                          sample_key: str | ObjectId,
                          lmdb_env=None) -> nn.Module:
    """
    Create model from configuration
    Args:
        model_configuration: model configuration
        sample_key: sample key to get sample data batch with
        lmdb_env: LMDB environment to use for getting sample data batch

    Returns:
        model: created model
    """
    sample_item = get_collated_batch_for_key(sample_key, model_configuration, lmdb_env=lmdb_env)
    input_size = sample_item[0].shape[2]
    output_size = sample_item[1].shape[2]
    model_class = model_configuration["model_class"]
    cnn_model = model_class(input_dim=input_size,
                            output_dim=output_size,
                            **model_configuration["model_parameters"], )
    return cnn_model


def simple_x_y_predict(model: nn.Module,
                       collated_batch: tuple[torch.Tensor, torch.Tensor],
                       device: str,
                       batch_size: int,
                       *args, **kwargs) -> torch.Tensor:
    """
    Predict y from x
    Args:
        model: model to use for prediction, must have 'predict' method
        collated_batch: collated data batch to predict with,
            must be a tuple of (x, y) where x is the input data and y is the target data
        device: device to use for prediction
        *args: additional arguments to pass to the model's predict method
        **kwargs: additional keyword arguments to pass to the model's predict method

    Returns:
        torch.Tensor: predicted y as numpy array on CPU
    """

    model.eval()
    x = collated_batch[0].to(device).float()
    batches = list()
    num_batches = math.ceil(len(x) / batch_size)
    for i in range(num_batches):
        start = i * batch_size
        end = (i + 1) * batch_size
        if end > len(x):
            end = len(x)
        batch_slice = x[start:end]
        batches.append(batch_slice)

    preds = list()
    with torch.no_grad():
        for batch in batches:
            batch = batch.to(device).float()
            pred = model(batch, *args, **kwargs)
            preds.append(pred.cpu().numpy())
    return np.concatenate(preds)


def simple_get_y(collated_batch: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor:
    """
    Get y from collated batch
    Args:
        collated_batch: collated batch to get y from

    Returns:
        torch.Tensor: y
    """
    return collated_batch[1]