import os
import torch
import json
import warnings
import prov.model as prov

from torch.utils.data import DataLoader, Subset, Dataset, RandomSampler
from typing import Any, Optional, Union

from yprov4ml.datamodel.attribute_type import LoggingItemKind
from yprov4ml.utils import energy_utils, flops_utils, system_utils, time_utils, funcs
from yprov4ml.datamodel.context import Context
from yprov4ml.constants import PROV4ML_DATA, VERBOSE
    
def log_metric(
        key: str, 
        value: float, 
        context: Optional[Context] = None, 
        step: Optional[int] = 0, 
        source: Optional[LoggingItemKind] = None, 
    ) -> None:
    """
    Logs a metric with the specified key, value, and context.

    Args:
        key (str): The key of the metric.
        value (float): The value of the metric.
        context (Context): The context in which the metric is recorded.
        step (Optional[int], optional): The step number for the metric. Defaults to None.
        source (LoggingItemKind, optional): The source of the logging item. Defaults to None.

    Returns:
        None
    """
    PROV4ML_DATA.add_metric(key, value, step, context=context, source=source)

def log_execution_start_time() -> None:
    """Logs the start time of the current execution. """
    return log_param("execution_start_time", time_utils.get_time())

def log_execution_end_time() -> None:
    """Logs the end time of the current execution."""
    return log_param("execution_end_time", time_utils.get_time())

def log_current_execution_time(label: str, context: Context, step: Optional[int] = None) -> None:
    """Logs the current execution time under the given label.
    
    Args:
        label (str): The label to associate with the logged execution time.
        context (mlflow.tracking.Context): The MLflow tracking context.
        step (Optional[int], optional): The step number for the logged execution time. Defaults to None.

    Returns:
        None
    """
    return log_metric(label, time_utils.get_time(), context, step=step, source=LoggingItemKind.EXECUTION_TIME)

def log_param(key: str, value: Any, context : Context = None) -> None:
    """Logs a single parameter key-value pair. 
    
    Args:
        key (str): The key of the parameter.
        value (Any): The value of the parameter.

    Returns:
        None
    """
    PROV4ML_DATA.add_parameter(key,value, context)

def _get_model_memory_footprint(model_name: str, model: Union[torch.nn.Module, Any]) -> dict:
    """Logs the memory footprint of the provided model.
    
    Args:
        model (Union[torch.nn.Module, Any]): The model whose memory footprint is to be logged.
        model_name (str, optional): Name of the model. Defaults to "default".

    Returns:
        None
    """
    ret = {"model_name": model_name}


    total_params = sum(p.numel() for p in model.parameters())
    try: 
        if hasattr(model, "trainer"): 
            precision_to_bits = {"64": 64, "32": 32, "16": 16, "bf16": 16}
            if hasattr(model.trainer, "precision"):
                precision = precision_to_bits.get(model.trainer.precision, 32)
            else: 
                precision = 32
        else: 
            precision = 32
    except RuntimeError: 
        if VERBOSE: 
            warnings.warn("Could not determine precision, defaulting to 32 bits. Please make sure to provide a model with a trainer attached, this is often due to calling this before the trainer.fit() method")
        precision = 32
    
    precision_megabytes = precision / 8 / 1e6

    memory_per_model = total_params * precision_megabytes
    memory_per_grad = total_params * 4 * 1e-6
    memory_per_optim = total_params * 4 * 1e-6
    
    ret[f"{PROV4ML_DATA.PROV_PREFIX}:total_params"] = total_params
    ret[f"{PROV4ML_DATA.PROV_PREFIX}:memory_of_model"] = memory_per_model
    ret[f"{PROV4ML_DATA.PROV_PREFIX}:total_memory_load_of_model"] = memory_per_model + memory_per_grad + memory_per_optim

    return ret

