"""Public API entry points: log_forward_pass and related user-facing functions."""

import collections.abc
import os
import random
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import pandas as pd
import torch
from torch import nn
from tqdm import tqdm

from .utils.introspection import get_vars_of_type_from_obj
from .utils.rng import set_random_seed
from .utils.display import warn_parallel
from .utils.arg_handling import safe_copy_args, safe_copy_kwargs, normalize_input_args
from .data_classes.model_log import (
    ModelLog,
)


def _unwrap_data_parallel(model: nn.Module) -> nn.Module:
    """Unwrap nn.DataParallel to get the underlying module."""
    if isinstance(model, nn.DataParallel):
        return model.module
    return model


def _move_tensors_to_device(obj, device):
    """Recursively move tensors in a nested structure to the given device."""
    if isinstance(obj, torch.Tensor):
        return obj.to(device)
    elif isinstance(obj, (list, tuple)):
        moved = [_move_tensors_to_device(item, device) for item in obj]
        return type(obj)(moved) if not isinstance(obj, tuple) else tuple(moved)
    elif isinstance(obj, collections.abc.MutableMapping):
        # Handles dict, UserDict, BatchEncoding, OrderedDict, etc.
        moved = {k: _move_tensors_to_device(v, device) for k, v in obj.items()}
        if type(obj) is dict:
            return moved
        try:
            return type(obj)(moved)
        except Exception:
            return moved
    return obj


def _run_model_and_save_specified_activations(
    model: nn.Module,
    input_args: Union[torch.Tensor, List[Any]],
    input_kwargs: Dict[Any, Any],
    layers_to_save: Optional[Union[str, List[Union[int, str]]]] = "all",
    keep_unsaved_layers: bool = True,
    output_device: str = "same",
    activation_postfunc: Optional[Callable] = None,
    mark_input_output_distances: bool = False,
    detach_saved_tensors: bool = False,
    save_function_args: bool = False,
    save_gradients: bool = False,
    random_seed: Optional[int] = None,
    num_context_lines: int = 7,
    optimizer=None,
) -> ModelLog:
    """Internal function that runs the given input through the given model, and saves the
    specified activations, as given by the tensor numbers (these will not be visible to the user;
    they will be generated from the nicer human-readable names and then fed in).

    Args:
        model: PyTorch model.
        input_args: Input arguments to the model's forward pass: either a single tensor, or a list of arguments.
        input_kwargs: Keyword arguments to the model's forward pass.
        layers_to_save: List of layers to save
        keep_unsaved_layers: Whether to keep layers in the ModelLog log if they don't have saved activations.
        output_device: device where saved tensors will be stored: either 'same' to keep unchanged, or
            'cpu' or 'cuda' to move to cpu or cuda.
        activation_postfunc: Function to apply to activations before saving them (e.g., any averaging)
        mark_input_output_distances: Whether to compute the distance of each layer from the input or output.
            This is computationally expensive for large networks, so it is off by default.
        detach_saved_tensors: whether to detach the saved tensors, so they remain attached to the computational graph
        save_function_args: whether to save the arguments to each function
        save_gradients: whether to save gradients from any subsequent backward pass
        random_seed: Which random seed to use.
        optimizer: Optional optimizer to tag which params are being optimized.

    Returns:
        ModelLog object with full log of the forward pass
    """
    # Move inputs to model device if needed (e.g. model on CUDA, inputs on CPU)
    model_device = next((p.device for p in model.parameters()), None)
    if model_device is not None:
        input_args = _move_tensors_to_device(input_args, model_device)
        if input_kwargs is not None:
            input_kwargs = _move_tensors_to_device(input_kwargs, model_device)

    model_name = str(type(model).__name__)
    model_log = ModelLog(
        model_name,
        output_device,
        activation_postfunc,
        keep_unsaved_layers,
        save_function_args,
        save_gradients,
        detach_saved_tensors,
        mark_input_output_distances,
        num_context_lines,
        optimizer,
    )
    model_log._run_and_log_inputs_through_model(
        model, input_args, input_kwargs, layers_to_save, random_seed
    )
    return model_log


