"""Interpretable input routing for the grouped MLP experiment.

Each first-layer unit receives one coherent family of features.  Later layers
remain dense, so the model can learn interactions *between* those family-level
representations without mixing unrelated raw signals in the first transform.
"""
from __future__ import annotations

from dataclasses import dataclass
import hashlib
from typing import Sequence

import numpy as np


@dataclass(frozen=True)
class FeatureFamily:
    """A named set of raw features routed to one or more first-layer units."""

    name: str
    features: frozenset[str]


LOCAL_TECHNICAL = FeatureFamily("local_technical", frozenset({
    "stoch_norm", "macd_pred_norm", "osc_norm", "m3_momentum_norm",
    "m2_tiny_norm", "stoch_div_norm", "vwap_div_norm", "stoch_peak_norm",
    "m3_div_norm", "m2_nooff_norm", "m2_div_norm", "rsi_subtf_norm",
    "bb_pct_b_norm",
}))
ON_CHAIN = FeatureFamily("on_chain", frozenset({
    "totalvol_norm", "newaddr_norm", "sendaddr_norm", "mvrv_zscore_value",
    "mvrv_zscore_cont", "nupl_norm", "sopr_norm", "active1y_norm", "hr_norm",
}))
CANDLESTICKS = FeatureFamily("candlesticks", frozenset({
    "bearish_engulfing_score", "bullish_hammer_score", "bullish_engulfing_score",
    "shooting_star_score",
}))
MACRO = FeatureFamily("macro", frozenset({
    "btc_spx_corr_30", "dxy_roc_norm", "vix_pctrank_inv", "btc_dom_roc_sign",
    "us10y_roc_inv_sign", "spy_above_200ema", "gold_roc_pctrank",
    "fed_net_liq_sign", "gc_position", "us2y_roc_inv_sign", "yield_curve_sign",
}))
RSID = FeatureFamily("rsid_structure", frozenset({
    "rsid_reg_bull_norm", "rsid_reg_bear_norm", "rsid_hid_bull_norm",
    "rsid_hid_bear_norm", "rsid_rt_bull_norm", "rsid_rt_bear_norm",
    "rsid_slow_bull_norm", "rsid_slow_bear_norm", "rsid_delayed_peak_norm",
    "rsid_delayed_dip_norm",
}))
MARKET_FLOW = FeatureFamily("market_flow", frozenset({
    "oi_roc_norm", "usdt_d_norm", "basis_norm", "cvd_norm", "rvol_norm",
}))
CROSS_ASSET = FeatureFamily("cross_asset", frozenset({"btc_gold_norm"}))

# 16 first-layer units: more capacity goes to broad, complementary families.
GROUPED_UNIT_ALLOCATION = (
    LOCAL_TECHNICAL, LOCAL_TECHNICAL, LOCAL_TECHNICAL,
    ON_CHAIN, ON_CHAIN, ON_CHAIN,
    MACRO, MACRO, MACRO,
    CANDLESTICKS,
    RSID, RSID, RSID,
    MARKET_FLOW, MARKET_FLOW,
    CROSS_ASSET,
)

# Hybrid v1 reserves half of the 16 first-layer units for semantic specialists
# and leaves the other half dense, preserving direct cross-feature capacity.
HYBRID_GROUPED_SPECIALIST_ALLOCATION = (
    LOCAL_TECHNICAL, LOCAL_TECHNICAL,
    ON_CHAIN,
    MACRO,
    CANDLESTICKS,
    RSID,
    MARKET_FLOW,
    CROSS_ASSET,
)
HYBRID_GROUPED_DENSE_UNIT_COUNT = 8
HYBRID_GROUPED_HIDDEN_SIZE = len(HYBRID_GROUPED_SPECIALIST_ALLOCATION) + HYBRID_GROUPED_DENSE_UNIT_COUNT
HYBRID_GROUPED_DENSE_ROW_INDICES = tuple(
    range(len(HYBRID_GROUPED_SPECIALIST_ALLOCATION), HYBRID_GROUPED_HIDDEN_SIZE)
)


def grouped_first_layer_mask(feature_cols: Sequence[str], hidden_size: int) -> np.ndarray:
    """Return the allowed ``(hidden_size, n_features)`` first-layer edges.

    The experiment deliberately targets the established 55→16→8→1 shape.
    Failing on a different width prevents a silently arbitrary allocation.
    """
    if hidden_size != len(GROUPED_UNIT_ALLOCATION):
        raise ValueError(
            f"Grouped input structure requires {len(GROUPED_UNIT_ALLOCATION)} first-layer units; "
            f"received {hidden_size}."
        )
    known = frozenset().union(*(family.features for family in GROUPED_UNIT_ALLOCATION))
    unknown = set(feature_cols) - known
    if unknown:
        raise ValueError(f"Grouped input structure has unknown features: {sorted(unknown)}")
    mask = np.array(
        [[feature in family.features for feature in feature_cols]
         for family in GROUPED_UNIT_ALLOCATION],
        dtype=bool,
    )
    if not mask.any(axis=0).all():
        missing = [feature for feature, present in zip(feature_cols, mask.any(axis=0)) if not present]
        raise ValueError(f"Grouped input structure left features unrouted: {missing}")
    return mask


