from __future__ import annotations

import typing
from typing import List, Optional, Union

import numpy

from .. import tracers
from openfisca_core.indexed_enums import EnumArray

if typing.TYPE_CHECKING:
    from numpy.typing import ArrayLike

    Array = Union[EnumArray, ArrayLike]


class ComputationLog:

    _full_tracer: tracers.FullTracer

    def __init__(self, full_tracer: tracers.FullTracer) -> None:
        self._full_tracer = full_tracer

    def display(
            self,
            value: Optional[Array],
            ) -> str:
        if isinstance(value, EnumArray):
            value = value.decode_to_str()

        return numpy.array2string(value, max_line_width = float("inf"))

    def _get_node_log(
            self,
            node: tracers.TraceNode,
            depth: int,
            aggregate: bool,
            depth_max: int = None,
            ) -> List[str]:

        def print_line(depth: int, node: tracers.TraceNode) -> str:
            indent = '  ' * depth
            value = node.value

            if value is None:
                formatted_value = "{'avg': '?', 'max': '?', 'min': '?'}"

            elif aggregate:
                try:
                    formatted_value = str({
                        'avg': numpy.mean(value),
                        'max': numpy.max(value),
                        'min': numpy.min(value),
                        })

                except TypeError:
                    formatted_value = "{'avg': '?', 'max': '?', 'min': '?'}"

            else:
                formatted_value = self.display(value)

            return f"{indent}{node.name}<{node.period}> >> {formatted_value}"

        if depth_max is not None and depth > depth_max:
            return

        node_log = [print_line(depth, node)]

        children_logs = [
            self._get_node_log(child, depth + 1, aggregate, depth_max = depth_max)
            for child
            in node.children
            if self._get_node_log(child, depth + 1, aggregate, depth_max = depth_max) is not None
            ]

        return node_log + self._flatten(children_logs)

    def _flatten(
            self,
            list_of_lists: List[List[str]],
            ) -> List[str]:
        return [item for _list in list_of_lists for item in _list]

    def lines(self, aggregate: bool = False, depth_max:int = None) -> List[str]:
        depth = 1

        lines_by_tree = [
            self._get_node_log(node, depth, aggregate, depth_max = depth_max)
            for node
            in self._full_tracer.trees
            ]

        return self._flatten(lines_by_tree)

    def print_log(self, aggregate = False, depth_max = None) -> None:
        """
        Print the computation log of a simulation.

        If ``aggregate`` is ``False`` (default), print the value of each
        computed vector.

        If ``aggregate`` is ``True``, only print the minimum, maximum, and
        average value of each computed vector.

        This mode is more suited for simulations on a large population.
        """
        for line in self.lines(aggregate, depth_max):
            print(line)  # noqa T001

@@ -1675,16 +1675,17 @@
      if 
+(
 depth_ma
@@ -1701,13 +1701,15 @@
 None
+)
  and 
+(
 dept
@@ -1721,16 +1721,17 @@
 epth_max
+)
 :%0A      
@@ -1742,16 +1742,21 @@
   return
+ %5B%22%22%5D
 %0A%0A      
@@ -2041,27 +2041,23 @@
 th_max) 
-is not None
+!= %5B%22%22%5D
 %0A       
@@ -2296,16 +2296,16 @@
 _list%5D%0A%0A
-
     def 
@@ -2351,16 +2351,17 @@
 pth_max:
+ 
 int = No

Fix typing