# List of categorical columns.
CAT_COLUMNS = [
 'geo_country', 
 'device_os',
 'last_run_end_reason',
]

"""
Explanation: Data preprocessing
End of explanation
"""


def read_files_into_df(file_list):
    li = []
    for filename in file_list:
        print("reading {}".format(filename))
        df = pd.read_csv(filename, index_col=None, header=0)
        df.sample(frac=1) # Shuffle
        li.append(df)
    return pd.concat(li, axis=0, ignore_index=True)

sample_data = read_files_into_df(training_files)
sample_data

# Go through all data to determine the category space for each column.
category_space_mapping = {}
for c in CAT_COLUMNS:
    category_space_mapping[c] = list(map(lambda v: c+'_'+str(v), list(sample_data[c].astype('category').cat.categories)))

category_space_mapping

def one_hot(df, cols):
  """Returns one-hot encoding of DataFrame df including columns in cols."""
  for col in cols:
    dummies = pd.get_dummies(df[col], prefix=col, drop_first=False)
    dummies = dummies.T.reindex(category_space_mapping[col]).T.fillna(0)
    df = pd.concat([df, dummies], axis=1)
    df = df.drop(col, axis=1)
  return df

# Preprocess a dataframe read from CSV into state action reward triples that can be used in model training.
def pre_process(df):
    # List of all numerical columns
    numerical_columns = list(filter(lambda x: x not in (CAT_COLUMNS + ['is_powerup_clicked', 'presented_powerup']), df.columns.values.tolist()))

    # Drop row if action or reward is na
    df = df[df['presented_powerup'].notna()]
    
    # Find all NAs 
    nans = df.isna().sum()
    nans= nans[nans!=0]
    columns_needs_filling = nans.index.tolist()

    # Determine which columns to fill
    numerical_columns_to_fill = list(set(columns_needs_filling) & set(numerical_columns))
    cat_columns_to_fill = list(set(columns_needs_filling) & set(CAT_COLUMNS))

    # Fill NAs
    df[cat_columns_to_fill] = df[cat_columns_to_fill].fillna(method='bfill').fillna(method='ffill')
    df[numerical_columns_to_fill] = df[numerical_columns_to_fill].fillna(value=0)

    # Verify all NA is filled
    assert df.isna().sum().sum() == 0
    
    # actions for on all data
    actions = df['presented_powerup'].astype('category').cat.codes
    # action space - list of all actions to choose from
    action_space = df['presented_powerup'].astype('category').cat.categories
    # mapping from category code back to string
    actions_mapping = dict(enumerate(df['presented_powerup'].astype('category').cat.categories))
    
    rewards = df['is_powerup_clicked']
    states = df.drop(['presented_powerup', 'is_powerup_clicked'], axis=1)
    
    states = one_hot(states, CAT_COLUMNS)
    
    # Assert all categorical state has been found from previously established mapping
    assert states.isna().sum().sum() == 0
    
    return (states, actions, rewards)

sample_states, sample_actions, sample_rewards = pre_process(sample_data)

# List of all numerical columns
NUM_COLUMNS = list(filter(lambda x: x not in (CAT_COLUMNS + ['is_powerup_clicked', 'presented_powerup']), sample_data.columns.values.tolist()))

sample_states

STATE_SIZE = sample_states.shape[1]
ACTION_SPACE = sample_data['presented_powerup'].astype('category').cat.categories
ACTIONS_MAPPING = dict(enumerate(sample_data['presented_powerup'].astype('category').cat.categories))
ACTION_SPACE_SIZE = len(ACTION_SPACE)

## scikit-learn scaler object to scale the states
scaler = StandardScaler()
scaler.fit(sample_states[NUM_COLUMNS])

def scale_transform(df):
    df[NUM_COLUMNS] = scaler.transform(df[NUM_COLUMNS])

"""
Explanation: Load sample data to initialize scaler and actions mapping
We need to understand the data so we can prepare it for preprocessing.
End of explanation
"""


# Calculate the sample weight based on frequency of reward showing up.
actions_distribution = sample_data['presented_powerup'].value_counts(normalize=True)

# Sample with reward of 1 is 2x the weight of those with reward of 0.
# This is decided based on the distribution of rewards.
# Then the samples with reward of 1 is scaled based on their distribution.
actions_weight = dict((2/actions_distribution) ** 0.5)

# Sample weight is 1 if reward is 0. Otherwise it's weighted by action's distribution frequency
def get_sample_weight(rewards, actions):
    weight_from_action = actions.map(lambda a: actions_weight[ACTIONS_MAPPING[a]])
    weight = weight_from_action.combine(rewards, (lambda w_a, r : 1 if r == 0 else w_a))
    return weight

