# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] tags=[]
# # FBSDE
#
# Ji, Shaolin, <NAME>, <NAME>, and <NAME>. “Three Algorithms for Solving High-Dimensional Fully-Coupled FBSDEs through Deep Learning.” ArXiv:1907.05327 [Cs, Math], February 2, 2020. http://arxiv.org/abs/1907.05327.
#
# For standard $d$-dimensional Brownian motion and $d$-dimensional Poisson process with intensity $\lambda$ (in every component) consider an FBSDE:
#
# $$
# \begin{aligned}
# dX_t^i &= \frac 1\alpha \exp\left(-\alpha X_t^i\right) \sum_j Z^{ij}_t dW^j_t + \frac 1\alpha \sum_j \log\left(1+r^{ij}_t \exp\left(-\alpha X^i_t\right)\right) dN^j_t & X_0 = x_0 \\
# dY_t^i &= \frac 12 \exp\left(-\alpha X^i_t\right) \sum_j \left(Z^{ij}_t\right)^2 dt + \sum_j Z^{ij}_t dW^j_t + \sum_j r^{ij}_t dN^j_t & Y^i_T= \exp\left(X^i_T\right)
# \end{aligned}
# $$
#
# One can see by Ito's lemma (Oksendal, p.9, 1.2.8) that
#
# $$
# Y^i_t = \exp(X^i_t)
# $$

# +
# working configuration without jumps: time_horizon = 0.1, n_paths = 2 ** 16, n_timesteps = 4, n_dimensions = 50
# working configuration with jumps: time_horizon = 0.1, n_paths = 2 ** 12, n_timesteps = 4, n_dimensions = 2
# -

import os
from makers.gpu_utils import *
os.environ["CUDA_VISIBLE_DEVICES"] = str(pick_gpu_lowest_memory())

import numpy as np
import tensorflow as tf
from keras.layers import Input, Dense, Lambda, Reshape, concatenate, Layer, BatchNormalization
from keras import Model, initializers
from keras.callbacks import ModelCheckpoint
from keras.metrics import mean_squared_error
import matplotlib.pyplot as plt
from datetime import datetime
from keras.metrics import mse
from keras.optimizers import Adam
from keras.callbacks import Callback
import pandas as pd
from tqdm import tqdm
from itertools import product

print("Num GPUs Available: ", len(tf.config.list_physical_devices("GPU")))

# # Inputs

# model parameters
n_paths = 2 ** 12
n_timesteps = 16
alpha = 0.1
time_horizon = 1.
n_dimensions = 2
n_diffusion_factors = 2
n_jump_factors = 2
intensity = 100.
dt = time_horizon / n_timesteps

# create dirs
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
model_name = f"{timestamp}__{n_dimensions}__{n_paths}__{n_timesteps}__{time_horizon}"
tb_log_dir = "/home/tmp/starokon/tensorboard/" + model_name
output_dir = f"_output/models/{model_name}"
os.makedirs(output_dir, exist_ok=True)


# +
def b(t, x, y, z, r):
    return [0. for _ in range(n_dimensions)]

def s(t, x, y, z, r):
    return [[tf.exp(-alpha * x[i]) * z[i][j] / alpha for j in range(n_diffusion_factors)] for i in range(n_dimensions)]

def v(t, x, y, z, r):
    return [[tf.math.log(1 + tf.math.maximum(r[i][j], 0.) * tf.exp(-alpha * x[i])) / alpha for j in range(n_jump_factors)] for i in range(n_dimensions)]

def f(t, x, y, z, r):
    return [tf.exp(-alpha * x[i]) * tf.reduce_sum(z[i] ** 2) / 2. for i in range(n_dimensions)]

def g(x):
    return [tf.exp(x[i]) for i in range(n_dimensions)]


# -

# # Custom layers and callbacks

class InitialValue(Layer):
    
    def __init__(self, y0, **kwargs):
        super().__init__(**kwargs)
        self.y0 = y0
    
    def call(self, inputs):
        return self.y0



# # Model

