#+------------------------------------------------------------------+
#|                                           train_export_model.py  |
#|                                       https://www.mql5.com       |
#+------------------------------------------------------------------+
# Trains the shared 8-16-8-1 MLP in plain float64 NumPy and exports the
# three files every MQL5 engine in this benchmark reads: model_weights.csv
# (CWeightLoader), test_vectors.csv (the EA's fixed comparison set) and
# model.onnx (the ONNX engine). No PyTorch, no framework-specific weight
# layout to reconcile -- training is full-batch gradient descent, by hand.
import numpy as np
from onnx import helper, TensorProto
from onnx.checker import check_model
from onnx import save

#--- fixed seed so this script's output is reproducible run to run
SEED = 42
rng = np.random.default_rng(SEED)

#--- network dimensions, must match BenchmarkConfig.mqh exactly
INPUT_DIM  = 8
H1         = 16
H2         = 8
OUTPUT_DIM = 1
N_TRAIN    = 2000
N_TEST     = 200
EPOCHS     = 4000
LR         = 0.02


#+------------------------------------------------------------------+
#| relu / relu_grad                                                 |
#+------------------------------------------------------------------+
def relu(z):
    return np.maximum(z, 0.0)


def relu_grad(z):
    return (z > 0.0).astype(np.float64)


#+------------------------------------------------------------------+
#| synthetic_target                                                 |
#| tanh of a linear combination of the inputs, plus a small          |
#| nonlinear cross-term and noise. The exact target is irrelevant   |
#| to the benchmark -- it only has to produce realistic, non-trivial|
#| floats for the network to learn.                                 |
#+------------------------------------------------------------------+
def synthetic_target(X, true_w, noise_std=0.05):
    linear = X @ true_w
    cross = 0.1 * np.sin(X[:, 0] * X[:, 1])
    y = np.tanh(linear) + cross + rng.normal(0.0, noise_std, size=X.shape[0])
    return y.reshape(-1, 1)


#+------------------------------------------------------------------+
#| forward                                                          |
#+------------------------------------------------------------------+
def forward(x, W1, b1, W2, b2, W3, b3):
    h1_pre = x @ W1 + b1
    h1 = relu(h1_pre)
    h2_pre = h1 @ W2 + b2
    h2 = relu(h2_pre)
    out = h2 @ W3 + b3
    return h1_pre, h1, h2_pre, h2, out


#+------------------------------------------------------------------+
#| train                                                            |
#| Full-batch gradient descent, manual backprop, plain MSE loss.    |
#+------------------------------------------------------------------+
def train():
    true_w = rng.normal(0.0, 1.0, size=INPUT_DIM)

    X_train = rng.normal(0.0, 1.0, size=(N_TRAIN, INPUT_DIM))
    y_train = synthetic_target(X_train, true_w)

    #--- He-scaled init, keeps the ReLU layers from dying at step one
    W1 = rng.normal(0.0, np.sqrt(2.0 / INPUT_DIM), size=(INPUT_DIM, H1))
    b1 = np.zeros(H1)
    W2 = rng.normal(0.0, np.sqrt(2.0 / H1), size=(H1, H2))
    b2 = np.zeros(H2)
    W3 = rng.normal(0.0, np.sqrt(2.0 / H2), size=(H2, OUTPUT_DIM))
    b3 = np.zeros(OUTPUT_DIM)

    n = X_train.shape[0]
    for epoch in range(EPOCHS):
        h1_pre, h1, h2_pre, h2, out = forward(X_train, W1, b1, W2, b2, W3, b3)

        #--- dL/dout for mean squared error, then plain chain rule back
        d_out = (out - y_train) / n
        dW3 = h2.T @ d_out
        db3 = d_out.sum(axis=0)

        d_h2 = d_out @ W3.T
        d_h2_pre = d_h2 * relu_grad(h2_pre)
        dW2 = h1.T @ d_h2_pre
        db2 = d_h2_pre.sum(axis=0)

        d_h1 = d_h2_pre @ W2.T
        d_h1_pre = d_h1 * relu_grad(h1_pre)
        dW1 = X_train.T @ d_h1_pre
        db1 = d_h1_pre.sum(axis=0)

        W1 -= LR * dW1; b1 -= LR * db1
        W2 -= LR * dW2; b2 -= LR * db2
        W3 -= LR * dW3; b3 -= LR * db3

        if epoch % 500 == 0 or epoch == EPOCHS - 1:
            loss = 0.5 * np.mean((out - y_train) ** 2)
            print(f"epoch {epoch:5d}  mse {loss:.6f}")

    return W1, b1, W2, b2, W3, b3