"""
Explanation: Computing the sample weights
End of explanation
"""


# Softmax to convert model prediction to action index.
def get_action(model_prediction):
    return np.argmax(model_prediction)

class ValidationCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        test_data = read_files_into_df(validation_files)
        test_states, test_actions, test_rewards = pre_process(test_data)
        scale_transform(test_states)
        
        test_actions.rename('action', inplace=True)
        test_data = pd.concat([test_states, test_actions, test_rewards], axis=1)

        # Filter only the ones clicked on
        positive_test_data = test_data[test_data['is_powerup_clicked']==1]
        # Find the count of the action with the smallest sample size
        down_sample_size = positive_test_data.groupby('action').count()['is_powerup_clicked'].min()
        # Downsample
        down_sampled_positive_test_data = positive_test_data.groupby('action').apply(lambda x: x.sample(down_sample_size))

        # Now we have the test dataset ready, run it through the model
        prediction_result = self.model.predict({
            'states': tf.convert_to_tensor(down_sampled_positive_test_data.iloc[:, :-2].values)
        })
        prediction_result = np.array(list(map(get_action, prediction_result)))

        # Construct all result into dataframe
        test_df = pd.DataFrame({'pred':prediction_result, 'real': down_sampled_positive_test_data['action']})
        test_df['is_match'] = test_df['pred'] == test_df['real']

        # Count number of matches
        test_counts = test_df['is_match'].value_counts()

        # Compute benchmark by randomly selecting across 8 actions
        benchmarking_result = test_counts.append(pd.Series([test_df.shape[0] / ACTION_SPACE_SIZE], index=['benchmark']))
        print("Model performance comparing to random: {}".format(benchmarking_result[1]/benchmarking_result[2]))

        # Log the distribution of all actions
        test_action_count = test_df['pred'].value_counts()
        
        log_for_current_epoch = {}
        
        for i in range(ACTION_SPACE_SIZE):
            log_for_current_epoch = {
                **log_for_current_epoch,
                ('action {}'.format(i)): test_action_count.get(i, 0),
            }
        optimizer = self.model.optimizer
        
        # To calculate statistical significance of the improvements,
        # try using the Binomial distribution calculator:
        # https://www.socscistatistics.com/tests/binomial/default2.aspx
        wandb.log({
            **log_for_current_epoch,
            'False prediction': benchmarking_result[0],
            'Correct prediction': benchmarking_result[1],
            'Benchmark': benchmarking_result[2],
                  })
        

"""
Explanation: Custom callback to evaluate model
After each epoch, we will test our model against validation data to see how it's performing relative to the benchmark. We will use random selection of actions as the benchmark.
To test our data, we will first filter for the samples that yield a positive reward, and see if our model can predict the action that generated that positive reward.
Because our dataset might be biased (one action appearing more frequently than others), we need to downsample all actions so they are evenly distributed, otherwise the test result will be biased. For example, if action_1 is appears twice as frequently as the other actions, a model that only predicts action_1 will be "twice as good as the benchmark", where in reality this will not be true.
End of explanation
"""


TRAIN_BATCH_SIZE = 2048

def train_generator():
    batch_size = TRAIN_BATCH_SIZE
    for filename in training_files * 100:
        print("reading {}".format(filename))
        df = pd.read_csv(filename, index_col=None, header=0)
        df.sample(frac=1) # Shuffle
        states, actions, rewards = pre_process(df)
        scale_transform(states)
        i = 0
        while i * batch_size < states.shape[0]:
            s = states[i * batch_size : (i + 1) * batch_size]
            a = actions[i * batch_size : (i + 1) * batch_size]
            r = rewards[i * batch_size : (i + 1) * batch_size]
            yield ({ 'states': tf.convert_to_tensor(s.values), 'actions': tf.convert_to_tensor(a.values)},
                   tf.convert_to_tensor(r.values),
                   tf.convert_to_tensor(get_sample_weight(r, a))
                  )
            i += 1


"""
Explanation: Defining the data generator
End of explanation
"""


