"""Source-faithful 2023 dual-network calculations and explicit dated adaptation.

The legacy graph is a historical reconstruction, not a valid trading backtest.
The monitoring graph uses one date's correlations and retains the input universe.
Neither graph establishes an economic direction of transmission.
"""
from dataclasses import dataclass
from collections import Counter
from itertools import combinations

import networkx as nx
import numpy as np
import pandas as pd


@dataclass(frozen=True)
class ResearchNetworkConfig:
    window: int = 7
    lag: int = 5
    percentile: float = 90
    seed: int = 7

    def __post_init__(self):
        if self.window < 2 or self.lag < 0 or not 0 <= self.percentile < 100:
            raise ValueError("Require window >= 2, lag >= 0 and 0 <= percentile < 100")


def legacy_pair_chunks(symbols, workers):
    """Reproduce the original omission of all cross-worker symbol pairs."""
    if workers < 1:
        raise ValueError("workers must be positive")
    return [pair for i in range(workers) for pair in combinations(symbols[i::workers], 2)]


def lagged_correlations(frames, config=ResearchNetworkConfig(), legacy_workers=None):
    """Preserve raw Close/Volume levels and one lag orientation per ordered pair.

    Input insertion order matters in the source method. This is deliberately
    exposed rather than silently interpreting the result as a directed network.
    Missing dates are intersected per pair before shift, as in the notebook.
    Complete pairs are the default correction; legacy_workers is audit-only.
    """
    symbols = list(frames)
    if any(not isinstance(s, str) for s in symbols):
        raise ValueError("Symbol identifiers must be strings")
    for frame in frames.values():
        if not frame.index.is_unique or not frame.index.is_monotonic_increasing:
            raise ValueError("Dates must be unique and increasing")
        if not {"Close", "Volume"}.issubset(frame.columns):
            raise ValueError("Both Close and Volume are required")
    pairs = (list(combinations(symbols, 2)) if legacy_workers is None
             else legacy_pair_chunks(symbols, legacy_workers))
    tables = {"Close": [], "Volume": []}
    for left, right in pairs:
        common = frames[left].index.intersection(frames[right].index)
        if common.empty:
            continue
        a, b = frames[left].loc[common], frames[right].loc[common]
        for field in tables:
            shifted = a[field].shift(config.lag).dropna()
            corr = shifted.rolling(config.window).corr(b[field].loc[shifted.index])
            result = corr.to_frame("Correlation_" + field)
            result.index.name = "Date"
            result["Stock1"], result["Stock2"] = left, right
            tables[field].append(result)
    empty = lambda field: pd.DataFrame(
        columns=["Correlation_" + field, "Stock1", "Stock2"],
        index=pd.DatetimeIndex([], name="Date"),
    )
    return {field: pd.concat(parts) if parts else empty(field)
            for field, parts in tables.items()}


def legacy_collapsed_graph(table, field, percentile=90):
    """Replicate percentile over all dates, duplicate-edge overwrite, then cleanup.

    A later qualifying negative observation can overwrite an earlier positive
    edge before cleanup removes it. Removed edges leave their endpoint nodes.
    This intentionally preserves the source behavior for comparison.
    """
    column = "Correlation_" + field
    values = table[column].dropna().to_numpy(dtype=float)
    if not len(values):
        return nx.Graph(), None
    threshold = float(np.percentile(np.abs(values), percentile))
    edges = [(row["Stock1"], row["Stock2"], {"weight": row[column]})
             for row in table.to_dict("records")
             if pd.notna(row[column]) and abs(row[column]) > threshold]
    graph = nx.Graph(edges)
    graph.remove_edges_from([(a, b) for a, b, d in graph.edges(data=True)
                             if np.isinf(d["weight"]) or d["weight"] < 0])
    return graph, threshold


def dated_graph(table, field, as_of, symbols, percentile=90):
    """Monitoring adaptation: a cross-section at one date, never future rows.

    The percentile uses that date's available pairs; it is not a calibrated
    anomaly threshold. All supplied universe nodes, including isolates, remain.
    Complete-pair correlations must be supplied by the caller.
    """
    date = pd.Timestamp(as_of)
    current = table.loc[table.index == date].copy()
    if current.duplicated(["Stock1", "Stock2"]).any():
        raise ValueError("Duplicate pair observations at the signal date")
    column = "Correlation_" + field
    current = current.loc[np.isfinite(current[column].astype(float))]
    graph = nx.Graph()
    graph.add_nodes_from(symbols)
    threshold = None
    if not current.empty:
        threshold = float(np.percentile(np.abs(current[column]), percentile))
        for row in current.to_dict("records"):
            if row["Stock1"] not in graph or row["Stock2"] not in graph:
                raise ValueError("Pair endpoint absent from declared universe")
            value = float(row[column])
            if value > threshold:
                graph.add_edge(row["Stock1"], row["Stock2"], weight=value)
    expected = len(symbols) * (len(symbols) - 1) // 2
    return graph, {"as_of": date.isoformat(), "threshold": threshold,
                   "finite_pairs": len(current), "expected_pairs": expected,
                   "pair_coverage": len(current) / expected if expected else 0.0,
                   "state": "available" if len(current) else "insufficient_data"}


