"""
MLP weight comparison across assets and timeframes.

For each input feature, computes the effective input sensitivity (L1 norm of
the column of W1 corresponding to that feature) and shows direction + magnitude
across all trained assets/TFs. This reveals which signals are universal vs
asset-specific.

Usage:
    python3 tools/mlp_compare_weights.py              # all winner artifacts
    python3 tools/mlp_compare_weights.py --tf 4H 8H  # specific timeframes
    python3 tools/mlp_compare_weights.py --top 10     # show top 10 features only
    python3 tools/mlp_compare_weights.py --csv        # emit CSV instead of table
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import sys
from pathlib import Path

import numpy as np

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO))

ALL_ASSETS = ["COINBASE_BTCUSD", "COINBASE_ETHUSD", "BINANCE_SOLUSD", "BINANCE_LINKUSD"]
ALL_TFS    = ["4H", "6H", "8H", "12H", "1D"]
WINNER_DIR = REPO / "results" / "winners"


def winner_artifact(asset: str, tf: str) -> str | None:
    p = WINNER_DIR / f"optimization_winner_strategy_mlp_scores_{asset}_{tf}.csv"
    if not p.exists():
        return None
    rows = list(csv.DictReader(open(p)))
    if not rows:
        return None
    path = rows[0].get("mlp_weights_file", "").strip()
    return path if path and os.path.exists(path) else None


def input_sensitivity(art: dict) -> tuple[list[str], np.ndarray]:
    """
    Returns (feature_cols, sensitivity_vector) where sensitivity[i] is the
    signed effective weight from input i to the output, approximated as:
        sign(mean output gradient) × L2 norm of W1[:, i]

    For a tanh MLP with small weights this is a reasonable first-order proxy.
    The sign is taken from the sum of W1[:, i] weighted by W2[0, :] — i.e.,
    the linear approximation of the output w.r.t. input i.
    """
    layers = art["layers"]
    W1 = np.asarray(layers[0][0], dtype=np.float64)  # (h1, n_in)
    b1 = np.asarray(layers[0][1], dtype=np.float64)
    W2 = np.asarray(layers[1][0], dtype=np.float64)  # (h2, h1)

    if len(layers) == 3:
        W3 = np.asarray(layers[2][0], dtype=np.float64)  # (1, h2)
        # Linear chain: output ≈ W3 @ W2 @ W1 @ x (ignoring tanh nonlinearity)
        combined = (W3 @ W2 @ W1).flatten()  # (n_in,)
    else:
        # 2-layer: output ≈ W2 @ W1 @ x
        combined = (W2 @ W1).flatten()  # (n_in,)

    feature_cols = art.get("feature_cols", [])
    return feature_cols, combined


def main() -> None:
    parser = argparse.ArgumentParser(description="MLP weight comparison across assets/TFs")
    parser.add_argument("--assets", nargs="*", default=ALL_ASSETS, choices=ALL_ASSETS)
    parser.add_argument("--tf", nargs="*", default=ALL_TFS, choices=ALL_TFS)
    parser.add_argument("--top", type=int, default=0, help="Show only top N features by variance")
    parser.add_argument("--csv", action="store_true", help="Emit CSV instead of table")
    args = parser.parse_args()

    # Load all winner artifacts
    entries: list[tuple[str, str, list[str], np.ndarray]] = []
    for asset in args.assets:
        for tf in args.tf:
            path = winner_artifact(asset, tf)
            if not path:
                continue
            try:
                with open(path) as f:
                    art = json.load(f)
                feat_cols, sens = input_sensitivity(art)
                entries.append((asset, tf, feat_cols, sens))
            except Exception as e:
                print(f"  ⚠️  {asset} {tf}: {e}", file=sys.stderr)

    if not entries:
        print("No winner artifacts found.", file=sys.stderr)
        sys.exit(1)

    # Collect union of feature columns
    all_features: list[str] = []
    seen: set[str] = set()
    for _, _, feat_cols, _ in entries:
        for f in feat_cols:
            if f not in seen:
                all_features.append(f)
                seen.add(f)

    # Build matrix: rows = asset/TF combos, cols = features
    labels = [f"{a} {tf}" for a, tf, _, _ in entries]
    mat = np.zeros((len(entries), len(all_features)))
    for i, (_, _, feat_cols, sens) in enumerate(entries):
        for j, feat in enumerate(all_features):
            if feat in feat_cols:
                fidx = feat_cols.index(feat)
                if fidx < len(sens):
                    mat[i, fidx] = sens[fidx]

    # Rank features by variance across combos (most discriminating first)
    variances = mat.var(axis=0)
    if args.top > 0:
        top_idx = np.argsort(-variances)[:args.top]
    else:
        top_idx = np.argsort(-variances)
    features_ordered = [all_features[i] for i in top_idx]

    if args.csv:
        writer = csv.writer(sys.stdout)
        writer.writerow(["feature"] + labels)
        for feat in features_ordered:
            fidx = all_features.index(feat)
            row = [feat] + [f"{mat[i, fidx]:.4f}" for i in range(len(entries))]
            writer.writerow(row)
        return

    # Pretty table
    col_w = max(12, max(len(l) for l in labels) + 2)
    feat_w = max(28, max(len(f) for f in features_ordered) + 2)

    header = f"{'Feature':<{feat_w}}" + "".join(f"{l:>{col_w}}" for l in labels) + f"{'Var':>{col_w}}"
    print(header)
    print("─" * len(header))

    for feat in features_ordered:
        fidx = all_features.index(feat)
        vals = mat[:, fidx]
        var  = variances[fidx]
        # Colour-code: + = bullish weight, - = bearish, ≈0 = near-zero
        def fmt(v: float) -> str:
            if abs(v) < 0.001:
                s = "≈0"
            else:
                sign = "+" if v > 0 else ""
                s = f"{sign}{v:.3f}"
            return f"{s:>{col_w}}"
        row = f"{feat:<{feat_w}}" + "".join(fmt(v) for v in vals) + f"{var:>{col_w}.4f}"
        print(row)

    print(f"\n{len(entries)} combos × {len(all_features)} features")
    print("Sorted by variance (most asset-specific features first)")
    print("Values = linearised input sensitivity (W3·W2·W1 approximation)")


if __name__ == "__main__":
    main()