def _get_nested_model_desc(m: torch.nn.Module):
    children = dict(m.named_children())
    output = {}
    if children == {}:
        node = {}
        funcs._safe_get_model_attr(node, m, "type", attr_label="layer_type")
        funcs._safe_get_model_attr(node, m, "in_features")
        funcs._safe_get_model_attr(node, m, "out_features")
        funcs._safe_get_model_attr(node, m, "in_channels")
        funcs._safe_get_model_attr(node, m, "out_channels")
        funcs._safe_get_model_attr(node, m, "kernel_size")
        funcs._safe_get_model_attr(node, m, "stride")
        funcs._safe_get_model_attr(node, m, "padding")
        funcs._safe_get_model_attr(node, m, "weight.dtype", attr_label="dtype")
        # safe_get_attr(node, m, "bias", attr_label="layer_bias")
        return node
    else:
        for name, child in children.items():
            output[f"{PROV4ML_DATA.PROV_PREFIX}:{name}"] = _get_nested_model_desc(child)

    return output

def _get_model_layers_description(model_name : str, model: Union[torch.nn.Module, Any]) -> prov.ProvEntity: 
    mo = _get_nested_model_desc(model)
    
    path = os.path.join(PROV4ML_DATA.ARTIFACTS_DIR, model_name)
    if not os.path.exists(path):
        os.makedirs(path, exist_ok=True)
    path = os.path.join(path, f"{model_name}_layers_description.json")
    with open(path, "w") as fp:
        json.dump(mo , fp) 

    return {f"{PROV4ML_DATA.PROV_PREFIX}:layers_description_path": path}

def log_model(
        model_name: str, 
        model: Union[torch.nn.Module, Any], 
        context : Optional[Context] = Context.MODELS, 
        log_model_info: bool = True, 
        log_model_layers : bool = False,
        is_input: bool = False,
    ) -> None:
    """Logs the provided model as artifact and logs memory footprint of the model. 
    
    Args:
        model (Union[torch.nn.Module, Any]): The model to be logged.
        model_name (str, optional): Name of the model. Defaults to "default".
        log_model_info (bool, optional): Whether to log model memory footprint. Defaults to True.
        log_model_layers (bool, optional): Whether to log model layers details. Defaults to False.
        log_as_artifact (bool, optional): Whether to log the model as an artifact. Defaults to True.
    """
    e = save_model_version(model_name, model, context, incremental=False, is_input=is_input)

    if log_model_info:
        d = _get_model_memory_footprint(model_name, model)
        e.add_attributes(d)

    if log_model_layers: 
        d = _get_model_layers_description(model_name, model)
        e.add_attributes(d)

     
def log_flops_per_epoch(label: str, model: Any, dataset: Any, context: Context, step: Optional[int] = None) -> None:
    """Logs the number of FLOPs (floating point operations) per epoch for the given model and dataset.
    
    Args:
        label (str): The label to associate with the logged FLOPs per epoch.
        model (Any): The model for which FLOPs per epoch are to be logged.
        dataset (Any): The dataset used for training the model.
        context (mlflow.tracking.Context): The MLflow tracking context.
        step (Optional[int], optional): The step number for the logged FLOPs per epoch. Defaults to None.

    Returns:
        None
    """
    return log_metric(label, flops_utils.get_flops_per_epoch(model, dataset), context, step=step, source=LoggingItemKind.FLOPS_PER_EPOCH)

def log_flops_per_batch(label: str, model: Any, batch: Any, context: Context, step: Optional[int] = None) -> None:
    """Logs the number of FLOPs (floating point operations) per batch for the given model and batch of data.
    
    Args:
        label (str): The label to associate with the logged FLOPs per batch.
        model (Any): The model for which FLOPs per batch are to be logged.
        batch (Any): A batch of data used for inference with the model.
        context (mlflow.tracking.Context): The MLflow tracking context.
        step (Optional[int], optional): The step number for the logged FLOPs per batch. Defaults to None.

    Returns:
        None
    """
    return log_metric(label, flops_utils.get_flops_per_batch(model, batch), context, step=step, source=LoggingItemKind.FLOPS_PER_BATCH)