def log_forward_pass(
    model: nn.Module,
    input_args: Union[torch.Tensor, List[Any], Tuple[Any]],
    input_kwargs: Optional[Dict[Any, Any]] = None,
    layers_to_save: Optional[Union[str, List]] = "all",
    keep_unsaved_layers: bool = True,
    output_device: str = "same",
    activation_postfunc: Optional[Callable] = None,
    mark_input_output_distances: bool = False,
    detach_saved_tensors: bool = False,
    save_function_args: bool = False,
    save_gradients: bool = False,
    vis_opt: str = "none",
    vis_nesting_depth: int = 1000,
    vis_outpath: str = "graph.gv",
    vis_save_only: bool = False,
    vis_fileformat: str = "pdf",
    vis_buffer_layers: bool = False,
    vis_direction: str = "bottomup",
    vis_graph_overrides: Optional[Dict] = None,
    vis_node_overrides: Optional[Dict] = None,
    vis_nested_node_overrides: Optional[Dict] = None,
    vis_edge_overrides: Optional[Dict] = None,
    vis_gradient_edge_overrides: Optional[Dict] = None,
    vis_module_overrides: Optional[Dict] = None,
    random_seed: Optional[int] = None,
    num_context_lines: int = 7,
    optimizer=None,
) -> ModelLog:
    """Runs a forward pass through a model given input x, and returns a ModelLog object containing a log
    (layer activations and accompanying layer metadata) of the forward pass for all layers specified in which_layers,
    and optionally visualizes the model graph if vis_opt is set to 'rolled' or 'unrolled'.

    In which_layers, can specify 'all', for all layers (default), or a list containing any combination of:
    1) desired layer names (e.g., 'conv2d_1_1'; if a layer has multiple passes, this includes all passes),
    2) a layer pass (e.g., conv2d_1_1:2 for just the second pass), 3) a module name to fetch the output of a particular
    module, 4) the ordinal index of a layer in the model (e.g. 3 for the third layer, -2 for the second to last, etc.),
    or 5) a desired substring with which to filter desired layers (e.g., 'conv2d' for all conv2d layers).

    Args:
        model: PyTorch model
        input_args: input arguments for model forward pass; as a list if multiple, else as a single tensor.
        input_kwargs: keyword arguments for model forward pass
        layers_to_save: list of layers to include (described above), or 'all' to include all layers.
        keep_unsaved_layers: whether to keep layers without saved activations in the log (i.e., with just metadata)
        activation_postfunc: Function to apply to tensors before saving them (e.g., channelwise averaging).
        mark_input_output_distances: whether to mark the distance of each layer from the input or output;
            False by default since this is computationally expensive.
        output_device: device where saved tensors are to be stored. Either 'same' to keep on the same device,
            or 'cpu' or 'cuda' to move them to cpu or cuda when saved.
        detach_saved_tensors: whether to detach the saved tensors, so they remain attached to the computational graph
        save_function_args: whether to save the arguments to each function involved in computing each tensor
        save_gradients: whether to save gradients from any subsequent backward pass
        vis_opt: whether, and how, to visualize the network; 'none' for
            no visualization, 'rolled' to show the graph in rolled-up format (i.e.,
            one node per layer if a recurrent network), or 'unrolled' to show the graph
            in unrolled format (i.e., one node per pass through a layer if a recurrent)
        vis_nesting_depth: How many levels of nested modules to show; 1 for only top-level modules, 2 for two
            levels, etc.
        vis_outpath: file path to save the graph visualization
        vis_save_only: whether to only save the graph visual without immediately showing it
        vis_fileformat: the format of the visualization (e.g,. 'pdf', 'jpg', etc.)
        vis_buffer_layers: whether to visualize the buffer layers
        vis_direction: either 'bottomup', 'topdown', or 'leftright'
        random_seed: which random seed to use in case model involves randomness

    Returns:
        ModelLog object with layer activations and metadata
    """
    warn_parallel()
    model = _unwrap_data_parallel(model)

    if vis_opt not in ["none", "rolled", "unrolled"]:
        raise ValueError("Visualization option must be either 'none', 'rolled', or 'unrolled'.")

    if output_device not in ["same", "cpu", "cuda"]:
        raise ValueError("output_device must be either 'same', 'cpu', or 'cuda'.")

    if type(layers_to_save) is str:
        layers_to_save = layers_to_save.lower()

    if layers_to_save in ["all", "none", None, []]:
        model_log = _run_model_and_save_specified_activations(
            model=model,
            input_args=input_args,  # type: ignore[arg-type]
            input_kwargs=input_kwargs,  # type: ignore[arg-type]
            layers_to_save=layers_to_save,
            keep_unsaved_layers=keep_unsaved_layers,
            output_device=output_device,
            activation_postfunc=activation_postfunc,
            mark_input_output_distances=mark_input_output_distances,
            detach_saved_tensors=detach_saved_tensors,
            save_function_args=save_function_args,
            save_gradients=save_gradients,
            random_seed=random_seed,
            num_context_lines=num_context_lines,
            optimizer=optimizer,
        )
    else:
        model_log = _run_model_and_save_specified_activations(
            model=model,
            input_args=input_args,  # type: ignore[arg-type]
            input_kwargs=input_kwargs,  # type: ignore[arg-type]
            layers_to_save=None,
            keep_unsaved_layers=True,
            output_device=output_device,
            activation_postfunc=activation_postfunc,
            mark_input_output_distances=mark_input_output_distances,
            detach_saved_tensors=detach_saved_tensors,
            save_function_args=save_function_args,
            save_gradients=save_gradients,
            random_seed=random_seed,
            num_context_lines=num_context_lines,
            optimizer=optimizer,
        )
        model_log.keep_unsaved_layers = keep_unsaved_layers
        model_log.save_new_activations(
            model=model,
            input_args=input_args,  # type: ignore[arg-type]
            input_kwargs=input_kwargs,
            layers_to_save=layers_to_save,  # type: ignore[arg-type]
            random_seed=random_seed,
        )

    # Visualize if desired.
    if vis_opt != "none":
        model_log.render_graph(
            vis_opt,
            vis_nesting_depth,
            vis_outpath,
            vis_graph_overrides,
            vis_node_overrides,
            vis_nested_node_overrides,
            vis_edge_overrides,
            vis_gradient_edge_overrides,
            vis_module_overrides,
            vis_save_only,
            vis_fileformat,
            vis_buffer_layers,
            vis_direction,
        )

    return model_log