class NeuralBanditModel(keras.Model):
    def __init__(self, input_dim, output_dim):
        super(NeuralBanditModel, self).__init__()
        self.nn = tf.keras.Sequential(
            [
                Input(shape=(input_dim,)),
                Dense(256, activation='relu'),
                Dense(512, activation='relu'),
                Dense(512, activation='relu'),
                Dense(256, activation='relu'),
                Dropout(0.2),
                Dense(128, activation='relu'),
                Dense(64, activation='relu'),
                Dropout(0.2),
                Dense(32, activation='relu'),
                Dense(output_dim, activation='relu'),
            ],
            name="neural_greedy",
        )
    
    def call(self, inputs):
        return self.nn(inputs['states'])
    
    # Override train_step to allow custom training logic
    def train_step(self, data):
        # Unpack the data. Its structure depends on your model and
        # on what you pass to `fit()`.
        x, y, sample_weight = data
        
        states = x['states']
        actions = x['actions']
        
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass

            # Start target from prediction result
            target = y_pred.numpy()
            # Find the ones that are affected by the actions and update them accordingly
            target[np.arange(states.shape[0]), actions] = y
            # Convert back to tensor
            target_tensor = tf.convert_to_tensor(target)
            
            # Compute the loss value
            # (the loss function is configured in `compile()`)
            loss = self.compiled_loss(target_tensor, y_pred, regularization_losses=self.losses, sample_weight=sample_weight,)

        # Compute gradients
        gradients = tape.gradient(loss, self.nn.trainable_weights)
        # Update weights
        self.optimizer.apply_gradients(zip(gradients, self.nn.trainable_weights))

        # Update metrics
        self.compiled_metrics.update_state(target_tensor, y_pred)

        # Return a dict mapping metric names to current value
        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        # Unpack the data
        x, y = data
        
        states = x['states']
        actions = x['actions']
        
        y_pred = self(x, training=True)  # Forward pass

        # Start target from prediction result
        target = y_pred.numpy()
        # Find the ones that are affected by the actions and update them accordingly
        target[np.arange(states.shape[0]), actions] = y
        # Convert back to tensor
        target_tensor = tf.convert_to_tensor(target)
        
        self.compiled_loss(target_tensor, y_pred, regularization_losses=self.losses)
        # Update the metrics.
        self.compiled_metrics.update_state(target_tensor, y_pred)
        
        # Return a dict mapping metric names to current value.
        # Note that it will include the loss (tracked in self.metrics).
        return {m.name: m.result() for m in self.metrics}


# TODO: (Optional) Replace with your own project ID on https://wandb.ai
# Alternatively you can remove wandb logging.
wandb.init(project="iap-optimization-codelab")

EPOCHS = 10
STEPS_PER_EPOCH = 256

initial_learning_rate = 0.0002
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate,
    decay_steps=10000,
    decay_rate=0.96,
    staircase=True)

model = NeuralBanditModel(STATE_SIZE, ACTION_SPACE.shape[0])
model.compile(optimizer=Adam(learning_rate = lr_schedule), 
              loss="mse", 
              run_eagerly=True)

wandb.config.update({"n_train": EPOCHS * STEPS_PER_EPOCH * TRAIN_BATCH_SIZE, 
                     "batch_size": TRAIN_BATCH_SIZE,
                     "epochs": EPOCHS,
                     "activation": 'relu',
                     'learning_rate': tf.keras.optimizers.schedules.serialize(lr_schedule),
                     'architecture': list(map(lambda l : l.get_output_at(0).get_shape().as_list(), model.nn.layers))
                    })

print("Training starts")
history = model.fit(x=train_generator(),
                    verbose=1, 
                    epochs=EPOCHS, 
                    steps_per_epoch=STEPS_PER_EPOCH,
                    shuffle=True, 
                    callbacks=[
                        WandbCallback(),
                        ValidationCallback(),
                    ]
                   )

"""
Explanation: Model Definition
Loss computation
The loss computation for our model is slightly different from how it's computed in a classification model.
Say performing action a3 given state s, gives the reward r of 1.0. This can only tells us a3 is a good option, but it doesn't tell us anything about the other actions. Therefore, we can only learn about action a3.
Putting this in the code means only the output node for a3 will have a non-zero loss.

To achieve this with keras, we subclassed keras.Model and wrote a customized train_step to use as a wrapper of our keras.Sequential model, for further details about this, you can read this tutorial here.
End of explanation
"""