def log_system_metrics(
    context: Context,
    step: Optional[int] = None,
    ) -> None:
    """Logs system metrics such as CPU usage, memory usage, disk usage, and GPU metrics.

    Args:
        context (mlflow.tracking.Context): The MLflow tracking context.
        step (Optional[int], optional): The step number for the logged metrics. Defaults to None.

    Returns:
        None
    """
    log_metric("cpu_usage", system_utils.get_cpu_usage(), context, step=step, source=LoggingItemKind.SYSTEM_METRIC)
    log_metric("memory_usage", system_utils.get_memory_usage(), context, step=step, source=LoggingItemKind.SYSTEM_METRIC)
    log_metric("disk_usage", system_utils.get_disk_usage(), context, step=step, source=LoggingItemKind.SYSTEM_METRIC)

    
    log_metric("gpu_memory_usage", system_utils.get_gpu_memory_usage(), context, step=step, source=LoggingItemKind.SYSTEM_METRIC)
    log_metric("gpu_usage", system_utils.get_gpu_usage(), context, step=step, source=LoggingItemKind.SYSTEM_METRIC)
    log_metric("gpu_temperature", system_utils.get_gpu_temperature(), context, step=step, source=LoggingItemKind.SYSTEM_METRIC)
    log_metric("gpu_power_usage", system_utils.get_gpu_power_usage(), context, step=step, source=LoggingItemKind.SYSTEM_METRIC)

def log_carbon_metrics(
    context: Context,
    step: Optional[int] = None,
    ):
    """Logs carbon emissions metrics such as energy consumed, emissions rate, and power consumption.
    
    Args:
        context (mlflow.tracking.Context): The MLflow tracking context.
        step (Optional[int], optional): The step number for the logged metrics. Defaults to None.
    
    Returns:
        None
    """    
    emissions = energy_utils.stop_carbon_tracked_block()
   
    log_metric("emissions", emissions.energy_consumed, context, step=step, source=LoggingItemKind.CARBON_METRIC)
    log_metric("emissions_rate", emissions.emissions_rate, context, step=step, source=LoggingItemKind.CARBON_METRIC)
    log_metric("cpu_power", emissions.cpu_power, context, step=step, source=LoggingItemKind.CARBON_METRIC)
    log_metric("gpu_power", emissions.gpu_power, context, step=step, source=LoggingItemKind.CARBON_METRIC)
    log_metric("ram_power", emissions.ram_power, context, step=step, source=LoggingItemKind.CARBON_METRIC)
    log_metric("cpu_energy", emissions.cpu_energy, context, step=step, source=LoggingItemKind.CARBON_METRIC)
    log_metric("gpu_energy", emissions.gpu_energy, context, step=step, source=LoggingItemKind.CARBON_METRIC)
    log_metric("ram_energy", emissions.ram_energy, context, step=step, source=LoggingItemKind.CARBON_METRIC)
    log_metric("energy_consumed", emissions.energy_consumed, context, step=step, source=LoggingItemKind.CARBON_METRIC)

def log_artifact(
        artifact_name : str, 
        artifact_path : str, 
        context: Optional[Context] = None,
        step: Optional[int] = None, 
        log_copy_in_prov_directory : bool = True, 
        is_model : bool = False, 
        is_input : bool = False, 
    ) -> prov.ProvEntity:
    """
    Logs the specified artifact to the given context.

    Parameters:
        artifact_path (str): The file path of the artifact to log.
        context (Context): The context in which the artifact is logged.
        step (Optional[int]): The step or epoch number associated with the artifact. Defaults to None.
        timestamp (Optional[int]): The timestamp associated with the artifact. Defaults to None.

    Returns:
        None
    """
    return PROV4ML_DATA.add_artifact(
        artifact_name=artifact_name, 
        artifact_path=artifact_path, 
        step=step, 
        context=context, 
        log_copy_in_prov_directory=log_copy_in_prov_directory, 
        is_model=is_model, 
        is_input=is_input, 
    )

