Source code for prefgraph.algorithms.harp

"""HARP (Homothetic Axiom of Revealed Preference) test for homotheticity."""

from __future__ import annotations

import time
from typing import cast

import numpy as np
from numpy.typing import NDArray

from prefgraph.core.session import ConsumerSession
from prefgraph.core.result import HARPResult
from prefgraph.core.types import Cycle
from prefgraph._kernels import floyd_warshall_max_log_numba, bfs_find_cycle_numba


def check_harp(
    session: ConsumerSession,
    tolerance: float = 1e-10,
) -> HARPResult:
    """
    Check if consumer data satisfies HARP (Homothetic Axiom of Revealed Preference).

    HARP tests whether preferences are homothetic - demand scales proportionally
    with income/wealth. This is a stronger condition than GARP.

    For homothetic preferences, the product of expenditure ratios around any
    cycle must be <= 1. Formally:
    - Define r_ij = (p_i @ x_i) / (p_i @ x_j) (expenditure ratio)
    - HARP is violated if there exists a cycle i_1 -> i_2 -> ... -> i_n -> i_1
      such that r_{i_1,i_2} * r_{i_2,i_3} * ... * r_{i_n,i_1} > 1

    The algorithm uses log-space Floyd-Warshall to find maximum product paths,
    then checks for positive log-sum cycles.

    Args:
        session: ConsumerSession with prices and quantities
        tolerance: Numerical tolerance for comparisons

    Returns:
        HARPResult with consistency flag and violation details

    Example:
        >>> import numpy as np
        >>> from prefgraph import ConsumerSession, check_harp
        >>> prices = np.array([[1.0, 2.0], [2.0, 1.0]])
        >>> quantities = np.array([[4.0, 1.0], [1.0, 4.0]])
        >>> session = ConsumerSession(prices=prices, quantities=quantities)
        >>> result = check_harp(session)
        >>> if result.is_homothetic:
        ...     print("Preferences are homothetic")

    References:
        Varian, H. R. (1983). Non-parametric tests of consumer behaviour.
        Review of Economic Studies, 50(1), 99-110.
    """
    start_time = time.perf_counter()

    from prefgraph.algorithms.garp import check_garp

    E = session.expenditure_matrix  # T x T where E[i,j] = p_i @ x_j
    own_exp = session.own_expenditures  # e_i = E[i,i]

    # First run GARP for comparison
    garp_result = check_garp(session, tolerance)

    # Compute expenditure ratio matrix: R[i,j] = e_i / E[i,j] = (p_i @ x_i) / (p_i @ x_j)
    # Avoid division by zero
    safe_E = np.where(E > tolerance, E, np.inf)
    ratio_matrix = own_exp[:, np.newaxis] / safe_E

    # Compute log-ratio matrix for numerical stability
    # log_R[i,j] = log(e_i) - log(E[i,j])
    log_own_exp = np.log(np.maximum(own_exp, tolerance))
    log_E = np.log(np.maximum(E, tolerance))
    log_ratio_matrix = log_own_exp[:, np.newaxis] - log_E

    # Set diagonal to 0 (log(1) = 0)
    np.fill_diagonal(log_ratio_matrix, 0.0)

    # Build adjacency matrix: edge i->j exists if ratio >= 1 (i.e., log_ratio >= 0)
    # This means bundle j was affordable when i was chosen
    adjacency = ratio_matrix >= 1.0 - tolerance

    # Modified Floyd-Warshall: track maximum log-sum of paths
    max_log_product = _floyd_warshall_max_log_product(log_ratio_matrix, adjacency)

    # Check for HARP violations:
    # Violation if max_log_product[i,i] > tolerance for any i (cycle with product > 1)
    diagonal_products = np.diag(max_log_product)
    is_consistent = not np.any(diagonal_products > tolerance)

    # HARP is a binary test only - Varian (1983) and C&E (2016, Thm 4.2)
    # define no severity metric. max_cycle_product is always 1.0.
    violations: list[tuple[Cycle, float]] = []
    max_cycle_product = 1.0

    if not is_consistent:
        violations, _ = _find_harp_violations(
            log_ratio_matrix, adjacency, max_log_product, tolerance
        )

    computation_time = (time.perf_counter() - start_time) * 1000

    return HARPResult(
        is_consistent=is_consistent,
        violations=violations,
        max_cycle_product=max_cycle_product,
        expenditure_ratio_matrix=ratio_matrix,
        log_ratio_matrix=log_ratio_matrix,
        garp_result=garp_result,
        computation_time_ms=computation_time,
    )