def get_model_metadata(
    model: nn.Module,
    input_args: Union[torch.Tensor, List[Any], Tuple[Any]],
    input_kwargs: Optional[Dict[Any, Any]] = None,
) -> ModelLog:
    """Logs all metadata for a given model and inputs without saving any activations. NOTE: this function
    will be removed in a future version of TorchLens, since calling it is identical to calling
    log_forward_pass without saving any layers.

    Args:
        model: model to inspect
        input_args: list of input positional arguments, or a single tensor
        input_kwargs: dict of keyword arguments
    Returns:
        ModelLog object with metadata about the model.
    """
    model_log = log_forward_pass(
        model,
        input_args,
        input_kwargs,
        layers_to_save=None,
        mark_input_output_distances=True,
    )
    return model_log


def show_model_graph(
    model: nn.Module,
    input_args: Union[torch.Tensor, List, Tuple],
    input_kwargs: Optional[Dict[Any, Any]] = None,
    vis_opt: str = "unrolled",
    vis_nesting_depth: int = 1000,
    vis_outpath: str = "graph.gv",
    vis_graph_overrides: Optional[Dict] = None,
    vis_node_overrides: Optional[Dict] = None,
    vis_nested_node_overrides: Optional[Dict] = None,
    vis_edge_overrides: Optional[Dict] = None,
    vis_gradient_edge_overrides: Optional[Dict] = None,
    vis_module_overrides: Optional[Dict] = None,
    save_only: bool = False,
    vis_fileformat: str = "pdf",
    vis_buffer_layers: bool = False,
    vis_direction: str = "bottomup",
    random_seed: Optional[int] = None,
) -> None:
    """Visualize the model graph without saving any activations.

    Args:
        model: PyTorch model
        input_args: Arguments for model forward pass
        input_kwargs: Keyword arguments for model forward pass
        vis_opt: whether, and how, to visualize the network; 'none' for
            no visualization, 'rolled' to show the graph in rolled-up format (i.e.,
            one node per layer if a recurrent network), or 'unrolled' to show the graph
            in unrolled format (i.e., one node per pass through a layer if a recurrent)
        vis_nesting_depth: How many levels of nested modules to show; 1 for only top-level modules, 2 for two
            levels, etc.
        vis_outpath: file path to save the graph visualization
        save_only: whether to only save the graph visual without immediately showing it
        vis_fileformat: the format of the visualization (e.g,. 'pdf', 'jpg', etc.)
        vis_buffer_layers: whether to visualize the buffer layers
        vis_direction: either 'bottomup', 'topdown', or 'leftright'
        random_seed: which random seed to use in case model involves randomness

    Returns:
        Nothing.
    """
    model = _unwrap_data_parallel(model)
    if not input_kwargs:
        input_kwargs = {}

    if vis_opt not in ["none", "rolled", "unrolled"]:
        raise ValueError("Visualization option must be either 'none', 'rolled', or 'unrolled'.")

    model_log = _run_model_and_save_specified_activations(
        model=model,
        input_args=input_args,  # type: ignore[arg-type]
        input_kwargs=input_kwargs,
        layers_to_save=None,
        activation_postfunc=None,
        mark_input_output_distances=False,
        detach_saved_tensors=False,
        save_gradients=False,
        random_seed=random_seed,
    )
    try:
        model_log.render_graph(
            vis_opt,
            vis_nesting_depth,
            vis_outpath,
            vis_graph_overrides,
            vis_node_overrides,
            vis_nested_node_overrides,
            vis_edge_overrides,
            vis_gradient_edge_overrides,
            vis_module_overrides,
            save_only,
            vis_fileformat,
            vis_buffer_layers,
            vis_direction,
        )
    finally:
        model_log.cleanup()