def network_state(graph, seed=7):
    """Keep original python-louvain semantics, now with an explicit seed."""
    from community import community_louvain
    if graph.number_of_edges():
        partition = community_louvain.best_partition(graph, random_state=seed)
    else:
        partition = {symbol: i for i, symbol in enumerate(sorted(graph))}
    return {"partition": partition, "degree_centrality": nx.degree_centrality(graph),
            "nodes": graph.number_of_nodes(), "edges": graph.number_of_edges(),
            "seed": seed}


def legacy_rank(price_state, volume_state):
    """Original lexicographic order: price degree first, volume degree second.

    The community grouping is followed by a global sort in the source, so it
    does not impose community quotas. Stable ties retain the input traversal.
    """
    partition = price_state["partition"]
    price = price_state["degree_centrality"]
    volume = volume_state["degree_centrality"]
    rows = []
    for group in set(partition.values()):
        members = [s for s in partition if partition[s] == group and s in price and s in volume]
        members.sort(key=lambda s: (-price[s], -volume[s]))
        rows.extend((s, price[s], volume[s]) for s in members)
    rows.sort(key=lambda row: (-row[1], -row[2]))
    return rows


def lagged_matrix_snapshot(values, as_of, config=ResearchNetworkConfig()):
    """All lag orientations on a shared trading calendar, using only dates <= cutoff.

    Each finite correlation requires config.window complete observations, as
    pandas rolling(window).corr does. The matrix entry i,j compares lagged i
    with contemporaneous j. It is not a structural causal effect.
    """
    if (not isinstance(values.index, pd.DatetimeIndex) or not values.index.is_unique
            or not values.index.is_monotonic_increasing or not values.columns.is_unique):
        raise ValueError("Require unique symbols and increasing unique trading dates")
    history = values.loc[:pd.Timestamp(as_of)]
    length = config.window + config.lag
    if len(history) < length:
        raise ValueError("Insufficient trailing history")
    block = history.iloc[-length:].to_numpy(dtype=float)
    a = block[:config.window]
    b = block[config.lag:config.lag + config.window]
    good_a = np.isfinite(a).all(axis=0)
    good_b = np.isfinite(b).all(axis=0)
    a = np.where(np.isfinite(a), a, 0.0)
    b = np.where(np.isfinite(b), b, 0.0)
    a = a - a.mean(axis=0)
    b = b - b.mean(axis=0)
    numerator = a.T @ b
    norm = np.sqrt(np.outer(np.sum(a * a, axis=0), np.sum(b * b, axis=0)))
    eligible = np.outer(good_a, good_b) & (norm > 0)
    result = np.full(numerator.shape, np.nan)
    np.divide(numerator, norm, out=result, where=eligible)
    result = np.clip(result, -1, 1)
    frame = pd.DataFrame(result, index=values.columns, columns=values.columns)
    frame.attrs.update({"as_of": history.index[-1].isoformat(), "window": config.window,
                        "lag": config.lag, "raw_levels": True, "shared_calendar": True})
    return frame


def graph_from_matrix(correlations, percentile=90):
    """Retain original i<j lag orientation and positive strong edges for one date.

    Unlike the old notebook, this evaluates all pairs and preserves all nodes.
    The resulting undirected graph deliberately does not claim directionality.
    """
    if (not correlations.index.equals(correlations.columns)
            or not correlations.index.is_unique or not 0 <= percentile < 100):
        raise ValueError("Require a square matrix with identical unique labels and a valid percentile")
    symbols = list(correlations.index)
    i, j = np.triu_indices(len(symbols), 1)
    weights = correlations.to_numpy()[i, j]
    valid = np.isfinite(weights)
    threshold = float(np.percentile(np.abs(weights[valid]), percentile)) if valid.any() else None
    selected = valid & (weights > threshold) if threshold is not None else np.zeros(len(weights), dtype=bool)
    graph = nx.Graph()
    graph.add_nodes_from(symbols)
    graph.add_weighted_edges_from((symbols[a], symbols[b], float(w))
                                 for a, b, w in zip(i[selected], j[selected], weights[selected]))
    return graph, {"threshold": threshold, "finite_pairs": int(valid.sum()),
                   "expected_pairs": len(weights), "pair_coverage": float(valid.mean()) if len(weights) else 0.0,
                   "lag_orientation": "original_input_order_i_before_j", "all_universe_nodes_retained": True}


def partition_disagreement(previous, current):
    """Fraction of common-node pairs whose co-membership changes; label invariant."""
    common = sorted(set(previous) & set(current))
    pairs = len(common) * (len(common) - 1) // 2
    if not pairs:
        return None
    choose2 = lambda counts: sum(v * (v - 1) // 2 for v in counts.values())
    old = choose2(Counter(previous[s] for s in common))
    new = choose2(Counter(current[s] for s in common))
    both = choose2(Counter((previous[s], current[s]) for s in common))
    return (old + new - 2 * both) / pairs