#!/usr/bin/env python3
"""
Generate new neurons by sampling in functional parameter space.

Each neuron is a piecewise-linear function fully described by 6 values:
    (boundary_x1, boundary_x2, left_slope, mid_slope, right_slope, y_boundary2)

We extract these from your existing neurons, fit a distribution over them,
sample new combinations, and reconstruct valid W1/b1/W2/b2 for each.
"""

import numpy as np
import torch
from safetensors.torch import load_file, save_file
from pathlib import Path
import json
import argparse
import os

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

NEURON_SOURCE = "multi"                         # "single" | "multi"
SINGLE_FILE   = "test_mlp_hf/model.safetensors"
MULTI_DIR     = "source_llm_neurons"

SINGLE_BOUNDARY_MODE = True  # Generate single-boundary neurons (2 active) instead of double-boundary (3 active)
N_GENERATE    = 500        # generate 500 neurons
OUTPUT_DIR    = "generated_neurons"
RANDOM_SEED   = 42

# Generation strategy:
#   "gaussian"     — fit mean/cov to existing neurons, sample from N(mu, sigma)
#   "interpolate"  — convex combinations of pairs of existing neurons
#   "grid"         — systematic grid over the observed parameter ranges
#   "all"          — produce all three sets
STRATEGY = "all"


# ---------------------------------------------------------------------------
# 1. Load existing neurons
# ---------------------------------------------------------------------------

def load_neurons(source, single_file, multi_dir):
    neurons = []
    if source == "single":
        w = load_file(single_file)
        neurons.append({k: v.float().numpy() for k, v in {
            "W1": w["layer1.weight"],
            "b1": w["layer1.bias"],
            "W2": w["layer2.weight"],
            "b2": w["layer2.bias"],
        }.items()})
    elif source == "multi":
        for f in sorted(Path(multi_dir).glob("neuron_*.safetensors")):
            w = load_file(str(f))
            neurons.append({k: v.float().numpy() for k, v in {
                "W1": w["layer1.weight"],
                "b1": w["layer1.bias"],
                "W2": w["layer2.weight"],
                "b2": w["layer2.bias"],
            }.items()})
    return neurons


# ---------------------------------------------------------------------------
# 2. Extract functional parameters from raw weights
# ---------------------------------------------------------------------------