def validate_forward_pass(
    model: nn.Module,
    input_args: Union[torch.Tensor, List[Any], Tuple[Any]],
    input_kwargs: Optional[Dict[Any, Any]] = None,
    random_seed: Union[int, None] = None,
    verbose: bool = False,
    validate_metadata: bool = True,
) -> bool:
    """Validate that the saved model activations correctly reproduce the ground truth output, and
    optionally check all metadata invariants across the ModelLog data structures.

    Works by running a forward pass through the model, saving all activations, re-running the
    forward pass starting from the saved activations in each layer, and checking that the resulting
    output matches the original output. Additionally, it substitutes in random activations and checks
    whether the output changes accordingly. When ``validate_metadata=True`` (default), also runs
    comprehensive invariant checks on all metadata cross-references.

    Args:
        model: PyTorch model.
        input_args: Input for which to validate the saved activations.
        input_kwargs: Keyword arguments for model forward pass
        random_seed: random seed in case model is stochastic
        verbose: whether to show verbose error messages
        validate_metadata: whether to run metadata invariant checks (default True)
    Returns:
        True if the saved activations correctly reproduce the ground truth output, false otherwise.
    """
    warn_parallel()
    model = _unwrap_data_parallel(model)
    if random_seed is None:  # set random seed
        random_seed = random.randint(1, 4294967294)
    set_random_seed(random_seed)
    input_args = normalize_input_args(input_args, model)
    if not input_kwargs:
        input_kwargs = {}
    input_args_copy = safe_copy_args(input_args)
    input_kwargs_copy = safe_copy_kwargs(input_kwargs)

    # Move inputs to model device if needed (e.g. model on CUDA, inputs on CPU)
    model_device = next((p.device for p in model.parameters()), None)
    if model_device is not None:
        input_args_copy = _move_tensors_to_device(input_args_copy, model_device)
        input_kwargs_copy = _move_tensors_to_device(input_kwargs_copy, model_device)

    state_dict = model.state_dict()
    ground_truth_output_all = get_vars_of_type_from_obj(
        model(*input_args_copy, **input_kwargs_copy),
        torch.Tensor,
        search_depth=5,
        return_addresses=True,
        allow_repeats=True,
    )
    # Deduplicate by address to match capture/trace.py output extraction
    addresses_used = []
    ground_truth_output_tensors = []
    for entry in ground_truth_output_all:
        if entry[1] in addresses_used:
            continue
        ground_truth_output_tensors.append(entry[0])
        addresses_used.append(entry[1])
    model.load_state_dict(state_dict)
    model_log = _run_model_and_save_specified_activations(
        model=model,
        input_args=input_args,
        input_kwargs=input_kwargs,
        layers_to_save="all",
        keep_unsaved_layers=True,
        activation_postfunc=None,
        mark_input_output_distances=False,
        detach_saved_tensors=False,
        save_gradients=False,
        save_function_args=True,
        random_seed=random_seed,
    )
    try:
        activations_are_valid = model_log.validate_saved_activations(
            ground_truth_output_tensors, verbose, validate_metadata=validate_metadata
        )
    finally:
        model_log.cleanup()
        del model_log
    return activations_are_valid