# + tags=[]
def build_model(n_dimensions, n_diffusion_factors, n_jump_factors, n_timesteps, time_horizon):
    
    dt = time_horizon / n_timesteps
    
    def dX(t, x, y, z, r, dW, dN):

        def drift(arg):
            x, y, z, r = arg
            return tf.math.multiply(b(t, x, y, z, r), dt)
        a0 = tf.vectorized_map(drift, (x, y, z, r))

        def noise(arg):
            x, y, z, r, dW = arg
            return tf.tensordot(s(t, x, y, z, r), dW, [[1], [0]])
        a1 = tf.vectorized_map(noise, (x, y, z, r, dW))

        def jump(arg):
            x, y, z, r, dN = arg
            return tf.tensordot(v(t, x, y, z, r), dN, [[1], [0]])
        a2 = tf.vectorized_map(jump, (x, y, z, r, dN))

        return a0 + a1 + a2

    def dY(t, x, y, z, r, dW, dN):

        def drift(arg):
            x, y, z, r = arg
            return tf.math.multiply(f(t, x, y, z, r), dt)
        a0 = tf.vectorized_map(drift, (x, y, z, r))

        def noise(arg):
            x, y, z, r, dW = arg
            return tf.tensordot(z, dW, [[1], [0]])
        a1 = tf.vectorized_map(noise, (x, y, z, r, dW))

        def jump(arg):
            x, y, z, r, dN = arg
            return tf.tensordot(r, dN, [[1], [0]])
        a2 = tf.vectorized_map(jump, (x, y, z, r, dN))        

        return a0 + a1 + a2

    @tf.function
    def hx(args):
        i, x, y, z, r, dW, dN = args
        return x + dX(i * dt, x, y, z, r, dW, dN)

    @tf.function
    def hy(args):
        i, x, y, z, r, dW, dN = args
        return y + dY(i * dt, x, y, z, r, dW, dN)
    
    paths = []

    n_hidden_units = n_dimensions + n_diffusion_factors + n_jump_factors + 10

    inputs_x0 = Input(shape=(n_dimensions))
    inputs_dW = Input(shape=(n_timesteps, n_diffusion_factors))
    inputs_dN = Input(shape=(n_timesteps, n_jump_factors))
    
    # variable x0

    y0 = Dense(n_hidden_units, activation='elu', kernel_initializer=initializers.RandomNormal(stddev=1e-1), name='y1_0')(inputs_x0)
    y0 = Dense(n_dimensions, activation='elu', kernel_initializer=initializers.RandomNormal(stddev=1e-1), name='y2_0')(y0)
    
    x = inputs_x0
    y = y0

    # constant x0
    
#     x0 = tf.Variable([[1. for _ in range(n_dimensions)]], trainable=False)
#     y0 = tf.Variable([[0. for _ in range(n_dimensions)]], trainable=True)

#     x = InitialValue(x0, trainable=False, name='x_0')(inputs_dW)
#     y = InitialValue(y0, trainable=True, name='y_0')(inputs_dW)
    
    z = concatenate([x, y])
    z = Dense(n_hidden_units, activation='relu', kernel_initializer=initializers.RandomNormal(stddev=1e-1), name='z1_0')(z)
    z = Dense(n_dimensions * n_diffusion_factors, activation='relu', kernel_initializer=initializers.RandomNormal(stddev=1e-1), name='z2_0')(z)
    # z = BatchNormalization(name='zbn_0')(z)
    z = Reshape((n_dimensions, n_diffusion_factors), name='zr_0')(z)

    r = concatenate([x, y])
    r = Dense(n_hidden_units, activation='relu', kernel_initializer=initializers.RandomNormal(stddev=1e-1), name='r1_0')(r)
    r = Dense(n_dimensions * n_jump_factors, activation='relu', kernel_initializer=initializers.RandomNormal(stddev=1e-1), name='r2_0')(r)
    # r = BatchNormalization(name='rbn_0')(r)
    r = Reshape((n_dimensions, n_jump_factors), name='rr_0')(r)

    paths += [[x, y, z, r]]

    # pre-compile lambda layers

    for i in range(n_timesteps):

        step = InitialValue(tf.Variable(i, dtype=tf.float32, trainable=False))(inputs_dW)

        dW = Lambda(lambda x: x[0][:, tf.cast(x[1], tf.int32)])([inputs_dW, step])
        dN = Lambda(lambda x: x[0][:, tf.cast(x[1], tf.int32)])([inputs_dN, step])

        x, y = (
            Lambda(hx, name=f'x_{i+1}')([step, x, y, z, r, dW, dN]),
            Lambda(hy, name=f'y_{i+1}')([step, x, y, z, r, dW, dN]),
        )

        # we don't train z for the last time step; keep for consistency
        z = concatenate([x, y])
        z = Dense(n_hidden_units, activation='relu', name=f'z1_{i+1}')(z)
        z = Dense(n_dimensions * n_diffusion_factors, activation='relu', name=f'z2_{i+1}')(z)
        z = Reshape((n_dimensions, n_diffusion_factors), name=f'zr_{i+1}')(z)
        # z = BatchNormalization(name=f'zbn_{i+1}')(z)

        # we don't train r for the last time step; keep for consistency
        r = concatenate([x, y])
        r = Dense(n_hidden_units, activation='relu', kernel_initializer=initializers.RandomNormal(stddev=1e-1), name=f'r1_{i+1}')(r)
        r = Dense(n_dimensions * n_jump_factors, activation='relu', kernel_initializer=initializers.RandomNormal(stddev=1e-1), name=f'r2_{i+1}')(r)
        r = Reshape((n_dimensions, n_jump_factors), name=f'rr_{i+1}')(r)
        # r = BatchNormalization(name=f'rbn_{i+1}')(r)

        paths += [[x, y, z, r]]

    outputs_loss = Lambda(lambda r: r[1] - tf.transpose(tf.vectorized_map(g, r[0])))([x, y])
    outputs_paths = tf.stack(
        [tf.stack([p[0] for p in paths[1:]], axis=1), tf.stack([p[1] for p in paths[1:]], axis=1)] + 
        [tf.stack([p[2][:, :, i] for p in paths[1:]], axis=1) for i in range(n_diffusion_factors)] +
        [tf.stack([p[3][:, :, i] for p in paths[1:]], axis=1) for i in range(n_jump_factors)], axis=2)


    model_loss = Model([inputs_x0, inputs_dW, inputs_dN], outputs_loss)

    # (n_sample, n_timestep, x/y/z_k, n_dimension)
    # skips the first time step
    model_paths = Model([inputs_x0, inputs_dW, inputs_dN], outputs_paths)
    
    model_y0 = Model([inputs_x0], y0)

    return model_loss, model_paths, model_y0