def test_model():
    test_data = read_files_into_df(test_files)
    test_states, test_actions, test_rewards = pre_process(test_data)
    scale_transform(test_states)

    test_actions.rename('action', inplace=True)
    test_data = pd.concat([test_states, test_actions, test_rewards], axis=1)

    # Filter only the ones clicked on
    positive_test_data = test_data[test_data['is_powerup_clicked']==1]
    # Find the count of the action with the smallest sample size
    down_sample_size = positive_test_data.groupby('action').count()['is_powerup_clicked'].min()
    # Downsample
    down_sampled_positive_test_data = positive_test_data.groupby('action').apply(lambda x: x.sample(down_sample_size))

    # Now we have the test dataset ready, run it through the model
    prediction_result = model.predict({
        'states': tf.convert_to_tensor(down_sampled_positive_test_data.iloc[:, :-2].values)
    })
    prediction_result = np.array(list(map(get_action, prediction_result)))

    # Construct all result into dataframe
    test_df = pd.DataFrame({'pred':prediction_result, 'real': down_sampled_positive_test_data['action']})
    test_df['is_match'] = test_df['pred'] == test_df['real']

    # Count number of matches
    test_counts = test_df['is_match'].value_counts()

    # Compute benchmark by randomly selecting across 13 actions
    benchmarking_result = test_counts.append(pd.Series([test_df.shape[0] / ACTION_SPACE_SIZE], index=['benchmark']))
    print("Model performance comparing to random: {}".format(benchmarking_result[1]/benchmarking_result[2]))

    # Log the distribution of all actions
    test_action_count = test_df['pred'].value_counts()

    log_for_current_epoch = {}

    for i in range(ACTION_SPACE_SIZE):
        log_for_current_epoch = {
            **log_for_current_epoch,
            ('Test - action {}'.format(i)): test_action_count.get(i, 0),
        }

    wandb.log({
        **log_for_current_epoch,
        'Test - False prediction': benchmarking_result[0],
        'Test - Correct prediction': benchmarking_result[1],
        'Test - Benchmark': benchmarking_result[2],
    })

test_model()

"""
Explanation: Model testing
Use data in the test set to test the model performance
End of explanation
"""


model.save('./model_checkpoint ' + str(datetime.now()))

"""
Explanation: Save model weights
End of explanation
"""


converter = tflite.TFLiteConverter.from_keras_model(model.nn)
tflite_model = converter.convert()

with tf.io.gfile.GFile('iap-optimizer.tflite', 'wb') as f:
  f.write(tflite_model)

"""
Explanation: Convert model to TFLite
End of explanation
"""


import json

data = {}

for i, col in enumerate(NUM_COLUMNS):
    data[col] = {}
    data[col]['type'] = 'numerical'
    data[col]['mean'] = scaler.mean_[i]
    data[col]['std'] = scaler.scale_[i]
    

"""
Explanation: Collecting model metadata
https://www.tensorflow.org/lite/convert/metadata
Serializing numerical columns
End of explanation
"""


for col in CAT_COLUMNS:
    dummies = pd.get_dummies(sample_data[col], drop_first=False)    
    data[col] = {}
    data[col]['type'] = 'categorical'
    data[col]['all_values'] = list(dummies.columns)

data['output_mapping'] = list(ACTIONS_MAPPING.values())

with open('preprocess.json', 'w') as f:
    json.dump(data, f)

print(json.dumps(data, indent=4, sort_keys=True))

"""
Explanation: Serializing categorical columns
End of explanation
"""


from tflite_support import flatbuffers
from tflite_support import metadata as _metadata
from tflite_support import metadata_schema_py_generated as _metadata_fb

# Creates model info.
model_meta = _metadata_fb.ModelMetadataT()
model_meta.name = "IAP optimizer"
model_meta.description = ("Determines the expected reward for each action given the state of a user")
model_meta.version = "v1"

b = flatbuffers.Builder(0)
b.Finish(
    model_meta.Pack(b),
    _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
metadata_buf = b.Output()

populator = _metadata.MetadataPopulator.with_model_file('./iap-optimizer.tflite')
populator.load_associated_files(["./preprocess.json"])
populator.populate()

"""
Explanation: Writing metadata to TFLite
End of explanation
"""


raw_test_data = read_files_into_df(test_files[:1])
test_data = raw_test_data.sample(frac=0.05) 
test_states, test_actions, test_rewards = pre_process(test_data)
scale_transform(test_states)

# Same input to the model to a csv for valiation
test_states[:10].to_csv('model_input_samples.csv')
test_actions.rename('action', inplace=True)
test_data = pd.concat([test_states, test_actions, test_rewards], axis=1)
# Now we have the test dataset ready, run it through the model
prediction_result = model.predict({
    'states': tf.convert_to_tensor(test_data
                                .iloc[:, :-2].values)
})

# Build raw data and output into one CSV for validation
s = raw_test_data[:10]
v = pd.DataFrame(prediction_result[:10], columns=list(ACTIONS_MAPPING.values()))
r = df_concat = pd.concat([s, v], axis=1)
r.to_csv('./integration_validation_samples.csv')

"""
Explanation: Run inference on model
Run inference on 10 rows to see what results we get.
We will run inference again on the TF-lite model on client side for comparison.
End of explanation
"""