def grouped_structure_metadata(feature_cols: Sequence[str], hidden_size: int) -> dict[str, object]:
    """Serializable provenance saved beside an experimental artifact."""
    grouped_first_layer_mask(feature_cols, hidden_size)
    families = tuple(dict.fromkeys(GROUPED_UNIT_ALLOCATION))
    return {
        "name": "grouped_first_layer_v1",
        "unit_families": [family.name for family in GROUPED_UNIT_ALLOCATION],
        "family_feature_counts": {
            family.name: sum(feature in family.features for feature in feature_cols)
            for family in families
        },
    }


def hybrid_grouped_first_layer_mask(feature_cols: Sequence[str], hidden_size: int) -> np.ndarray:
    """Return the v1 specialist-plus-dense first-layer connectivity mask.

    Rows 0--7 are hard-routed specialists, in the documented family order.
    Rows 8--15 remain fully trainable across every input feature.  Keeping the
    architecture check strict makes the allocation auditable rather than an
    accidental choice for another hidden width.
    """
    if hidden_size != HYBRID_GROUPED_HIDDEN_SIZE:
        raise ValueError(
            f"Hybrid grouped input structure requires {HYBRID_GROUPED_HIDDEN_SIZE} "
            f"first-layer units; received {hidden_size}."
        )
    known = frozenset().union(*(family.features for family in HYBRID_GROUPED_SPECIALIST_ALLOCATION))
    unknown = set(feature_cols) - known
    if unknown:
        raise ValueError(f"Hybrid grouped input structure has unknown features: {sorted(unknown)}")
    specialist_rows = np.array(
        [[feature in family.features for feature in feature_cols]
         for family in HYBRID_GROUPED_SPECIALIST_ALLOCATION],
        dtype=bool,
    )
    dense_rows = np.ones((len(HYBRID_GROUPED_DENSE_ROW_INDICES), len(feature_cols)), dtype=bool)
    return np.vstack((specialist_rows, dense_rows))


def hybrid_grouped_structure_metadata(feature_cols: Sequence[str], hidden_size: int) -> dict[str, object]:
    """Return complete, reproducible provenance for the hybrid v1 mask."""
    mask = hybrid_grouped_first_layer_mask(feature_cols, hidden_size)
    feature_text = "\x1f".join(feature_cols)
    specialist_count = len(HYBRID_GROUPED_SPECIALIST_ALLOCATION)
    return {
        "name": "hybrid_grouped_first_layer_v1",
        "version": 1,
        "specialist_row_families": {
            str(row): family.name
            for row, family in enumerate(HYBRID_GROUPED_SPECIALIST_ALLOCATION)
        },
        "dense_row_indices": list(HYBRID_GROUPED_DENSE_ROW_INDICES),
        "feature_cols_sha256": hashlib.sha256(feature_text.encode()).hexdigest(),
        "active_edges": int(mask.sum()),
        "blocked_edges": int(mask.size - mask.sum()),
        "row_active_edges": [int(count) for count in mask.sum(axis=1)],
        "row_feature_indices": [np.flatnonzero(row).tolist() for row in mask],
        "specialist_row_count": specialist_count,
        "dense_row_count": len(HYBRID_GROUPED_DENSE_ROW_INDICES),
    }


def random_sparse_first_layer_mask(
    feature_cols: Sequence[str], hidden_size: int, mask_seed: int,
) -> np.ndarray:
    """Return a seeded random control mask with the grouped row capacities.

    This deliberately samples features independently per hidden unit.  It is a
    capacity-matched control, not a second semantic feature routing scheme.
    Calling the grouped builder first preserves its architecture and feature
    validation while providing the exact per-row active-edge counts.
    """
    grouped_mask = grouped_first_layer_mask(feature_cols, hidden_size)
    rng = np.random.default_rng(mask_seed)
    mask = np.zeros_like(grouped_mask, dtype=bool)
    for row_index, active_edges in enumerate(grouped_mask.sum(axis=1)):
        selected = rng.choice(len(feature_cols), size=int(active_edges), replace=False)
        mask[row_index, selected] = True
    return mask


def random_sparse_structure_metadata(
    feature_cols: Sequence[str], hidden_size: int, mask_seed: int,
) -> dict[str, object]:
    """Return complete, serializable provenance for a random sparse mask."""
    mask = random_sparse_first_layer_mask(feature_cols, hidden_size, mask_seed)
    feature_text = "\x1f".join(feature_cols)
    return {
        "name": "random_sparse_first_layer_v1",
        "mask_seed": mask_seed,
        "generator": "numpy.default_rng(seed).choice(n_features, row_count, replace=False)",
        "feature_cols_sha256": hashlib.sha256(feature_text.encode()).hexdigest(),
        "active_edges": int(mask.sum()),
        "blocked_edges": int(mask.size - mask.sum()),
        "row_active_edges": [int(count) for count in mask.sum(axis=1)],
        "row_feature_indices": [np.flatnonzero(row).tolist() for row in mask],
    }