"""ModelLog: the top-level object that stores the full representation of a logged forward pass."""

import copy
from collections import OrderedDict, defaultdict
from typing import Any, Callable, Dict, List, Optional, Set, TYPE_CHECKING, Tuple

if TYPE_CHECKING:
    from .buffer_log import BufferAccessor
    from .layer_log import LayerAccessor

from .cleanup import _remove_log_entry, _batch_remove_log_entries, cleanup
from .module_log import ModuleAccessor
from .param_log import ParamAccessor
from ..utils.display import human_readable_size
from .interface import (
    _getitem_after_pass,
    _getitem_during_pass,
    _str_after_pass,
    _str_during_pass,
    to_pandas,
    print_all_fields,
)
from ..capture.trace import save_new_activations
from ..postprocess import postprocess
from .layer_log import LayerLog
from .layer_pass_log import LayerPassLog, TensorLog
from ..capture.trace import run_and_log_inputs_through_model
from ..validation import validate_saved_activations
from ..validation import check_metadata_invariants
from ..visualization.rendering import render_graph


def _init_module_build_data() -> dict:
    """Create the transient dict used to accumulate module hierarchy data during logging.

    Consumed by ``_build_module_logs`` (step 17) and then cleared.
    """
    return {
        "module_addresses": [],
        "module_types": {},
        "module_passes": [],
        "module_num_passes": defaultdict(lambda: 1),
        "top_level_modules": [],
        "top_level_module_passes": [],
        "module_children": defaultdict(list),
        "module_pass_children": defaultdict(list),
        "module_nparams": defaultdict(lambda: 0),
        "module_nparams_trainable": defaultdict(lambda: 0),
        "module_nparams_frozen": defaultdict(lambda: 0),
        "module_num_tensors": defaultdict(lambda: 0),
        "module_pass_num_tensors": defaultdict(lambda: 0),
        "module_layers": defaultdict(list),
        "module_pass_layers": defaultdict(list),
        "module_layer_argnames": defaultdict(list),
        "module_training_modes": {},
    }


