import torch
import torch.nn as nn
from joblib import load as joblib_load

# ---- load NN model once (at import time, not every iteration) ----
class MyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.input   = nn.Linear(2, 10)
        self.hidden1 = nn.Linear(10, 10)
        self.hidden2 = nn.Linear(10, 1)

    def forward(self, x):
        x = nn.functional.relu(self.input(x))
        x = nn.functional.relu(self.hidden1(x))
        return self.hidden2(x)

NN_bool = True   # set False to use standard AKN analytic f2

if NN_bool:
    _NN_f2      = torch.load('model-neural-k-omega-f_2.pth', weights_only=False)
    _scaler_yp  = joblib_load('scaler-yplus-k-omega-f_2.bin')
    _scaler_ys  = joblib_load('scaler-ystar-k-omega-f_2.bin')
    _yplus_min, _yplus_max, _ystar_min, _ystar_max, _f2_min, _f2_max = \
        np.loadtxt('min-max-model-f_2.txt')


def calceps(su2d, sp2d, eps2d, gen):
    if iter == 0:
        print('calceps (NN version) called')

    ueps  = (eps2d * viscos) ** 0.25
    ystar = ueps * dist / viscos
    rt    = k2d ** 2 / eps2d / viscos

    # ---- compute f2 ------------------------------------------------
    if NN_bool:
        # friction velocity from south wall (same as rans-k-eps-NN.py)
        ustar = (viscos * u2d[:, 0] / yp2d[:, 0]) ** 0.5   # shape (ni,)

        # yplus and ystar for every cell, clipped to training range
        yplus_2d = np.minimum(yp2d, yp2d[:, -1:] - yp2d) * ustar[:, None] / viscos
        ystar_2d = ystar.copy()

        yplus_2d = np.clip(yplus_2d, _yplus_min, _yplus_max)
        ystar_2d = np.clip(ystar_2d, _ystar_min, _ystar_max)

        # scale and build input matrix
        X = np.zeros((ni * nj, 2))
        X[:, 0] = _scaler_yp.transform(yplus_2d.reshape(-1, 1))[:, 0]
        X[:, 1] = _scaler_ys.transform(ystar_2d.reshape(-1, 1))[:, 0]

        with torch.no_grad():
            preds = _NN_f2(torch.tensor(X, dtype=torch.float32))

        f2 = preds.numpy().reshape(ni, nj)
        f2 = np.clip(f2, _f2_min, _f2_max)
    else:
        f2 = ((1 - np.exp(-ystar / 3.1)) ** 2) * \
             (1. - 0.3 * np.exp(-(rt / 6.5) ** 2))

    # ---- fmu (always analytic) -------------------------------------
    fmu2d = ((1 - np.exp(-ystar / 14)) ** 2) * \
            (1 + 5 / rt ** 0.75 * np.exp(-(rt / 200) ** 2))
    fmu2d = np.minimum(fmu2d, 1)

    # ---- production term -------------------------------------------
    vist  = vis2d - viscos
    su2d  = su2d + c_eps_1 * cmu * fmu2d * gen * k2d * vol

    # ---- dissipation term ------------------------------------------
    sp2d  = sp2d - c_eps_2 * f2 * eps2d * vol / k2d

    # ---- modify su & sp (case-specific hooks) ----------------------
    su2d, sp2d = modify_eps(su2d, sp2d)

    ap2d = aw2d + ae2d + as2d + an2d - sp2d

    # ---- under-relaxation ------------------------------------------
    ap2d = ap2d / urf_eps
    su2d = su2d + (1 - urf_eps) * ap2d * eps2d

    return su2d, sp2d, ap2d, fmu2d