full_loop.py

#!/usr/bin/env sage -python
"""
run_full_loop_nonlinear.py

One-button runner + nonlinear gearbox fitter (interaction & quadratic terms).

- Generates dynamic problem classes
- Runs AdaptiveTuner.tune_class on each
- Collects best complexity_threshold per class
- Fits degree-2 model with interaction: threshold ≈ f(n, bit)
- Validates on held-out classes and prints summary
- Saves gearbox_report_nonlinear.json
"""

import os
import json
import time
import random
import math
from typing import List, Tuple, Dict

import numpy as np

# Adjust this import to your tuner module
try:
    from ai_analyst import AdaptiveTuner, eval_cache
except Exception:
    try:
        from ai_tuner3 import AdaptiveTuner, eval_cache
    except Exception as e:
        raise ImportError("Could not import AI_analyst. Ensure adaptive_ai_tuner_v2.py or ai_tuner3.py is present.") from e

# ----------------------------
# Settings
# ----------------------------
RNG_SEED = 12345
random.seed(RNG_SEED)
np.random.seed(RNG_SEED)

OUTPUT_JSON = "gearbox_report_nonlinear.json"
NUM_DYNAMIC_CLASSES = 12
N_BINS = [(2, 8), (6, 13), (9, 16), (12, 19)]
BIT_BINS = [(6, 12), (12, 16), (16, 20), (20, 24)]

CLASSES_PER_RUN = min(NUM_DYNAMIC_CLASSES, 6)

# Ridge regularization strength for fitting (tweakable)
RIDGE_LAMBDA = 1e-2

# ----------------------------
# Helpers
# ----------------------------
def sample_dynamic_classes(num_classes: int) -> List[dict]:
    combos = [(n,b) for n in N_BINS for b in BIT_BINS]
    random.shuffle(combos)
    selected = combos[:num_classes]
    classes = []
    for i, (nr, br) in enumerate(selected):
        name = f"Dyn_{nr[0]}-{nr[1]}__B{br[0]}-{br[1]}"
        classes.append({"name": name, "n_range": nr, "bit_range": br})
    return classes

def center_of_range(r: Tuple[int,int]) -> float:
    return (r[0] + r[1]) / 2.0

# ----------------------------
# Nonlinear fitter
# ----------------------------
def build_design_matrix(records: List[Dict]) -> Tuple[np.ndarray, np.ndarray, List[str]]:
    """
    Build design matrix X and target y from records list.
    Each record has 'n_center', 'bit_center', 'threshold'.
    Features: [1, n, bit, n*bit, n^2, bit^2]
    """
    X = []
    y = []
    for r in records:
        n = float(r["n_center"])
        b = float(r["bit_center"])
        X.append([1.0, n, b, n*b, n*n, b*b])
        y.append(float(r["threshold"]))
    X = np.array(X, dtype=float)
    y = np.array(y, dtype=float)
    feature_names = ["intercept", "n", "bit", "n*bit", "n^2", "bit^2"]
    return X, y, feature_names

def fit_gearbox_formula_nonlinear(records: List[Dict], ridge_lambda: float = RIDGE_LAMBDA):
    """
    Fit ridge regression to design matrix from records.
    Returns dict with coefficients and a predict function.
    """
    if not records:
        # default trivial formula
        return {"coeffs": [0,1,0,0,0,0], "feature_names": ["intercept","n","bit","n*bit","n^2","bit^2"], "predict": lambda n,b: n}

    X, y, feat_names = build_design_matrix(records)
    # Normal equation with ridge: (X^T X + lambda I) theta = X^T y
    d = X.shape[1]
    A = X.T.dot(X)
    A += ridge_lambda * np.eye(d)
    rhs = X.T.dot(y)
    try:
        theta = np.linalg.solve(A, rhs)
    except np.linalg.LinAlgError:
        theta, *_ = np.linalg.lstsq(A, rhs, rcond=None)
    # prediction function
    def predict(n_val: float, b_val: float) -> float:
        v = np.array([1.0, n_val, b_val, n_val*b_val, n_val*n_val, b_val*b_val], dtype=float)
        return float(np.dot(theta, v))
    return {"coeffs": theta.tolist(), "feature_names": feat_names, "predict": predict}

# Pretty format formula
def formula_to_string(model_dict):
    names = model_dict["feature_names"]
    coeffs = model_dict["coeffs"]
    parts = []
    for name, c in zip(names, coeffs):
        if abs(c) < 1e-8:
            continue
        if name == "intercept":
            parts.append(f"{c:+.4f}")
        else:
            parts.append(f"{c:+.4f}*{name}")
    return " ".join(parts)