def _floyd_warshall_max_log_product(
    log_ratios: NDArray[np.float64],
    adjacency: NDArray[np.bool_],
) -> NDArray[np.float64]:
    """
    Modified Floyd-Warshall to find maximum log-product paths.

    Instead of boolean reachability, tracks the maximum sum of log-ratios
    on any path from i to j. Uses Numba JIT for 10-50x speedup.

    Args:
        log_ratios: T x T matrix of log(expenditure ratios)
        adjacency: T x T boolean matrix of direct edges

    Returns:
        T x T matrix where result[i,j] = max sum of log_ratios on path i->j
        -inf if no path exists
    """
    log_ratios_c = np.ascontiguousarray(log_ratios, dtype=np.float64)
    adjacency_c = np.ascontiguousarray(adjacency, dtype=np.bool_)
    # numba-jitted kernel is untyped (returns Any); the kernel returns a float64
    # matrix by construction, so annotate the contract here.
    return cast(
        "NDArray[np.float64]",
        floyd_warshall_max_log_numba(log_ratios_c, adjacency_c),
    )


def _find_harp_violations(
    log_ratios: NDArray[np.float64],
    adjacency: NDArray[np.bool_],
    max_log_product: NDArray[np.float64],
    tolerance: float,
) -> tuple[list[tuple[Cycle, float]], float]:
    """
    Find cycles that violate HARP (product of ratios > 1).

    Returns:
        Tuple of (list of (cycle, product), maximum product found)
    """
    T = log_ratios.shape[0]
    violations: list[tuple[Cycle, float]] = []
    max_product = 1.0
    seen_cycles: set[frozenset[int]] = set()

    # Find nodes with positive log-product cycles
    for i in range(T):
        if max_log_product[i, i] > tolerance:
            # Reconstruct cycle through i
            cycle = _reconstruct_cycle(adjacency, log_ratios, i, tolerance)
            if cycle is not None:
                cycle_set = frozenset(cycle[:-1])
                if cycle_set not in seen_cycles:
                    seen_cycles.add(cycle_set)

                    # Compute actual product for this cycle
                    product = _compute_cycle_product(cycle, log_ratios)
                    violations.append((tuple(cycle), product))
                    max_product = max(max_product, product)

    return violations, max_product


def _compute_cycle_product(
    cycle: list[int],
    log_ratios: NDArray[np.float64],
) -> float:
    """Compute product of ratios around a cycle."""
    log_sum = 0.0
    for i in range(len(cycle) - 1):
        log_sum += log_ratios[cycle[i], cycle[i + 1]]
    return float(np.exp(log_sum))


def _reconstruct_cycle(
    adjacency: NDArray[np.bool_],
    log_ratios: NDArray[np.float64],
    start: int,
    tolerance: float,
) -> list[int] | None:
    """
    Reconstruct a cycle through the given node using BFS.

    Uses Numba JIT for fast cycle detection.
    """
    adjacency_c = np.ascontiguousarray(adjacency, dtype=np.bool_)
    cycle_arr = bfs_find_cycle_numba(adjacency_c, np.int64(start))

    if len(cycle_arr) == 0 or cycle_arr[0] == -1:
        return None

    return list(cycle_arr)


# =============================================================================
# TECH-FRIENDLY ALIASES
# =============================================================================

# validate_proportional_scaling: Tech-friendly name for check_harp
validate_proportional_scaling = check_harp
"""
Validate that user preferences scale proportionally with budget.

This is the tech-friendly alias for check_harp (HARP = Homothetic Axiom
of Revealed Preference).

Proportional preferences mean the user's relative preferences don't change
with their budget - they just scale up proportionally. This is useful for:
- User segmentation (different budget levels have same relative preferences)
- Demand prediction (can extrapolate to different spending levels)
- Aggregating users across income levels

Example:
    >>> from prefgraph import BehaviorLog, validate_proportional_scaling
    >>> result = validate_proportional_scaling(user_log)
    >>> if result.is_consistent:
    ...     print("User has proportional preferences")
"""