# 1) Generate synthetic dataset and save to CSV
import numpy as np, pandas as pd
np.random.seed(42)
N = 5000
dVtn = np.random.normal(0, 20, N)   # mV
dVtp = np.random.normal(0, 20, N)   # mV
dL   = np.random.normal(0, 0.8, N)  # nm
Vdd  = np.random.normal(0.95, 0.03, N)  # V
Temp = np.random.normal(25, 30, N)  # degC
base = 200 * (Vdd / 1.0) - 0.4 * (Temp - 25)
snm = base - 0.8 * dVtn + 0.7 * dVtp - 15 * (dL) + np.random.normal(0, 8, N)
data = pd.DataFrame({
    "dVtn": dVtn, "dVtp": dVtp, "dL": dL, "Vdd": Vdd, "Temp": Temp, "SNM": snm
})
data.to_csv("synthetic_sram_variability.csv", index=False)
print("Saved synthetic dataset to synthetic_sram_variability.csv (rows = {})".format(len(data)))
data.head()

# 2) Gradient Boosting surrogate + quantile models (fast, scalable for large data)
import pandas as pd, numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.metrics import mean_absolute_error, r2_score

data = pd.read_csv("synthetic_sram_variability.csv")
X = data[["dVtn","dVtp","dL","Vdd","Temp"]].values
y = data["SNM"].values
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=1)

gbr = GradientBoostingRegressor(n_estimators=300, max_depth=5, random_state=0)
gbr.fit(X_train, y_train)
y_pred = gbr.predict(X_test)
print("GBR MAE (mV):", mean_absolute_error(y_test, y_pred))
print("GBR R2:", r2_score(y_test, y_pred))

# Quantile models for approximate intervals
lower_gbr = GradientBoostingRegressor(loss="quantile", alpha=0.05, n_estimators=300, max_depth=5, random_state=2)
upper_gbr = GradientBoostingRegressor(loss="quantile", alpha=0.95, n_estimators=300, max_depth=5, random_state=3)
lower_gbr.fit(X_train, y_train)
upper_gbr.fit(X_train, y_train)

# Save models with joblib for later reuse (optional)
import joblib
joblib.dump(gbr, "gbr_model.joblib")
joblib.dump(lower_gbr, "gbr_lower_q05.joblib")
joblib.dump(upper_gbr, "gbr_upper_q95.joblib")
print('Saved GBR and quantile models as joblib files.')

# 3) Gaussian Process surrogate (good UQ for small-to-medium datasets)
# NOTE: GPR scales poorly with large N (O(N^3)). Use on subsets or with sparse approximations.
import numpy as np, pandas as pd
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, ConstantKernel as C, WhiteKernel

data = pd.read_csv("synthetic_sram_variability.csv")
# Use a subset for GP demo to keep runtime reasonable
subset = data.sample(n=800, random_state=10).reset_index(drop=True)
X_sub = subset[["dVtn","dVtp","dL","Vdd","Temp"]].values
y_sub = subset["SNM"].values

# Kernel: constant * RBF + white noise
kernel = C(1.0, (1e-3, 1e3)) * RBF([20,20,1.0,0.1,10], (1e-2,1e2)) + WhiteKernel(noise_level=1, noise_level_bounds=(1e-3,1e2))
gpr = GaussianProcessRegressor(kernel=kernel, alpha=0.0, normalize_y=True, n_restarts_optimizer=5, random_state=0)
gpr.fit(X_sub, y_sub)
print("Trained GPR on subset (n={})".format(len(X_sub)))

# Example predictions with uncertainty
X_test_sample = X_sub[:10]
y_mean, y_std = gpr.predict(X_test_sample, return_std=True)
for i, (m,s) in enumerate(zip(y_mean, y_std)):
    print(f"Sample {i}: mean={m:.2f} mV, std={s:.2f} mV")
joblib.dump(gpr, "gpr_model.joblib")
print("Saved GPR model as joblib file.")

# 4) Helper functions: yield curve and PVT grid using a given surrogate model
import numpy as np, pandas as pd