# + [markdown] tags=[]
# # Training
# -

np.set_printoptions(edgeitems=11, linewidth=90, formatter=dict(float=lambda x: "%7.5g" % x))

# create model
dt = time_horizon / n_timesteps
model_loss, model_paths, model_y0 = build_model(n_dimensions=n_dimensions, n_diffusion_factors=n_dimensions, n_jump_factors=n_dimensions, n_timesteps=n_timesteps, time_horizon=time_horizon)

# +
# model_loss.summary()
# -

# x0 = tf.constant(np.full((n_paths, n_dimensions), 1.))
x0 = 1. + 0. * tf.random.normal((n_paths, n_dimensions))
dW = tf.sqrt(dt) * tf.random.normal((n_paths, n_timesteps, n_diffusion_factors))
dN = tf.random.poisson((n_paths, n_timesteps), tf.constant(dt * np.array([intensity for _ in range(n_jump_factors)])))
target = tf.zeros((n_paths, n_dimensions))


class Y0Callback(Callback):
    
    def __init__(self, filepath=None):
        super(Y0Callback, self).__init__()
        self.filepath = filepath
        self.y0s = []
    
    def on_epoch_end(self, epoch, logs=None):
        y0 = model_y0(tf.constant([[1. for _ in range(n_dimensions)]]))
        self.y0s += [y0[0]]
        print(f"{y0}\n")
        
    def on_train_end(self, logs=None):
        if self.filepath is not None:
            pd.DataFrame(self.y0s).to_csv(self.filepath)


# +
# define callbacks
callbacks = []

callbacks += [Y0Callback(filepath=os.path.join(output_dir, "y0.csv"))]
callbacks += [ModelCheckpoint(os.path.join(output_dir, "model.h5"), monitor="loss", save_weights_only=True, save_best_only=True, overwrite=True)]
# callbacks += [tf.keras.callbacks.TensorBoard(log_dir=tb_log_dir, histogram_freq=1, profile_batch=1)]
callbacks += [tf.keras.callbacks.TerminateOnNaN()]
callbacks += [tf.keras.callbacks.EarlyStopping(monitor="loss", min_delta=1e-4, patience=10)]
# -

# leave it here to readjust learning rate on the fly
adam = Adam(learning_rate=1e-3) 
model_loss.compile(loss='mse', optimizer=adam)

history = model_loss.fit([x0, dW, dN], target, batch_size=128, initial_epoch=0, epochs=1000, callbacks=callbacks)
df_loss = pd.DataFrame(history.history['loss'])
df_loss.to_csv(os.path.join(output_dir, 'loss.csv'))

# + [markdown] tags=[]
# # Transfer weights (if needed)

# +
# try transfer learning from another starting point