#+------------------------------------------------------------------+
#| write_section                                                    |
#| Tagged-section CSV format CWeightLoader reads. 17 significant    |
#| digits round-trips a float64 exactly -- fewer and the benchmark's |
#| error floor would come from the file, not from the engines.      |
#+------------------------------------------------------------------+
def write_section(f, tag, arr, rows, cols):
    f.write(f"{tag},{rows},{cols}\n")
    f.write(",".join(f"{v:.17g}" for v in arr.flatten()) + "\n")


#+------------------------------------------------------------------+
#| const_tensor                                                     |
#| Every initializer cast to float32 -- the deliberate source of    |
#| the fidelity gap this whole benchmark measures.                  |
#+------------------------------------------------------------------+
def const_tensor(name, arr):
    return helper.make_tensor(name, TensorProto.FLOAT, arr.shape, arr.astype(np.float32).flatten())


#+------------------------------------------------------------------+
#| main                                                              |
#+------------------------------------------------------------------+
if __name__ == "__main__":
    W1, b1, W2, b2, W3, b3 = train()

    #--- weights -> model_weights.csv, read by CWeightLoader
    with open("model_weights.csv", "w") as f:
        write_section(f, "#W1", W1, INPUT_DIM, H1)
        write_section(f, "#b1", b1, 1, H1)
        write_section(f, "#W2", W2, H1, H2)
        write_section(f, "#b2", b2, 1, H2)
        write_section(f, "#W3", W3, H2, OUTPUT_DIM)
        write_section(f, "#b3", b3, 1, OUTPUT_DIM)
    print("wrote model_weights.csv")

    #--- 200 fresh test vectors + float64 reference output -> test_vectors.csv
    X_test = rng.normal(0.0, 1.0, size=(N_TEST, INPUT_DIM))
    _, _, _, _, y_test = forward(X_test, W1, b1, W2, b2, W3, b3)

    with open("test_vectors.csv", "w") as f:
        f.write(f"{INPUT_DIM}\n{N_TEST}\n")
        for row, ref in zip(X_test, y_test):
            f.write(",".join(f"{v:.17g}" for v in row) + f",{ref[0]:.17g}\n")
    print("wrote test_vectors.csv")

    #--- same weights, built by hand into an ONNX graph -- no exporter-added
    #--- fusions or layout changes between the weights and the graph
    nodes = [
        helper.make_node("Gemm", ["input", "W1", "b1"], ["h1_pre"], alpha=1.0, beta=1.0),
        helper.make_node("Relu", ["h1_pre"], ["h1"]),
        helper.make_node("Gemm", ["h1", "W2", "b2"], ["h2_pre"], alpha=1.0, beta=1.0),
        helper.make_node("Relu", ["h2_pre"], ["h2"]),
        helper.make_node("Gemm", ["h2", "W3", "b3"], ["output"], alpha=1.0, beta=1.0),
    ]
    initializers = [
        const_tensor("W1", W1), const_tensor("b1", b1),
        const_tensor("W2", W2), const_tensor("b2", b2),
        const_tensor("W3", W3), const_tensor("b3", b3),
    ]
    graph = helper.make_graph(
        nodes, "MathFidelityBenchmarkNet",
        [helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, INPUT_DIM])],
        [helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, OUTPUT_DIM])],
        initializers,
    )
    model = helper.make_model(graph, producer_name="MathFidelityBenchmark",
                               opset_imports=[helper.make_opsetid("", 13)])
    model.ir_version = 8  # newer onnx releases default to an IR version older runtimes may not support
    check_model(model)
    save(model, "model.onnx")
    print("wrote model.onnx")