"""Centralized tensor presentation formatting obeying the 16-element budget.

Tensor-value invariant: Wherever the visual grammar claims to display the value
of a quantity, adjoint, gradient contribution, or numerical backward computation,
it displays a representation of the actual numerical value. Shape information is
metadata and must never substitute for a value.
"""
from dataclasses import dataclass
import numpy as np
from number_format import format_num, DEFAULT_FORMAT_CONFIG, NumberFormatConfig

DEFAULT_TENSOR_ELEMENT_BUDGET = 16


@dataclass
class TensorDisplay:
    text: str
    shape: tuple[int, ...]
    truncated: bool
    element_count: int


def _format_vector(arr, truncated=False, config: NumberFormatConfig | None = None):
    if not truncated:
        items = [format_num(float(x), config=config) for x in arr]
        return f"[{', '.join(items)}]"
    n = len(arr)
    # Symmetric truncation: 2 at start, 1 at end (or 2 at end if length allows)
    # E.g. [a0, a1, …, an-1]
    first = [format_num(float(x), config=config) for x in arr[:2]]
    last = [format_num(float(x), config=config) for x in arr[-1:]]
    return f"[{', '.join(first)}, …, {', '.join(last)}]"


def _format_matrix(arr, truncated=False, config: NumberFormatConfig | None = None):
    m, n = arr.shape
    if not truncated:
        rows = []
        for r in arr:
            row_items = [format_num(float(x), config=config) for x in r]
            rows.append(f"[{', '.join(row_items)}]")
        return f"[{', '.join(rows)}]"
    
    # Truncated matrix: preserve beginning and ending rows/columns
    # Rows to show: if m > 2, show row 0, row 1, ..., row m-1
    # Cols to show: if n > 2, show col 0, col 1, ..., col n-1
    def format_row(r):
        if n > 2:
            return f"[{format_num(float(r[0]), config=config)}, {format_num(float(r[1]), config=config)}, …, {format_num(float(r[-1]), config=config)}]"
        else:
            cols = [format_num(float(x), config=config) for x in r]
            return f"[{', '.join(cols)}]"

    row_strs = []
    if m <= 2:
        for r in arr:
            row_strs.append(format_row(r))
    else:
        # First row (or two rows if m >= 4)
        num_top = 2 if m >= 4 else 1
        for i in range(num_top):
            row_strs.append(format_row(arr[i]))
        v_ellipsis = ["⋮", "⋮", "⋱", "⋮"] if n > 2 else ["⋮"] * n
        row_strs.append(f"[{', '.join(v_ellipsis)}]")
        row_strs.append(format_row(arr[-1]))
        
    return f"[{', '.join(row_strs)}]"


def _format_nd(arr, truncated=False, config: NumberFormatConfig | None = None):
    if arr.ndim == 1:
        return _format_vector(arr, truncated=truncated, config=config)
    if arr.ndim == 2:
        return _format_matrix(arr, truncated=truncated, config=config)
    
    if not truncated:
        sub_strs = [_format_nd(sub, truncated=False, config=config) for sub in arr]
        return f"[{', '.join(sub_strs)}]"
    else:
        if len(arr) <= 2:
            sub_strs = [_format_nd(sub, truncated=True, config=config) for sub in arr]
            return f"[{', '.join(sub_strs)}]"
        sub_first = _format_nd(arr[0], truncated=True, config=config)
        sub_last = _format_nd(arr[-1], truncated=True, config=config)
        return f"[{sub_first}, …, {sub_last}]"


def make_tensor_display(val, element_budget=DEFAULT_TENSOR_ELEMENT_BUDGET, config: NumberFormatConfig | None = None) -> TensorDisplay:
    """Produces structured TensorDisplay distinguishing text representation from shape metadata."""
    if val is None:
        return TensorDisplay(text="—", shape=(), truncated=False, element_count=0)
    
    if isinstance(val, (int, float, np.floating, np.integer)):
        val_f = float(val)
        return TensorDisplay(text=format_num(val_f, config=config), shape=(), truncated=False, element_count=1)
    
    if not isinstance(val, np.ndarray):
        val = np.asarray(val, dtype=np.float64)
        
    if val.ndim == 0:
        return TensorDisplay(text=format_num(float(val), config=config), shape=(), truncated=False, element_count=1)
        
    size = val.size
    shape = val.shape
    truncated = size > element_budget
    
    if val.ndim == 1:
        text = _format_vector(val, truncated=truncated, config=config)
    elif val.ndim == 2:
        text = _format_matrix(val, truncated=truncated, config=config)
    else:
        text = _format_nd(val, truncated=truncated, config=config)
        
    return TensorDisplay(text=text, shape=shape, truncated=truncated, element_count=size)


def format_tensor_value(val, element_budget=DEFAULT_TENSOR_ELEMENT_BUDGET, include_shape_if_truncated=True, config: NumberFormatConfig | None = None) -> str:
    """Returns canonical string representation of tensor value.
    
    If truncated and include_shape_if_truncated is True, appends shape metadata.
    Shape information is NEVER substituted for the value itself.
    """
    disp = make_tensor_display(val, element_budget=element_budget, config=config)
    if disp.truncated and include_shape_if_truncated:
        return f"{disp.text} shape {disp.shape}"
    return disp.text


def format_tensor_shape(val) -> str:
    """Format shape metadata only."""
    if val is None:
        return ""
    if isinstance(val, (int, float, np.floating, np.integer)):
        return "()"
    arr = np.asarray(val)
    return str(arr.shape)