def validate_saved_activations(
    model: nn.Module,
    input_args: Union[torch.Tensor, List[Any], Tuple[Any]],
    input_kwargs: Optional[Dict[Any, Any]] = None,
    random_seed: Union[int, None] = None,
    verbose: bool = False,
) -> bool:
    """Deprecated: use ``validate_forward_pass`` instead."""
    import warnings

    warnings.warn(
        "validate_saved_activations is deprecated, use validate_forward_pass instead",
        DeprecationWarning,
        stacklevel=2,
    )
    return validate_forward_pass(
        model, input_args, input_kwargs, random_seed=random_seed, verbose=verbose
    )


def validate_batch_of_models_and_inputs(
    models_and_inputs_dict: Dict[str, Dict[str, Union[str, Callable, Dict]]],
    out_path: str,
    redo_model_if_already_run: bool = True,
) -> pd.DataFrame:
    """Given multiple models and several inputs for each, validates the saved activations for all of them
    and returns a Pandas dataframe summarizing the validation results.

    Args:
        models_and_inputs_dict: Dict mapping each model name to a dict of model info:
                model_category: category of model (e.g., torchvision; just for directory bookkeeping)
                model_loading_func: function to load the model
                model_sample_inputs: dict of example inputs {name: input}
        out_path: Path to save the validation results to
        redo_model_if_already_run: If True, will re-run the validation for a model even if already run

    Returns:
        Pandas dataframe with validation information for each model and input
    """
    if os.path.exists(out_path):
        current_csv = pd.read_csv(out_path)
    else:
        current_csv = pd.DataFrame.from_dict(
            {
                "model_category": [],
                "model_name": [],
                "input_name": [],
                "validation_success": [],
            }
        )
    models_already_run = current_csv["model_name"].unique()
    for model_name, model_info in tqdm(models_and_inputs_dict.items(), desc="Validating models"):
        print(f"Validating model {model_name}")
        if model_name in models_already_run and not redo_model_if_already_run:
            continue
        model_category = model_info["model_category"]
        model_loading_func = model_info["model_loading_func"]
        model = model_loading_func()
        model_sample_inputs = model_info["model_sample_inputs"]
        for input_name, input_data in model_sample_inputs.items():
            validation_success = validate_forward_pass(model, input_data)
            current_csv = pd.concat(
                [
                    current_csv,
                    pd.DataFrame(
                        [
                            {
                                "model_category": model_category,
                                "model_name": model_name,
                                "input_name": input_name,
                                "validation_success": validation_success,
                            }
                        ]
                    ),
                ],
                ignore_index=True,
            )
        current_csv.to_csv(out_path, index=False)
        del model
    return current_csv