# ----------------------------
# Main orchestrator (similar to previous)
# ----------------------------
def run_full_loop_nonlinear(num_dynamic_classes: int = NUM_DYNAMIC_CLASSES):
    print(f"[run_full_loop_nonlinear] Generating {num_dynamic_classes} dynamic classes...")
    dynamic_classes = sample_dynamic_classes(num_dynamic_classes)
    for c in dynamic_classes:
        print(f" - {c['name']}  n={c['n_range']}  bits={c['bit_range']}")

    tuner = AdaptiveTuner(dynamic_classes)
    tuner.load_meta_models()

    per_class_results = []
    start_all = time.time()

    for idx, pclass in enumerate(dynamic_classes):
        print("\n" + "="*70)
        print(f"[{idx+1}/{len(dynamic_classes)}] Tuning class: {pclass['name']}")
        t0 = time.time()
        try:
            tuner.tune_class(pclass)
            best_cfg = tuner.best_configs.get(pclass['name'])
            if best_cfg is None:
                print("  -> No best config found; skipping.")
                continue
            ct = best_cfg.get("complexity_threshold", None)
            v2_knobs = best_cfg.get("v2_knobs", {})
            elapsed = time.time() - t0
            n_center = center_of_range(pclass['n_range'])
            bit_center = center_of_range(pclass['bit_range'])
            rec = {"class_name": pclass["name"], "n_center": n_center, "bit_center": bit_center, "best_threshold": ct, "v2_knobs": v2_knobs, "elapsed_s": elapsed}
            per_class_results.append(rec)
            print(f"  -> Best threshold: {ct}; elapsed: {elapsed:.1f}s")
        except KeyboardInterrupt:
            print("Interrupted by user; saving progress so far.")
            break
        except Exception as e:
            print(f"  -> Error tuning class {pclass['name']}: {e}")
            continue

    # Fit nonlinear gearbox
    fit_records = [{"n_center": r["n_center"], "bit_center": r["bit_center"], "threshold": r["best_threshold"]} for r in per_class_results if r.get("best_threshold") is not None]
    model = fit_gearbox_formula_nonlinear(fit_records, ridge_lambda=RIDGE_LAMBDA)
    model_str = formula_to_string(model)
    print("\n" + "="*70)
    print("[RESULT] Learned nonlinear gearbox formula:")
    print(f"    threshold ≈ {model_str}")
    print("="*70)

    # Validation: hold-out some synthetic classes (not used in fitting)
    holdout = sample_dynamic_classes(max(3, int(len(dynamic_classes)/3)))
    val_results = []
    for pclass in holdout:
        print(f"[validate] Testing holdout class {pclass['name']}")
        # Tune briefly to get ground truth threshold
        try:
            tuner.tune_class(pclass)
            best_cfg = tuner.best_configs.get(pclass['name'])
            if not best_cfg:
                continue
            true_ct = best_cfg.get("complexity_threshold")
            pred_ct = model["predict"](center_of_range(pclass["n_range"]), center_of_range(pclass["bit_range"]))
            val_results.append({"class": pclass["name"], "true_ct": true_ct, "pred_ct": pred_ct})
            print(f"  true={true_ct}, pred={pred_ct:.3f}")
        except Exception as e:
            print(f"  validation error: {e}")

    # Save report
    report = {"generated_classes": dynamic_classes, "per_class_results": per_class_results, "model": {"coeffs": model["coeffs"], "features": model["feature_names"]}, "validation": val_results, "timestamp": time.time()}
    try:
        tmp = OUTPUT_JSON + ".tmp"
        with open(tmp, "w") as f:
            json.dump(report, f, indent=2)
        os.replace(tmp, OUTPUT_JSON)
        print(f"[run_full_loop_nonlinear] Saved report to {OUTPUT_JSON}")
    except Exception as e:
        print(f"[run_full_loop_nonlinear] Failed to write report: {e}")

    elapsed_total = time.time() - start_all
    print(f"[run_full_loop_nonlinear] Total elapsed: {elapsed_total:.1f}s")
    return report

# ----------------------------
# CLI entrypoint
# ----------------------------
if __name__ == "__main__":
    rpt = run_full_loop_nonlinear(NUM_DYNAMIC_CLASSES)
    print("\nCompact summary:")
    for rec in rpt.get("per_class_results", []):
        print(f" - {rec['class_name']}: n_center={rec['n_center']}, bit_center={rec['bit_center']}, threshold={rec['best_threshold']}, time={rec['elapsed_s']:.1f}s")
    print("\nDone.")