"""Measurement engine — entropy, replicator detection, classification."""

import zlib
from collections import Counter
from dataclasses import dataclass, field
from enum import Enum

import numpy as np

from .interpreter import execute_traced


# ---------------------------------------------------------------------------
# Compression / entropy
# ---------------------------------------------------------------------------

def compression_ratio(soup: np.ndarray) -> float:
    """Ratio of compressed size to raw size (lower = more structure)."""
    raw = soup.tobytes()
    compressed = zlib.compress(raw, level=6)
    return len(compressed) / len(raw)


def shannon_entropy(soup: np.ndarray) -> float:
    """Shannon entropy H of the empirical distribution over distinct tapes."""
    hist = population_histogram(soup)
    n = soup.shape[0]
    probs = np.array(list(hist.values())) / n
    # Filter out zero probabilities to avoid log(0)
    probs = probs[probs > 0]
    if len(probs) == 0:
        return 0.0
    return float(-np.sum(probs * np.log2(probs)))


# ---------------------------------------------------------------------------
# Population histogram
# ---------------------------------------------------------------------------

def population_histogram(soup: np.ndarray) -> dict[bytes, int]:
    """Count identical tapes in the soup."""
    counts: dict[bytes, int] = Counter()
    for i in range(soup.shape[0]):
        counts[soup[i].tobytes()] += 1
    return dict(counts)


def top_sequences(soup: np.ndarray, n: int = 20) -> list[tuple[bytes, int]]:
    """Return the N most common tapes."""
    hist = population_histogram(soup)
    return sorted(hist.items(), key=lambda x: -x[1])[:n]


# ---------------------------------------------------------------------------
# Replicator detection via rolling hash
# ---------------------------------------------------------------------------

def find_repeated_subsequences(soup: np.ndarray, min_length: int = 4,
                                min_count: int = 3) -> dict[bytes, list[tuple[int, int]]]:
    """Find byte subsequences appearing on multiple tapes.

    Returns dict mapping subsequence → list of (tape_idx, start_pos).
    Only returns subsequences appearing >= min_count times.
    """
    n_tapes, tape_len = soup.shape
    occurrences: dict[bytes, list[tuple[int, int]]] = {}

    # Check multiple window sizes, starting from min_length
    for win_len in range(min_length, tape_len // 2 + 1):
        window_counts: dict[bytes, list[tuple[int, int]]] = {}
        for t in range(n_tapes):
            tape_bytes = soup[t].tobytes()
            for start in range(tape_len - win_len + 1):
                window = tape_bytes[start:start + win_len]
                if window not in window_counts:
                    window_counts[window] = []
                window_counts[window].append((t, start))

        for seq, locs in window_counts.items():
            # Count distinct tapes
            distinct_tapes = len(set(loc[0] for loc in locs))
            if distinct_tapes >= min_count:
                occurrences[seq] = locs

    # Remove subsequences that are substrings of longer found sequences
    keys = sorted(occurrences.keys(), key=len, reverse=True)
    filtered: dict[bytes, list[tuple[int, int]]] = {}
    for seq in keys:
        # Keep if no longer sequence in filtered contains this one
        is_sub = False
        for longer in filtered:
            if seq in longer:
                is_sub = True
                break
        if not is_sub:
            filtered[seq] = occurrences[seq]

    return filtered


# ---------------------------------------------------------------------------
# Replicator classification
# ---------------------------------------------------------------------------

class ReplicatorType(Enum):
    INANIMATE = "inanimate"  # executed ∩ copied = ∅
    VIRAL = "viral"          # executed ∩ copied ≠ ∅, copied ⊄ executed
    CELLULAR = "cellular"    # copied ⊆ executed (full autopoiesis)


@dataclass(slots=True)
class Replicator:
    sequence: bytes
    length: int
    count: int             # number of distinct tapes it appears on
    classification: ReplicatorType
    locations: list[tuple[int, int]] = field(default_factory=list)


def classify_replicator(sequence: bytes, tape_len: int = 128) -> ReplicatorType:
    """Classify a replicator by running it in isolation and checking overlap.

    The sequence is placed at the start of a tape of `tape_len` bytes,
    rest zeroed. We track which positions are executed vs written.
    """
    tape = bytearray(tape_len)
    tape[:len(sequence)] = sequence

    _, _, executed, written = execute_traced(
        tape, max_steps=10_000,
        dp_offset=len(sequence),          # DP starts just past the code
        wp_offset=len(sequence) + tape_len // 4  # WP further out
    )

    if not written:
        return ReplicatorType.INANIMATE

    overlap = executed & written
    if not overlap:
        return ReplicatorType.INANIMATE
    elif written <= executed:
        return ReplicatorType.CELLULAR
    else:
        return ReplicatorType.VIRAL


def find_replicators(soup: np.ndarray, min_length: int = 4,
                     min_count: int = 3) -> list[Replicator]:
    """Find and classify replicators in the soup."""
    subsequences = find_repeated_subsequences(soup, min_length, min_count)

    replicators = []
    for seq, locs in subsequences.items():
        distinct_tapes = len(set(loc[0] for loc in locs))
        classification = classify_replicator(seq)
        replicators.append(Replicator(
            sequence=seq,
            length=len(seq),
            count=distinct_tapes,
            classification=classification,
            locations=locs,
        ))

    # Sort by count descending, then length descending
    replicators.sort(key=lambda r: (-r.count, -r.length))
    return replicators


# ---------------------------------------------------------------------------
# Depth distribution
# ---------------------------------------------------------------------------

def depth_histogram(depths: np.ndarray) -> dict[int, int]:
    """Count tapes at each depth level."""
    unique, counts = np.unique(depths, return_counts=True)
    return {int(d): int(c) for d, c in zip(unique, counts)}