def compute_yield_from_model(model, sigma_scale=1.0, Nsamps=20000, threshold=100.0, seed=0):
    rng = np.random.RandomState(seed)
    dVtn_s = rng.normal(0, 20*sigma_scale, Nsamps)
    dVtp_s = rng.normal(0, 20*sigma_scale, Nsamps)
    dL_s   = rng.normal(0, 0.8*sigma_scale, Nsamps)
    Vdd_s  = np.full(Nsamps, 0.95)
    Temp_s = np.full(Nsamps, 25)
    X_s = np.stack([dVtn_s, dVtp_s, dL_s, Vdd_s, Temp_s], axis=1)
    snm_pred = model.predict(X_s)
    return np.mean(snm_pred > threshold)

def pvt_yield_grid(model, Vdd_grid, Temp_grid, Nsamps=5000, threshold=100.0):
    grid = np.zeros((len(Temp_grid), len(Vdd_grid)))
    rng = np.random.RandomState(1)
    for i,T in enumerate(Temp_grid):
        for j,V in enumerate(Vdd_grid):
            dVtn_s = rng.normal(0, 20, Nsamps)
            dVtp_s = rng.normal(0, 20, Nsamps)
            dL_s   = rng.normal(0, 0.8, Nsamps)
            Vdd_s  = np.full(Nsamps, V)
            Temp_s = np.full(Nsamps, T)
            X_s = np.stack([dVtn_s, dVtp_s, dL_s, Vdd_s, Temp_s], axis=1)
            snm_pred = model.predict(X_s)
            grid[i,j] = np.mean(snm_pred > threshold)
    return grid

# Example quick run using the saved gbr model
import joblib
gbr_loaded = joblib.load("gbr_model.joblib")
print("Yield at sigma=1.0 (example):", compute_yield_from_model(gbr_loaded, sigma_scale=1.0))

# 5) Active learning loop skeleton (identify high-uncertainty regions and run targeted SPICE)
# This is a template ??adapt to your infrastructure that runs SPICE/TCAD jobs.
import numpy as np

def find_high_uncertainty_points(model, candidate_X, top_k=100):
    # For models that provide std (e.g., GPR), use predicted std
    # For models without std, use ensemble disagreement or quantile spread
    if hasattr(model, 'predict') and 'return_std' in model.predict.__code__.co_varnames:
        mu, std = model.predict(candidate_X, return_std=True)
        idx = np.argsort(-std)[:top_k]
        return candidate_X[idx], std[idx]
    else:
        # fallback: use quantile models to estimate spread
        try:
            lower = lower_gbr.predict(candidate_X)
            upper = upper_gbr.predict(candidate_X)
            spread = upper - lower
            idx = np.argsort(-spread)[:top_k]
            return candidate_X[idx], spread[idx]
        except Exception as e:
            print("No UQ available on model:", e)
            return None, None

# Example: create a candidate pool and pick top uncertain points for SPICE runs
rng = np.random.RandomState(2)
candidate_pool = np.column_stack([rng.normal(0,20,20000), rng.normal(0,20,20000), rng.normal(0,0.8,20000), np.full(20000,0.95), np.full(20000,25)])
topX, top_unc = find_high_uncertainty_points(gpr, candidate_pool, top_k=50)
print("Selected top uncertain candidates shape:", None if topX is None else topX.shape)

# 6) HOW TO ADAPT THIS NOTEBOOK TO YOUR REAL DATA (CSV)
# - Your CSV should have columns matching the features used here (dVtn, dVtp, dL, Vdd, Temp) and target SNM
# - Example: 'your_sram_data.csv' with those columns
# - Load your CSV, split to train/test, and replace the training cells above with your data.
#
# Example snippet:
# import pandas as pd
# data = pd.read_csv('your_sram_data.csv')
# X = data[['dVtn','dVtp','dL','Vdd','Temp']].values
# y = data['SNM'].values
# ... follow training steps (GBR/GPR) above ...
#
print('Place your CSV next to this notebook and follow the snippet comments to retrain models on your data.')