import numpy as np
import matplotlib.pyplot as plt
from naginterfaces.library import opt

# --- 1. データ準備 (W_in, W_res も返すように変更) ---
def prepare_data(filename="mackey_glass.txt", n_neurons=50):
    try:
        raw_data = np.loadtxt(filename)
    except OSError:
        print("エラー: データファイルが見つかりません。")
        return None, None, None, None

    # リザバー生成 (シード固定で再現性を確保)
    np.random.seed(42)
    
    # 内部結合 W_res
    W_res = np.random.normal(0, 1, (n_neurons, n_neurons))
    rho = np.max(np.abs(np.linalg.eigvals(W_res)))
    W_res *= 0.9 / rho  # スペクトル半径調整
    
    # 入力結合 W_in
    W_in = np.random.normal(0, 1, (n_neurons, 1))
    
    states = []
    x = np.zeros((n_neurons, 1))
    washout = 50
    
    # 状態生成ループ
    for u in raw_data:
        u_in = np.array([[u]])
        x = np.tanh(np.dot(W_in, u_in) + np.dot(W_res, x))
        states.append(x.flatten())
        
    X_full = np.array(states)
    X_train = X_full[washout:-1]
    y_train = raw_data[washout+1:]
    
    # W_in, W_res も一緒に返す
    return X_train, y_train, W_in, W_res

# --- 2. C言語ヘッダ出力関数 (新規追加) ---
def export_to_c_header(filename, W_in, W_res, w_out, threshold=1e-4):
    """
    マイコン実装用の params.h を生成する
    """
    n_neurons = len(w_out)
    
    # スパースな重みを抽出 (閾値以上のものだけ)
    indices = np.where(np.abs(w_out) > threshold)[0]
    values = w_out[indices]
    
    print(f"ヘッダファイル '{filename}' を生成中...")
    
    with open(filename, "w") as f:
        f.write("#ifndef PARAMS_H\n")
        f.write("#define PARAMS_H\n\n")
        
        f.write(f"// ニューロン数\n")
        f.write(f"#define N_NEURONS {n_neurons}\n\n")
        
        # 1. 入力結合行列 W_in (これは密行列として持つのが一般的)
        f.write("// 入力結合重み (W_in)\n")
        f.write(f"const float W_in[{n_neurons}] = {{\n")
        for i, val in enumerate(W_in.flatten()):
            f.write(f"    {val:.6f}f,")
            if (i+1) % 5 == 0: f.write("\n")
        f.write("};\n\n")
        
        # 2. 内部結合行列 W_res (今回は全結合の密行列として出力)
        # ※マイコンのメモリがきつい場合はここもスパース化するのが定石です
        f.write("// リザバー内部結合重み (W_res)\n")
        f.write(f"const float W_res[{n_neurons} * {n_neurons}] = {{\n")
        for i, val in enumerate(W_res.flatten()):
            f.write(f"    {val:.6f}f,")
            if (i+1) % 5 == 0: f.write("\n")
        f.write("};\n\n")

        # 3. 学習済み出力重み W_out (スパース形式で出力！)
        #    ゼロじゃない重みの「場所」と「値」だけを書き出します
        f.write(f"// --- 学習結果 (L1正則化によるスパース重み) ---\n")
        f.write(f"#define N_ACTIVE_WEIGHTS {len(indices)}\n\n")
        
        f.write("// 有効なニューロンのインデックス\n")
        f.write(f"const int W_out_idx[N_ACTIVE_WEIGHTS] = {{\n    ")
        for idx in indices:
            f.write(f"{idx}, ")
        f.write("\n};\n\n")
        
        f.write("// 有効なニューロンの重み値\n")
        f.write(f"const float W_out_val[N_ACTIVE_WEIGHTS] = {{\n    ")
        for val in values:
            f.write(f"{val:.6f}f, ")
        f.write("\n};\n\n")
        
        f.write("#endif // PARAMS_H\n")
    
    print(f"完了: {len(indices)}個の重みをエクスポートしました。")

# --- 3. コールバック関数 (変更なし) ---
def lsqfun(w, nres, inform, data):
    X = data['X']
    y = data['y']
    pred = np.dot(X, w)
    rx = pred - y
    return rx, inform

def lsqgrd(w, nres, rdx, inform, data):
    X = data['X']
    rdx[:] = X.flatten()
    return inform

# --- メイン処理 ---
def main():
    # データ準備 (W_in, W_res も受け取る)
    n_neurons = 50
    X, y, W_in, W_res = prepare_data(n_neurons=n_neurons)
    if X is None: return

    n_samples = X.shape[0]
    n_weights = X.shape[1]

    print(f"学習開始: サンプル数={n_samples}, 重み数={n_weights}")

    # ソルバー設定
    handle = opt.handle_init(nvar=n_weights)
    opt.handle_set_nlnls(handle, nres=n_samples)

    # ★スパース化のための設定 (強め)
    for option in [
        "NLDF Loss Function Type = L2",
        "Reg Term Type = L1",
        "Reg Coefficient = 0.1",   # 強めの正則化
        "Print Level = 0"
    ]:
        opt.handle_opt_set(handle, option)

    user_data = {'X': X, 'y': y}
    w_init = np.zeros(n_weights)
    
    try:
        results = opt.handle_solve_nldf(
            handle=handle, 
            lsqfun=lsqfun, 
            lsqgrd=lsqgrd,
            x=w_init, 
            nres=n_samples,
            data=user_data
        )
    except Exception as e:
        print(f"ソルバー実行エラー: {e}")
        opt.handle_free(handle)
        return
    
    w_optimized = results.x
    opt.handle_free(handle)

    # 結果表示
    threshold = 1e-4
    n_zeros = np.sum(np.abs(w_optimized) < threshold)
    print("-" * 40)
    print(f"学習完了: ゼロ重み = {n_zeros} / {n_weights}")
    print(f"有効な重み = {n_weights - n_zeros} 個")
    print("-" * 40)

    # --- ★ここでヘッダファイルを出力 ---
    export_to_c_header("params.h", W_in, W_res, w_optimized, threshold)
    
    # グラフ表示 (確認用)
    indices = np.where(np.abs(w_optimized) > threshold)[0]
    plt.figure(figsize=(10, 4))
    plt.stem(w_optimized)
    plt.title(f"Final Weights (Active: {len(indices)}/{n_neurons})")
    plt.xlabel("Neuron Index")
    plt.grid()
    plt.show()

if __name__ == "__main__":
    main()