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

import sys
sys.path.insert(1,"/home1/07064/tg863631/anaconda3/envs/CbrainCustomLayer/lib/python3.6/site-packages") #work around for h5py
from cbrain.imports import *
from cbrain.cam_constants import *
from cbrain.utils import *
from cbrain.layers import *
from cbrain.data_generator import DataGenerator
import tensorflow as tf
from tensorflow import math as tfm
from tensorflow.keras.layers import *
from tensorflow.keras.models import *
# import tensorflow_probability as tfp
import xarray as xr
import numpy as np
from cbrain.model_diagnostics import ModelDiagnostics
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.image as imag
import scipy.integrate as sin
import matplotlib.ticker as mticker
import pickle
from tensorflow.keras import layers
from tensorflow.keras.losses import *
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
import datetime
from cbrain.climate_invariant import *
import yaml

# ## Data Generators

from cbrain.imports import *
from cbrain.utils import *
from cbrain.normalization import *
import h5py
from sklearn.preprocessing import OneHotEncoder


class DataGeneratorClassification(tf.keras.utils.Sequence):
    def __init__(self, data_fn, input_vars, output_vars, percentile_path, data_name,
                 norm_fn=None, input_transform=None, output_transform=None,
                 batch_size=1024, shuffle=True, xarray=False, var_cut_off=None, normalize_flag=True, bin_size=100):
        # Just copy over the attributes
        self.data_fn, self.norm_fn = data_fn, norm_fn
        self.input_vars, self.output_vars = input_vars, output_vars
        self.batch_size, self.shuffle = batch_size, shuffle
        self.bin_size = bin_size
        self.percentile_bins = load_pickle(percentile_path)['Percentile'][data_name]
        self.enc = OneHotEncoder(sparse=False)
        classes = np.arange(self.bin_size+2)
        self.enc.fit(classes.reshape(-1,1))
        # Open datasets
        self.data_ds = xr.open_mfdataset(data_fn)
        if norm_fn is not None: self.norm_ds = xr.open_dataset(norm_fn)
     # Compute number of samples and batches
        self.n_samples = self.data_ds.X.shape[0]
        self.n_batches = int(np.floor(self.n_samples) / self.batch_size)

        self.n_inputs, self.n_outputs = 64, 64
        
                # Initialize input and output normalizers/transformers
        if input_transform is None:
            self.input_transform = Normalizer()
        elif type(input_transform) is tuple:
            ## normalize flag added by Ankitesh
            self.input_transform = InputNormalizer(
                self.norm_ds,normalize_flag, input_vars, input_transform[0], input_transform[1], var_cut_off)
        else:
            self.input_transform = input_transform  # Assume an initialized normalizer is passed
            
            
        if output_transform is None:
            self.output_transform = Normalizer()
        elif type(output_transform) is dict:
            self.output_transform = DictNormalizer(self.norm_ds, output_vars, output_transform)
        else:
            self.output_transform = output_transform  # Assume an initialized normalizer is passed

        # Now close the xarray file and load it as an h5 file instead
        # This significantly speeds up the reading of the data...
        if not xarray:
            self.data_ds.close()
            self.data_ds = h5py.File(data_fn, 'r')
    
    def __len__(self):
        return self.n_batches
     
    def __getitem__(self, index):
        # Compute start and end indices for batch
        start_idx = index * self.batch_size
        end_idx = start_idx + self.batch_size

        # Grab batch from data
        batch_X = self.data_ds['X'][start_idx:end_idx]
        Y = self.data_ds['Y'][start_idx:end_idx]
        # Normalize
        X = self.input_transform.transform(batch_X)
        return X, Y

    def on_epoch_end(self):
        self.indices = np.arange(self.n_batches)
        if self.shuffle: np.random.shuffle(self.indices)

scale_dict = load_pickle('/export/nfs0home/ankitesg/CBrain_project/CBRAIN-CAM/nn_config/scale_dicts/009_Wm2_scaling.pkl')

TRAINFILE = 'CI_SP_M4K_train_shuffle.nc'
TRAIN_FILE_ONEHOT = '/scratch/ankitesh/data/new_data_for_v2_*'
VALIDFILE = 'CI_SP_M4K_valid.nc'
NORMFILE = 'CI_SP_M4K_NORM_norm.nc'
data_path = '/fast/ankitesh/data/'

train_gen = DataGeneratorClassification(
    data_fn=f'{TRAIN_FILE_ONEHOT}', 
    input_vars=['QBP','TBP','PS', 'SOLIN', 'SHFLX', 'LHFLX'], 
    output_vars=['PHQ','TPHYSTND','FSNT', 'FSNS', 'FLNT', 'FLNS'], 
    percentile_path='/export/nfs0home/ankitesg/data/percentile_data.pkl', 
    data_name = 'M4K',
    input_transform = ('mean', 'maxrs'),
    output_transform = scale_dict,
    norm_fn = f'{data_path}{NORMFILE}',
    batch_size=1024,
    bin_size=100
)