model_loss.get_layer('y_0').set_weights(m_large.get_layer('y_0').get_weights())

for i in range(n_timesteps):
    model_loss.get_layer(f'z1_{i}').set_weights(m_large.get_layer(f'z1_{i}').get_weights())
    model_loss.get_layer(f'z2_{i}').set_weights(m_large.get_layer(f'z2_{i}').get_weights())

# + tags=[]
# transfer learning from cruder discretization

model_loss.get_layer('y_0').set_weights(m_small.get_layer('y_0').get_weights())

n_small = 4

for i in range(n_small):
    for j in range(n_timesteps // n_small):
        model_loss.get_layer(f'z1_{n_timesteps // n_small * i}').set_weights(m_small.get_layer(f'z1_{i}').get_weights())
        model_loss.get_layer(f'z2_{n_timesteps // n_small * i}').set_weights(m_small.get_layer(f'z2_{i}').get_weights())
        model_loss.get_layer(f'z1_{n_timesteps // n_small * i + j}').set_weights(m_small.get_layer(f'z1_{i}').get_weights())
        model_loss.get_layer(f'z2_{n_timesteps // n_small * i + j}').set_weights(m_small.get_layer(f'z2_{i}').get_weights())

# +
# check for exploding gradients before training

with tf.GradientTape() as tape:
    loss = mse(model_loss([dW, dN]), target)

# bias of the last dense layer
variables = model_loss.variables[-1]
tape.gradient(loss, variables)

# + [markdown] tags=[]
# # Training loop
# -

for time_horizon, n_paths, n_timesteps, n_dimensions in tqdm(list(product([0.1], [2**12], [4], [2]))):
    
    # create model
    dt = time_horizon / n_timesteps
    model_loss, model_paths = build_model(n_dimensions=n_dimensions, n_diffusion_factors=n_dimensions, n_jump_factors=n_dimensions, n_timesteps=n_timesteps, time_horizon=time_horizon)
    
    # generate paths
    dW = tf.sqrt(dt) * tf.random.normal((n_paths, n_timesteps, n_diffusion_factors))
    dN = tf.random.poisson((n_paths, n_timesteps), tf.constant(dt * np.array([intensity for _ in range(n_jump_factors)]).transpose().reshape(-1)))
    target = tf.zeros((n_paths, n_dimensions))
    
    # create dirs
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    model_name = f"{timestamp}__{n_dimensions}__{n_paths}__{n_timesteps}__{time_horizon}"
    tb_log_dir = "/home/tmp/starokon/tensorboard/" + model_name
    output_dir = f"_output/models/{model_name}"
    os.makedirs(output_dir, exist_ok=True)
    
    # define callbacks
    y0_callback = Y0Callback(filepath=os.path.join(output_dir, "y0.csv"))
    checkpoint_callback = ModelCheckpoint(os.path.join(output_dir, "model.h5"), monitor="loss", save_weights_only=True, save_best_only=True, overwrite=True)
    tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=tb_log_dir, histogram_freq=1, profile_batch=1)
    nan_callback = tf.keras.callbacks.TerminateOnNaN()
    loss_callback = tf.keras.callbacks.EarlyStopping(monitor="loss", min_delta=1e-4, patience=20)
    
    # callbacks
    callbacks = [
        nan_callback,
        checkpoint_callback,
        # tensorboard_callback,
        y0_callback,
        loss_callback,
    ]
    
    # leave it here to readjust learning rate on the fly
    adam = Adam(learning_rate=1e-3) 
    model_loss.compile(loss='mse', optimizer=adam)
    
    # train
    history = model_loss.fit([dW, dN], target, batch_size=128, initial_epoch=0, epochs=3000, callbacks=callbacks)
    df_loss = pd.DataFrame(history.history['loss'])
    df_loss.to_csv(os.path.join(output_dir, 'loss.csv'))

# # Validate

# +
dW_test = tf.sqrt(dt) * tf.random.normal((n_paths//8, n_timesteps, n_diffusion_factors))
dN_test = tf.random.poisson((n_paths//8, n_timesteps), tf.constant(dt * np.array([1., 1.]).transpose().reshape(-1)))
target_test = tf.zeros((n_paths//8, n_dimensions))

model_loss.evaluate([dW_test, dN_test], target_test)
# -

# # Display paths and loss

# load bad model
model_loss.load_weights('_models/weights0011.h5')

loss = model_loss([dW, dN]).numpy()
loss

paths = model_paths([dW, dN]).numpy()

StarcoderdataJupyterScriptsDedupFiltered