class ModelLog:
    def __init__(
        self,
        model_name: str,
        output_device: str = "same",
        activation_postfunc: Optional[Callable] = None,
        keep_unsaved_layers: bool = True,
        save_function_args: bool = False,
        save_gradients: bool = False,
        detach_saved_tensors: bool = False,
        mark_input_output_distances: bool = True,
        num_context_lines: int = 7,
        optimizer=None,
    ):
        """Object that stores the history of a model's forward pass.
        Both logs the history in real time, and stores a nice
        representation of the full history for the user afterward.
        """
        # Setup:
        activation_postfunc = copy.deepcopy(activation_postfunc)

        # General info
        self.model_name = model_name
        self.num_context_lines = num_context_lines
        self._optimizer = optimizer
        self._pass_finished = False
        self.logging_mode = "exhaustive"
        self._all_layers_logged = False
        self._all_layers_saved = False
        self.keep_unsaved_layers = keep_unsaved_layers
        self.activation_postfunc = activation_postfunc
        self.current_function_call_barcode = None
        self.random_seed_used = None
        self.output_device = output_device
        self.detach_saved_tensors = detach_saved_tensors
        self.save_function_args = save_function_args
        self.save_gradients = save_gradients
        self.has_saved_gradients = False
        self.mark_input_output_distances = mark_input_output_distances

        # Model structure info
        self.model_is_recurrent = False
        self.model_max_recurrent_loops = 1
        self.model_has_conditional_branching = False
        self.model_is_branching = False

        # Tensor Tracking:
        self.layer_list: List[LayerPassLog] = []
        self.layer_dict_main_keys: Dict[str, LayerPassLog] = OrderedDict()
        self.layer_dict_all_keys: Dict[str, LayerPassLog] = OrderedDict()
        self.layer_logs: Dict[str, LayerLog] = OrderedDict()
        self.layer_labels: List[str] = []
        self.layer_labels_w_pass: List[str] = []
        self.layer_labels_no_pass: List[str] = []
        self.layer_num_passes: Dict[str, int] = OrderedDict()
        self._raw_layer_dict: Dict[str, LayerPassLog] = OrderedDict()
        self._raw_layer_labels_list: List[str] = []
        self._layer_nums_to_save: List[int] = []
        self._layer_counter: int = 0
        self.num_operations: int = 0
        self._raw_layer_type_counter: Dict[str, int] = defaultdict(lambda: 0)
        self._unsaved_layers_lookup_keys: Set[str] = set()

        # Mapping from raw to final layer labels:
        self._raw_to_final_layer_labels: Dict[str, str] = {}
        self._final_to_raw_layer_labels: Dict[str, str] = {}
        self._lookup_keys_to_layer_num_dict: Dict[str, int] = {}
        self._layer_num_to_lookup_keys_dict: Dict[int, List[str]] = defaultdict(list)

        # Special Layers:
        self.input_layers: List[str] = []
        self.output_layers: List[str] = []
        self.buffer_layers: List[str] = []
        self.buffer_num_passes: Dict = {}
        self._buffer_accessor = None
        self.internally_initialized_layers: List[str] = []
        self._layers_where_internal_branches_merge_with_input: List[str] = []
        self.internally_terminated_layers: List[str] = []
        self.internally_terminated_bool_layers: List[str] = []
        self.conditional_branch_edges: List[Tuple[str, str]] = []
        self.layers_with_saved_activations: List[str] = []
        self.orphan_layers: List[str] = []
        self.unlogged_layers: List[str] = []
        self.layers_with_saved_gradients: List[str] = []
        self._saved_gradients_set: set = set()
        self.layers_computed_with_params: Dict[str, List] = defaultdict(list)
        self.equivalent_operations: Dict[str, set] = defaultdict(set)

        # Tensor info:
        self.num_tensors_total: int = 0
        self.tensor_fsize_total: int = 0
        self.tensor_fsize_total_nice: str = human_readable_size(0)
        self.num_tensors_saved: int = 0
        self.tensor_fsize_saved: int = 0
        self.tensor_fsize_saved_nice: str = human_readable_size(0)

        # Param info:
        self.param_logs: "ParamAccessor" = ParamAccessor({})
        self.total_param_tensors: int = 0
        self.total_param_layers: int = 0
        self.total_params: int = 0
        self.total_params_trainable: int = 0
        self.total_params_frozen: int = 0
        self.total_params_fsize: int = 0
        self.total_params_fsize_nice: str = human_readable_size(0)

        # Transient module build data (consumed by _build_module_logs, then cleared):
        self._module_build_data: Dict = _init_module_build_data()

        # Structured module info:
        self._module_logs: ModuleAccessor = ModuleAccessor({})

        # Temporary storage for module metadata capture (consumed by _build_module_logs):
        self._module_metadata: Dict = {}
        self._module_forward_args: Dict = {}

        # Time elapsed:
        self.pass_start_time: float = 0
        self.pass_end_time: float = 0
        self.elapsed_time_setup: float = 0
        self.elapsed_time_forward_pass: float = 0
        self.elapsed_time_cleanup: float = 0
        self.elapsed_time_total: float = 0
        self.elapsed_time_function_calls: float = 0
        self.elapsed_time_torchlens_logging: float = 0

    # ********************************************
    # ************ Built-in Methods **************
    # ********************************************

    def __len__(self):
        if self._pass_finished:
            return len(self.layer_list)
        else:
            return len(self._raw_layer_dict)

    def __getitem__(self, ix) -> LayerPassLog:
        """Returns an object logging a model layer given an index. If the pass is finished,
        it'll do this intelligently; if not, it simply queries based on the layer's raw barcode.

        Args:
            ix: desired index

        Returns:
            Tensor log entry object with info about specified layer.
        """
        if self._pass_finished:
            return _getitem_after_pass(self, ix)
        else:
            return _getitem_during_pass(self, ix)

    def __str__(self) -> str:
        """Human-readable summary; delegates to post-pass or mid-pass formatter."""
        if self._pass_finished:
            return _str_after_pass(self)
        else:
            return _str_during_pass(self)

    def __repr__(self) -> str:
        """Same as __str__."""
        return self.__str__()

    def __iter__(self):
        """Loops through all tensors in the log."""
        if self._pass_finished:
            return iter(self.layer_list)
        else:
            return iter(list(self._raw_layer_dict.values()))

    # ********************************************
    # ************* FLOPs Properties *************
    # ********************************************

    @property
    def total_flops_forward(self) -> int:
        """Total forward FLOPs across all layers (skipping None/unknown)."""
        return sum(
            entry.flops_forward for entry in self.layer_list if entry.flops_forward is not None
        )

    @property
    def total_flops_backward(self) -> int:
        """Total backward FLOPs across all layers (skipping None/unknown)."""
        return sum(
            entry.flops_backward for entry in self.layer_list if entry.flops_backward is not None
        )

    @property
    def total_flops(self) -> int:
        """Total FLOPs (forward + backward)."""
        return self.total_flops_forward + self.total_flops_backward

    def flops_by_type(self) -> Dict[str, Dict[str, int]]:
        """Group FLOPs by layer type.

        Returns:
            Dict mapping layer_type to {"forward": int, "backward": int, "count": int}.
        """
        result: Dict[str, Dict[str, int]] = {}
        for entry in self.layer_list:
            lt = entry.layer_type
            if lt not in result:
                result[lt] = {"forward": 0, "backward": 0, "count": 0}
            result[lt]["count"] += 1
            if entry.flops_forward is not None:
                result[lt]["forward"] += entry.flops_forward
            if entry.flops_backward is not None:
                result[lt]["backward"] += entry.flops_backward
        return result

    # ********************************************
    # ************* Params Accessor **************
    # ********************************************

    @property
    def params(self) -> ParamAccessor:
        """Access parameter metadata by address, short name, or index."""
        return self.param_logs

    @property
    def layers(self) -> "LayerAccessor":
        """Access aggregate per-layer metadata by label, index, or pass notation."""
        from .layer_log import LayerAccessor

        return LayerAccessor(self.layer_logs, source_model_log=self)

    @property
    def modules(self) -> "ModuleAccessor":
        """Access structured per-module metadata by address, index, or pass notation."""
        return self._module_logs

    @property
    def root_module(self):
        """The root module (the model itself)."""
        return self._module_logs["self"]

    @property
    def buffers(self) -> "BufferAccessor":
        """Access buffer metadata by address, short name, or index."""
        return self._buffer_accessor  # type: ignore[return-value]

    # ********************************************
    # ******** Assign Imported Methods ***********
    # ********************************************

    render_graph = render_graph
    print_all_fields = print_all_fields
    to_pandas = to_pandas
    save_new_activations = save_new_activations
    validate_saved_activations = validate_saved_activations
    validate_forward_pass = validate_saved_activations  # alias
    check_metadata_invariants = check_metadata_invariants
    cleanup = cleanup
    _postprocess = postprocess
    _run_and_log_inputs_through_model = run_and_log_inputs_through_model
    _remove_log_entry = _remove_log_entry
    _batch_remove_log_entries = _batch_remove_log_entries