def save_model_version(
        model_name: str, 
        model: Union[torch.nn.Module, Any], 
        context: Optional[Context] = None, 
        step: Optional[int] = None, 
        incremental : bool = True, 
        is_input : bool =False, 
    ) -> prov.ProvEntity:
    """
    Saves the state dictionary of the provided model and logs it as an artifact.
    
    Parameters:
        model (torch.nn.Module): The PyTorch model to be saved.
        model_name (str): The name under which to save the model.
        context (Context): The context in which the model is saved.
        step (Optional[int]): The step or epoch number associated with the saved model. Defaults to None.
        timestamp (Optional[int]): The timestamp associated with the saved model. Defaults to None.

    Returns:
        None
    """

    path = os.path.join(PROV4ML_DATA.ARTIFACTS_DIR, model_name)
    if not os.path.exists(path):
        os.makedirs(path, exist_ok=True)

    # count all models with the same name stored at "path"
    if incremental: 
        num_files = len([file for file in os.listdir(path) if str(file).startswith(model_name)])
        fn = os.path.join(path, f"{model_name}_{num_files}.pt")
        torch.save(model.state_dict(), fn)
        return log_artifact(f"{model_name}_{num_files}", fn, context=context, step=step, log_copy_in_prov_directory=False, is_model=True, is_input=is_input)
    else: 
        fn = os.path.join(path, f"{model_name}.pt")
        torch.save(model.state_dict(), fn)
        return log_artifact(model_name, fn, context=context, step=step, log_copy_in_prov_directory=False, is_model=True, is_input=is_input)

def log_dataset(
        dataset_label : str, 
        dataset : Union[DataLoader, Subset, Dataset], 
        context : Optional[Context] = Context.DATASETS, 
        log_dataset_info : bool = True, 
        ): 
    """
    Logs dataset statistics such as total samples and total steps.

    Args:
        dataset (Union[DataLoader, Subset, Dataset]): The dataset for which statistics are to be logged.
        label (str): The label to associate with the logged dataset statistics.

    Returns:
        None
    """

    e = log_artifact(f"{dataset_label}", "", context=context, log_copy_in_prov_directory=False, is_model=False, is_input=True)
    
    if not log_dataset_info: return

    e.add_attributes({f"{PROV4ML_DATA.PROV_PREFIX}:{dataset_label}_stat_total_samples": len(dataset)})
    # handle datasets from DataLoader
    if isinstance(dataset, DataLoader):
        dl = dataset
        dataset = dl.dataset
        attrs = {
            f"{PROV4ML_DATA.PROV_PREFIX}:{dataset_label}_stat_batch_size": dl.batch_size, 
            f"{PROV4ML_DATA.PROV_PREFIX}:{dataset_label}_stat_num_workers": dl.num_workers, 
            f"{PROV4ML_DATA.PROV_PREFIX}:{dataset_label}_stat_shuffle": isinstance(dl.sampler, RandomSampler), 
            f"{PROV4ML_DATA.PROV_PREFIX}:{dataset_label}_stat_total_steps": len(dl), 
        }
        e.add_attributes(attrs)

    elif isinstance(dataset, Subset):
        dl = dataset
        dataset = dl.dataset
        e.add_attributes({f"{dataset_label}_stat_total_steps": len(dl)})

def log_execution_command(cmd: str, path : str) -> None:
    """
    Logs the execution command.
    
    Args:
        cmd (str): The command to be logged.
    """
    path = os.path.join("", "workspace", f"{PROV4ML_DATA.CLEAN_EXPERIMENT_NAME}_{PROV4ML_DATA.RUN_ID}", "artifacts", path)
    log_param("execution_command", cmd + " " + path)

def log_source_code(path: Optional[str] = None) -> None:
    """
    Logs the source code location, either from a Git repository or a specified path.
    
    Args: 
        path (Optional[str]): The path to the source code. If None, attempts to retrieve from Git.
    """
    PROV4ML_DATA.add_source_code(path)

def create_context(context : str, is_subcontext_of=None): 
    PROV4ML_DATA.add_context(context, is_subcontext_of=is_subcontext_of)