import numpy as np
import xarray as xr
import os


# ============================================================
# Experiment settings from shell script
# ============================================================
REST = %REST
TSTR = %TSTR
TEND = %TEND
TSTEP = %TSTEP

RDIR = "%RDIR"
RFILE = "%RFILE"
ODIR = "%ODIR"
OUTPUTF = "%OUTPUTF"

DEBUG_OUTPUT = %DEBUG_OUTPUT


# ============================================================
# Global model parameters
# ============================================================
DT = 0.01

SIGMA = 10.0
RHO = 28.0
BETA = 8.0 / 3.0

AMP_X = %AMP_X
AMP_Y = %AMP_Y
AMP_Z = %AMP_Z


# ============================================================
# Internal tendency: Lorenz model
# ============================================================
def internal_tendency(X, Y, Z):

    dXdt = -SIGMA * X + SIGMA * Y
    dYdt = -X * Z + RHO * X - Y
    dZdt = X * Y - BETA * Z

    return dXdt, dYdt, dZdt


# ============================================================
# External forcing tendency
# Students can modify this function later
# ============================================================
def external_forcing_tendency(time):

    # Constant external forcing
    Fx = AMP_X
    Fy = AMP_Y
    Fz = AMP_Z

    # Seasonal-cycle external forcing
    #PERIOD = 12.0
    #OMEGA = 2.0 * np.pi / PERIOD
    #Fx = AMP_X * np.cos(OMEGA * time)
    #Fy = AMP_Y * np.cos(OMEGA * time)
    #Fz = AMP_Z * np.cos(OMEGA * time)

    return Fx, Fy, Fz


# ============================================================
# Time integration
# ============================================================
NT = int(TSTEP / DT)

os.makedirs(ODIR, exist_ok=True)
os.makedirs(RDIR, exist_ok=True)

print(" == Run the Lorenz climate model ==")
print(" DEBUG_OUTPUT =", DEBUG_OUTPUT)

for it in range(TSTR, TEND, TSTEP):

    time = np.float64(it) + DT * np.arange(NT)
    print(" time =", time[0])

    # --------------------------------------------------------
    # Initial condition / restart
    # --------------------------------------------------------
    if REST:
        rin = RFILE.replace("dddd", f"{it:04d}")
        fin = xr.open_dataset(RDIR + rin)

        Xint = fin.Xint
        Yint = fin.Yint
        Zint = fin.Zint

        Xext = fin.Xext
        Yext = fin.Yext
        Zext = fin.Zext

        fin.close()

    else:
        Xint = xr.DataArray([0.0], dims=["time"], coords={"time": [time[0]]})
        Yint = xr.DataArray([10.0], dims=["time"], coords={"time": [time[0]]})
        Zint = xr.DataArray([0.0], dims=["time"], coords={"time": [time[0]]})

        Xext = xr.DataArray([0.0], dims=["time"], coords={"time": [time[0]]})
        Yext = xr.DataArray([0.0], dims=["time"], coords={"time": [time[0]]})
        Zext = xr.DataArray([0.0], dims=["time"], coords={"time": [time[0]]})

    # --------------------------------------------------------
    # Output arrays
    # --------------------------------------------------------
    Xout = xr.DataArray(np.zeros(NT), dims=["time"], coords={"time": time})
    Yout = xr.DataArray(np.zeros(NT), dims=["time"], coords={"time": time})
    Zout = xr.DataArray(np.zeros(NT), dims=["time"], coords={"time": time})

    if DEBUG_OUTPUT:
        Xint_out = xr.DataArray(np.zeros(NT), dims=["time"], coords={"time": time})
        Yint_out = xr.DataArray(np.zeros(NT), dims=["time"], coords={"time": time})
        Zint_out = xr.DataArray(np.zeros(NT), dims=["time"], coords={"time": time})

        Xext_out = xr.DataArray(np.zeros(NT), dims=["time"], coords={"time": time})
        Yext_out = xr.DataArray(np.zeros(NT), dims=["time"], coords={"time": time})
        Zext_out = xr.DataArray(np.zeros(NT), dims=["time"], coords={"time": time})

    # --------------------------------------------------------
    # Model integration
    # --------------------------------------------------------
    for i in range(NT):

        # Save current total climate state
        Xout[i] = Xint[0] + Xext[0]
        Yout[i] = Yint[0] + Yext[0]
        Zout[i] = Zint[0] + Zext[0]

        if DEBUG_OUTPUT:
            Xint_out[i] = Xint[0]
            Yint_out[i] = Yint[0]
            Zint_out[i] = Zint[0]

            Xext_out[i] = Xext[0]
            Yext_out[i] = Yext[0]
            Zext_out[i] = Zext[0]

        # Internal tendency
        dXdt, dYdt, dZdt = internal_tendency(Xint, Yint, Zint)

        # External forcing tendency
        Fx, Fy, Fz = external_forcing_tendency(time[i])

        # Forward Euler update
        Xint = Xint + DT * dXdt
        Yint = Yint + DT * dYdt
        Zint = Zint + DT * dZdt

        Xext = Xext + DT * Fx
        Yext = Yext + DT * Fy
        Zext = Zext + DT * Fz

    etime = time[-1] + DT

    Xint = Xint.assign_coords(time=[etime])
    Yint = Yint.assign_coords(time=[etime])
    Zint = Zint.assign_coords(time=[etime])

    Xext = Xext.assign_coords(time=[etime])
    Yext = Yext.assign_coords(time=[etime])
    Zext = Zext.assign_coords(time=[etime])

    # --------------------------------------------------------
    # Model output
    # --------------------------------------------------------
    fout = OUTPUTF.replace("dddd", f"{it:04d}")

    if DEBUG_OUTPUT:
        dsout = xr.Dataset(
            {
                "X": Xout,
                "Y": Yout,
                "Z": Zout,
                "Xint": Xint_out,
                "Yint": Yint_out,
                "Zint": Zint_out,
                "Xext": Xext_out,
                "Yext": Yext_out,
                "Zext": Zext_out,
            }
        )
    else:
        dsout = xr.Dataset(
            {
                "X": Xout,
                "Y": Yout,
                "Z": Zout,
            }
        )

    dsout.to_netcdf(ODIR + fout)
    dsout.close()

    # --------------------------------------------------------
    # Restart output
    # --------------------------------------------------------
    rout = RFILE.replace("dddd", f"{int(etime):04d}")

    dsrest = xr.Dataset(
        {
            "Xint": Xint,
            "Yint": Yint,
            "Zint": Zint,
            "Xext": Xext,
            "Yext": Yext,
            "Zext": Zext,
        }
    )

    dsrest.to_netcdf(RDIR + rout)
    dsrest.close()

    REST = True