train_gen.n_samples

# ## Model (flattened and reshaped)

inp_shape = train_gen[0][0].shape
output_shape = train_gen[0][1].shape

print(f"Input shape {inp_shape} \nOutput Shape {output_shape}")

bin_size = 100

# +
inp = Input(shape=(64,))
densout = Dense(128, activation='linear')(inp)
densout = LeakyReLU(alpha=0.3)(densout)
for i in range (4):
    densout = Dense(128, activation='linear')(densout)
    densout = LeakyReLU(alpha=0.3)(densout)

densout = Dense(256, activation='linear')(densout)
densout = LeakyReLU(alpha=0.3)(densout)
densout = Dense(512, activation='linear')(densout)
densout = LeakyReLU(alpha=0.3)(densout)
densout = Dense(1024, activation='linear')(densout)
densout = LeakyReLU(alpha=0.3)(densout)

densout = Dense(64*(bin_size+2))(densout)
reshaped = tf.keras.layers.Reshape((64,bin_size+2))(densout)
out = tf.keras.layers.Softmax()(reshaped) 
model = tf.keras.models.Model(inp, out)

# -

model.summary()

model.compile(tf.keras.optimizers.Adam(), loss='categorical_crossentropy', metrics=["accuracy"])
path_HDF5 = '/scratch/ankitesh/models/'
earlyStopping = EarlyStopping(monitor='val_loss', patience=10, verbose=0, mode='min')
mcp_save = ModelCheckpoint(path_HDF5+'BF_Classification_bin_size_100_model_v2_fast.hdf5',save_best_only=True, monitor='val_loss', mode='min')

with tf.device('/gpu:1'):
    Nep = 5
    model.fit_generator(train_gen, epochs=Nep,
                      callbacks=[earlyStopping, mcp_save])

# ## Model type 2 (all flattened)

bin_size = 20

inp = Input(shape=(64,))
densout = Dense(128, activation='linear')(inp)
densout = LeakyReLU(alpha=0.3)(densout)
for i in range (4):
    densout = Dense(128, activation='linear')(densout)
    densout = LeakyReLU(alpha=0.3)(densout)
densout = Dense(64*(bin_size+2), activation='sigmoid')(densout)
model = tf.keras.models.Model(inp, densout)

model.summary()

model.compile(tf.keras.optimizers.Adam(), loss='categorical_crossentropy', metrics=["accuracy"])
path_HDF5 = '/scratch/ankitesh/models/'
earlyStopping = EarlyStopping(monitor='val_loss', patience=10, verbose=0, mode='min')
mcp_save = ModelCheckpoint(path_HDF5+'BF_Classification_bin_size_20_flat.hdf5',save_best_only=True, monitor='val_loss', mode='min')

with tf.device('/gpu:1'):
    Nep = 5
    model.fit_generator(train_gen, epochs=Nep, validation_data=valid_gen,\
                      callbacks=[earlyStopping, mcp_save])

# ## Change Preprocessing way

# Pre-dump the preprocessed data and simply load it during run time

# load the data nc file
ds = xr.open_dataset('/scratch/ankitesh/data/CI_SP_M4K_train_shuffle.nc')

ds['time']



# ## Model (mult-output classification)

bin_size = 20


#this defines a single branch out of 64 branches
def define_single_output_branch(densout,out_index):
    out = Dense(bin_size+2, activation='softmax',name=f"output_{out_index}")(densout)
    return out


inp = Input(shape=(64,))
densout = Dense(128, activation='linear')(inp)
densout = LeakyReLU(alpha=0.3)(densout)
for i in range (4):
    densout = Dense(128, activation='linear')(densout)
    densout = LeakyReLU(alpha=0.3)(densout)
densout = Dense(32, activation='linear')(densout)
densout = LeakyReLU(alpha=0.3)(densout)
all_outputs = [define_single_output_branch(densout,i) for i in range(64)]
model = tf.keras.models.Model(inputs=inp, outputs=all_outputs)

model.summary()

losses = {}
for i in range(64):
    losses[f'output_{i}'] = "categorical_crossentropy"

model.compile(tf.keras.optimizers.Adam(), loss=losses, metrics=["accuracy"])
path_HDF5 = '/scratch/ankitesh/models/'
earlyStopping = EarlyStopping(monitor='val_loss', patience=10, verbose=0, mode='min')
mcp_save = ModelCheckpoint(path_HDF5+'BF_Classification_bin_size_20.hdf5',save_best_only=True, monitor='val_loss', mode='min')

with tf.device('/gpu:1'):
    Nep = 5
    model.fit_generator(train_gen, epochs=Nep, validation_data=valid_gen,\
                      callbacks=[earlyStopping, mcp_save])

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())


tf.__version__

data_temp = xr.open_mfdataset("/scratch/ankitesh/data/new_data_for_v2_*")

data_temp.to_netcdf("/scratch/ankitesh/data/train.nc")



StarcoderdataJupyterScriptsDedupFiltered