def weights_to_functional(W1, b1, W2, b2, x_probe_range=(-2.0, 2.0), n_probe=200000):
    xs = np.linspace(x_probe_range[0], x_probe_range[1], n_probe)

    def forward(x_scalar):
        x  = np.array([[x_scalar]], dtype=np.float32)
        h  = np.maximum(0, x @ W1.T + b1)
        y  = h @ W2.T + b2
        return float(y.squeeze())

    ys = np.array([forward(x) for x in xs])

    slopes = np.gradient(ys, xs)
    slope_changes = np.abs(np.gradient(slopes, xs))
    
    peak_window = int(n_probe * 0.1) 
    idx1 = int(np.argmax(slope_changes))
    
    masked_changes = slope_changes.copy()
    l_mask = max(0, idx1 - peak_window)
    r_mask = min(n_probe, idx1 + peak_window)
    masked_changes[l_mask:r_mask] = 0.0
    
    idx2 = int(np.argmax(masked_changes))
    
    if idx1 > idx2:
        idx1, idx2 = idx2, idx1
        
    boundary_x1 = float(xs[idx1])
    boundary_x2 = float(xs[idx2])
    
    margin = int(n_probe * 0.03)
    
    idx_l = max(0, idx1 - margin)
    idx_m1 = min(n_probe - 1, idx1 + margin)
    idx_m2 = max(0, idx2 - margin)
    idx_r = min(n_probe - 1, idx2 + margin)
    
    left_slope = float(np.mean(slopes[:idx_l])) if idx_l > 0 else float(slopes[0])
    
    if idx_m2 > idx_m1:
        mid_slope = float(np.mean(slopes[idx_m1:idx_m2]))
    else:
        mid_slope = float(slopes[(idx1 + idx2) // 2])
        
    right_slope = float(np.mean(slopes[idx_r:])) if idx_r < n_probe - 1 else float(slopes[-1])
    y_boundary2 = float(ys[idx2])

    return {
        "boundary_x1":  boundary_x1,
        "boundary_x2":  boundary_x2,
        "left_slope":   left_slope,
        "mid_slope":    mid_slope,
        "right_slope":  right_slope,
        "y_boundary2":  y_boundary2,
    }


# ---------------------------------------------------------------------------
# 3. Reconstruct weights from functional parameters
# ---------------------------------------------------------------------------

def functional_to_weights(boundary_x1, boundary_x2, left_slope, mid_slope, right_slope, y_boundary2,
                           n_hidden=8):
    if boundary_x1 > boundary_x2:
        boundary_x1, boundary_x2 = boundary_x2, boundary_x1

    W1 = np.zeros((n_hidden, 1), dtype=np.float32)
    b1 = np.zeros(n_hidden,     dtype=np.float32)
    W2 = np.zeros((1, n_hidden), dtype=np.float32)
    b2 = np.zeros(1,             dtype=np.float32)

    # Neuron 0: always active, pure slope carrier
    W1[0, 0] = 1.0
    b1[0]    = 100.0  # Ensures carrier stability during extreme negative activation outliers
    W2[0, 0] = right_slope

    # Neuron 1: active left of boundary_x1
    W1[1, 0]  = -1.0
    b1[1]     = boundary_x1
    W2[0, 1]  = -(left_slope - mid_slope)

    # Neuron 2: active left of boundary_x2
    W1[2, 0]  = -1.0
    b1[2]     = boundary_x2
    W2[0, 2]  = -(mid_slope - right_slope)

    target_y = y_boundary2
    neuron0_out = W2[0, 0] * (W1[0, 0] * boundary_x2 + b1[0])
    b2[0]       = target_y - neuron0_out

    return W1, b1, W2, b2


def functional_to_weights_single(boundary_x, left_slope, right_slope, y_at_boundary,
                                  n_hidden=8):
    """Single-boundary version: only 2 active neurons (carrier + 1 transition)"""
    W1 = np.zeros((n_hidden, 1), dtype=np.float32)
    b1 = np.zeros(n_hidden,     dtype=np.float32)
    W2 = np.zeros((1, n_hidden), dtype=np.float32)
    b2 = np.zeros(1,             dtype=np.float32)

    # Neuron 0: always active, pure slope carrier (carries right_slope)
    W1[0, 0] = 1.0
    b1[0]    = 100.0
    W2[0, 0] = right_slope

    # Neuron 1: active left of boundary_x (adds left_slope - right_slope)
    W1[1, 0]  = -1.0
    b1[1]     = boundary_x
    W2[0, 1]  = -(left_slope - right_slope)

    # Calculate b2 for continuity at boundary
    target_y = y_at_boundary
    neuron0_out = W2[0, 0] * (W1[0, 0] * boundary_x + b1[0])
    b2[0]       = target_y - neuron0_out

    return W1, b1, W2, b2


# ---------------------------------------------------------------------------
# 4. Validate a generated neuron (analytical, not numerical gradient)
# ---------------------------------------------------------------------------

def _mlp_forward(x_scalar, W1, b1, W2, b2):
    x = np.array([[x_scalar]], dtype=np.float32)
    h = np.maximum(0.0, x @ W1.T + b1)
    return float((h @ W2.T + b2).squeeze())


def validate_neuron(W1, b1, W2, b2, params, tol=0.05):
    bx1 = params["boundary_x1"]
    bx2 = params["boundary_x2"]
    
    # Dynamically scale probes so we don't accidentally step over boundaries
    # when random generation places bx1 and bx2 extremely close together.
    dist = max(abs(bx2 - bx1), 1e-6)
    eps = min(1e-3, dist / 10.0)
    gap = min(0.05, dist / 4.0)

    y_at_bx2 = _mlp_forward(bx2, W1, b1, W2, b2)

    slope_left = (_mlp_forward(bx1 - gap,       W1, b1, W2, b2) -
                  _mlp_forward(bx1 - gap - eps,  W1, b1, W2, b2)) / eps

    x_mid = (bx1 + bx2) / 2
    slope_mid = (_mlp_forward(x_mid + eps, W1, b1, W2, b2) -
                 _mlp_forward(x_mid,       W1, b1, W2, b2)) / eps

    slope_right = (_mlp_forward(bx2 + gap + eps,  W1, b1, W2, b2) -
                   _mlp_forward(bx2 + gap,        W1, b1, W2, b2)) / eps

    recovered = {
        "boundary_x1": bx1,
        "boundary_x2": bx2,
        "left_slope":  slope_left,
        "mid_slope":   slope_mid,
        "right_slope": slope_right,
        "y_boundary2": y_at_bx2,
    }

    checks = {
        "left_slope":  abs(slope_left  - params["left_slope"])  < tol,
        "mid_slope":   abs(slope_mid   - params["mid_slope"])   < tol,
        "right_slope": abs(slope_right - params["right_slope"]) < tol,
        "y_boundary2": abs(y_at_bx2    - params["y_boundary2"]) < tol * 5,
    }
    return all(checks.values()), checks, recovered


def validate_neuron_single(W1, b1, W2, b2, params, tol=0.05):
    """Validate single-boundary neuron (only 2 slopes)"""
    bx = params["boundary_x"]
    eps = 1e-3
    gap = 0.05

    y_at_bx = _mlp_forward(bx, W1, b1, W2, b2)

    slope_left = (_mlp_forward(bx - gap,       W1, b1, W2, b2) -
                  _mlp_forward(bx - gap - eps,  W1, b1, W2, b2)) / eps

    slope_right = (_mlp_forward(bx + gap + eps,  W1, b1, W2, b2) -
                   _mlp_forward(bx + gap,        W1, b1, W2, b2)) / eps

    recovered = {
        "boundary_x": bx,
        "left_slope":  slope_left,
        "right_slope": slope_right,
        "y_at_boundary": y_at_bx,
    }

    checks = {
        "left_slope":  abs(slope_left  - params["left_slope"])  < tol,
        "right_slope": abs(slope_right - params["right_slope"]) < tol,
        "y_at_boundary": abs(y_at_bx - params["y_at_boundary"]) < tol * 5,
    }
    return all(checks.values()), checks, recovered


# ---------------------------------------------------------------------------
# 5. Generation strategies
# ---------------------------------------------------------------------------

def strategy_gaussian(functional_params, n, rng):
    mat = np.array([
        [p["boundary_x1"], p["boundary_x2"], p["left_slope"], p["mid_slope"], p["right_slope"], p["y_boundary2"]]
        for p in functional_params
    ])

    mu  = mat.mean(axis=0)
    cov = np.cov(mat.T) if len(mat) > 1 else np.eye(6) * 0.1
    cov += np.eye(6) * 1e-4

    samples = rng.multivariate_normal(mu, cov, size=n)
    return [
        {"boundary_x1": s[0], "boundary_x2": s[1], "left_slope": s[2],
         "mid_slope": s[3], "right_slope": s[4], "y_boundary2": s[5]}
        for s in samples
    ]


def strategy_interpolate(functional_params, n, rng):
    results = []
    fp = functional_params
    for _ in range(n):
        i, j = rng.choice(len(fp), size=2, replace=True)
        t    = rng.uniform(0, 1)
        results.append({
            k: (1 - t) * fp[i][k] + t * fp[j][k]
            for k in fp[i]
        })
    return results


def strategy_grid(functional_params, n, rng):
    def get_range(vals, margin=0.2):
        v_min, v_max = min(vals), max(vals)
        if v_min == v_max:
            # Prevent 0-variance collapse by injecting a spread for single neurons
            offset = abs(v_min) * margin if v_min != 0 else margin
            return v_min - offset, v_max + offset
        return v_min, v_max

    bx1_min, bx1_max = get_range([p["boundary_x1"] for p in functional_params])
    bx2_min, bx2_max = get_range([p["boundary_x2"] for p in functional_params])
    ls_min, ls_max   = get_range([p["left_slope"]  for p in functional_params])
    ms_min, ms_max   = get_range([p["mid_slope"]   for p in functional_params])
    rs_min, rs_max   = get_range([p["right_slope"] for p in functional_params])
    yb_min, yb_max   = get_range([p["y_boundary2"] for p in functional_params])

    side = max(2, int(n ** (1.0/6.0)) + 1)

    grid = []
    for bx1i in np.linspace(bx1_min, bx1_max, side):
        for bx2i in np.linspace(bx2_min, bx2_max, side):
            for lsi in np.linspace(ls_min, ls_max, side):
                for msi in np.linspace(ms_min, ms_max, side):
                    for rsi in np.linspace(rs_min, rs_max, side):
                        for ybi in np.linspace(yb_min, yb_max, side):
                            grid.append({
                                "boundary_x1": bx1i, "boundary_x2": bx2i,
                                "left_slope": lsi, "mid_slope": msi,
                                "right_slope": rsi, "y_boundary2": ybi,
                            })

    rng.shuffle(grid)
    while len(grid) < n:
        grid += grid
    return grid[:n]


# ---------------------------------------------------------------------------
# 6. Main
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    rng = np.random.default_rng(RANDOM_SEED)
    out = Path(OUTPUT_DIR)
    out.mkdir(exist_ok=True)

    print("=" * 60)
    print("Generating new neurons from existing ones (Multi-Boundary)")
    print("=" * 60)

    print("\n[1] Loading existing neurons...")
    neurons = load_neurons(NEURON_SOURCE, SINGLE_FILE, MULTI_DIR)
    print(f"    {len(neurons)} source neuron(s)")

    print("\n[2] Extracting functional parameters...")
    functional_params = []
    for k, n in enumerate(neurons):
        p = weights_to_functional(n["W1"], n["b1"], n["W2"], n["b2"])
        functional_params.append(p)
        print(f"    Neuron {k}: boundary1={p['boundary_x1']:+.4f}  "
              f"boundary2={p['boundary_x2']:+.4f}  "
              f"left_slope={p['left_slope']:+.4f}  "
              f"mid_slope={p['mid_slope']:+.4f}  "
              f"right_slope={p['right_slope']:+.4f}  "
              f"y@boundary2={p['y_boundary2']:+.4f}")

    strategies = (
        ["gaussian", "interpolate", "grid"] if STRATEGY == "all"
        else [STRATEGY]
    )

    total_saved = 0
    summary = {}

    for strat in strategies:
        print(f"\n[3] Generating {N_GENERATE} neurons via '{strat}'...")

        if strat == "gaussian":
            new_params = strategy_gaussian(functional_params, N_GENERATE, rng)
        elif strat == "interpolate":
            new_params = strategy_interpolate(functional_params, N_GENERATE, rng)
        elif strat == "grid":
            new_params = strategy_grid(functional_params, N_GENERATE, rng)
        else:
            raise ValueError(f"Unknown strategy: {strat}")

        strat_dir = out / strat
        strat_dir.mkdir(exist_ok=True)

        n_valid = 0
        for idx, p in enumerate(new_params):
            if SINGLE_BOUNDARY_MODE:
                # Convert double-boundary params to single-boundary
                # Use boundary_x1 as the single boundary, ignore boundary_x2
                # Use left_slope and right_slope, ignore mid_slope
                # Estimate y_at_boundary from y_boundary2
                W1, b1, W2, b2 = functional_to_weights_single(
                    p["boundary_x1"], p["left_slope"], p["right_slope"],
                    p["y_boundary2"],
                )
                # Create single-boundary params for validation
                p_single = {
                    "boundary_x": p["boundary_x1"],
                    "left_slope": p["left_slope"],
                    "right_slope": p["right_slope"],
                    "y_at_boundary": p["y_boundary2"],
                }
                valid, checks, recovered = validate_neuron_single(W1, b1, W2, b2, p_single)
            else:
                W1, b1, W2, b2 = functional_to_weights(
                    p["boundary_x1"], p["boundary_x2"], p["left_slope"],
                    p["mid_slope"], p["right_slope"], p["y_boundary2"],
                )
                valid, checks, recovered = validate_neuron(W1, b1, W2, b2, p)

            if valid:
                save_file(
                    {
                        "layer1.weight": torch.tensor(W1),
                        "layer1.bias":   torch.tensor(b1),
                        "layer2.weight": torch.tensor(W2),
                        "layer2.bias":   torch.tensor(b2),
                    },
                    # Padded to 6 digits (06d) to prevent python alphabetical sorting issues downstream
                    str(strat_dir / f"neuron_{idx:06d}.safetensors"),
                )
                n_valid += 1
            else:
                failed = [k for k, v in checks.items() if not v]
                if idx < 10 or idx % 50000 == 0:
                    print(f"    [skip] neuron_{idx:06d}: failed checks {failed}")

        pct = 100 * n_valid / N_GENERATE
        print(f"    Saved {n_valid}/{N_GENERATE} valid neurons ({pct:.0f}%) to {strat_dir}/")
        summary[strat] = {"generated": N_GENERATE, "valid": n_valid, "path": str(strat_dir)}
        total_saved += n_valid

    meta = {
        "source_neurons": len(neurons),
        "source_functional_params": functional_params,
        "strategies": summary,
        "total_saved": total_saved,
    }
    with open(out / "generation_meta.json", "w") as f:
        json.dump(meta, f, indent=2)

    print(f"\n{'=' * 60}")
    print(f"Total neurons generated: {total_saved}")
    print(f"Metadata saved to {out}/generation_meta.json")
    print(f"\nTo use generated neurons in append_neurons_to_t5.py:")
    print(f"  NEURON_SOURCE = 'multi'")
    print(f"  MULTI_DIR     = '{out}/gaussian'  # or interpolate / grid")
    print(f"{'=' * 60}")