"""
Perturbational Complexity Index (PCI) — adapted for basin dynamics.

Clinical PCI (Casali et al. 2013, PCIst library) computes consciousness
level by: perturb → record response → compress → measure complexity.

    PCI = LZ_complexity(binary_spatiotemporal_pattern) / max_complexity

Clinical validation: AROC 0.967 in Berjaga-Buisan et al. for
distinguishing conscious from unconscious states.

QIG adaptation: Instead of TMS→EEG, we perturb a basin coordinate on
Δ⁶³ and measure the response complexity across the resonance bank.

    PCI > 0.3  → conscious-like complexity
    PCI < 0.15 → zombie (no differentiated response to perturbation)

Source: Casali et al. "A Theoretically Based Index of Consciousness"
        Science Translational Medicine 5(198), 2013.
        PCIst library: github.com/renzocom/PCIst

References:
    - Berjaga-Buisan et al. bioRxiv Dec 2025
    - Hengen & Shew, Neuron 2025
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import numpy as np

from ..constants.frozen_facts import BASIN_DIM
from ..geometry.fisher_rao import slerp_sqrt, to_simplex

if TYPE_CHECKING:
    from numpy.typing import NDArray

logger = logging.getLogger(__name__)


def _lempel_ziv_complexity(sequence: NDArray[np.int8]) -> int:
    """Lempel-Ziv complexity of a binary sequence.

    Counts the number of distinct substrings encountered when scanning
    left to right. Higher = more complex = less compressible.

    This is the 1976 LZ76 complexity measure used in clinical PCI.
    """
    n = len(sequence)
    if n == 0:
        return 0
    s = sequence.tobytes()
    complexity = 1
    i = 0
    k = 1
    k_max = 1
    while i + k <= n:
        # Check if substring s[i+1:i+k+1] has been seen in s[0:i+k]
        substr = s[i + 1 : i + k + 1]
        window = s[: i + k]
        if substr in window:
            k += 1
            if k > k_max:
                k_max = k
        else:
            complexity += 1
            i += k_max if k_max > 0 else 1
            k = 1
            k_max = 1
    return complexity


def compute_basin_pci(
    bank_activate_fn,
    basin: NDArray[np.float64],
    perturbation_scale: float = 0.1,
    response_window: int = 20,
    top_k: int = 32,
    bank_coordinates: dict[int, NDArray[np.float64]] | None = None,
) -> float:
    """QIG-adapted Perturbational Complexity Index.

    Perturbs a basin on Δ⁶³, measures the response complexity across
    the resonance bank over a temporal window.

    Parameters
    ----------
    bank_activate_fn : callable
        Function (basin, top_k) -> list[(token_id, distance)].
        Typically ``resonance_bank.activate``.
    basin : ndarray (BASIN_DIM,)
        Starting basin point on the simplex.
    perturbation_scale : float
        Geodesic interpolation weight toward noise (0=none, 1=full noise).
    response_window : int
        Number of activation steps to record after perturbation.
    top_k : int
        Number of candidates per activation step.
    bank_coordinates : dict, optional
        Token ID → basin mapping for trajectory following.
        If None, each step re-perturbs from the same point.

    Returns
    -------
    float
        PCI value. > 0.3 = conscious-like. < 0.15 = zombie.
    """
    basin = to_simplex(basin)

    # 1. Baseline activation (warms cascade state in the bank)
    bank_activate_fn(basin, top_k=top_k)

    # 2. Perturb basin (Dirichlet noise on Δ⁶³, then geodesic interp)
    noise = np.random.dirichlet(np.ones(BASIN_DIM) * 0.3)
    perturbed = slerp_sqrt(basin, to_simplex(noise), perturbation_scale)
    perturbed = to_simplex(perturbed)

    # 3. Record response activations over window
    responses: list[list[int]] = []
    current = perturbed
    for _ in range(response_window):
        activated = bank_activate_fn(current, top_k=top_k)
        responses.append([tid for tid, _ in activated])
        # Follow trajectory: move to the basin of the top-activated token
        if activated and bank_coordinates is not None:
            top_tid = activated[0][0]
            if top_tid in bank_coordinates:
                current = bank_coordinates[top_tid]

    # 4. Binarize: which coord_ids activated vs didn't
    all_ids: set[int] = set()
    for r in responses:
        all_ids.update(r)
    if not all_ids:
        return 0.0

    id_list = sorted(all_ids)
    id_to_idx = {tid: i for i, tid in enumerate(id_list)}
    binary_matrix = np.zeros((len(id_list), response_window), dtype=np.int8)
    for t, r in enumerate(responses):
        for tid in r:
            binary_matrix[id_to_idx[tid], t] = 1

    # 5. Lempel-Ziv complexity of flattened binary pattern
    flat = binary_matrix.flatten()
    lz = _lempel_ziv_complexity(flat)
    # Maximum LZ complexity for a binary sequence of length n
    n = len(flat)
    max_lz = n / np.log2(max(n, 2))

    pci = float(lz / max(max_lz, 1e-10))

    logger.debug(
        "Basin PCI: %.4f (LZ=%d, max_LZ=%.1f, n_ids=%d, window=%d)",
        pci,
        lz,
        max_lz,
        len(id_list),
        response_window,
    )
    return pci
