# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Environment (conda_pytorch_p36)
#     language: python
#     name: conda_pytorch_p36
# ---

# +
from spotlight.cross_validation import random_train_test_split
from spotlight.datasets.movielens import get_movielens_dataset
from spotlight.evaluation import mrr_score
from spotlight.factorization.implicit import ImplicitFactorizationModel
from spotlight.interactions import Interactions
import pandas as pd
from spotlight.evaluation import *
from torch.utils.tensorboard import SummaryWriter
import torch
from torch.utils.tensorboard import SummaryWriter
import logging
from spotlight.factorization.representations import *
import os
import pandas as pd
import numpy as np
from spotlight.cross_validation import random_train_test_split
from spotlight.datasets.movielens import get_movielens_dataset
from spotlight.evaluation import mrr_score
from spotlight.factorization.implicit import ImplicitFactorizationModel
from spotlight.interactions import Interactions
import pandas as pd
from spotlight.evaluation import *
from torch.utils.tensorboard import SummaryWriter
import torch
from torch.utils.tensorboard import SummaryWriter
import logging
from spotlight.factorization.representations import *
import os
import collections
from spotlight.sequence.implicit import ImplicitSequenceModel
from spotlight.sequence.representations import CNNNet, PoolNet

input_config = os.environ

hyperparameters = {
    'loss': "adaptive_hinge",
    'batch': 128,
    'lr': 1e-3,
    'l2': 1e-06,
    'n_iter': 50,
    'emb_dim': 128,
    'type': 'pool',
    "layers": [int(x) for x in "256-128-64-1".split("-")],
    "nonlin": 'tanh',
    "max_seq": 50,
    "min_seq": 2,
    "step": 2,
    "fmax_seq": 50
}

h = hyperparameters

if 'SUFFIX' in input_config:
    suffix = input_config['SUFFIX']
else:
    suffix = "1"

if 'LOSS' in input_config:
    h['loss'] = input_config['LOSS']

if 'LR' in input_config:
    h['lr'] = float(input_config['LR'])

if 'L2' in input_config:
    h['l2'] = float(input_config['L2'])

if 'MOM' in input_config:
    h['mom'] = float(input_config['MOM'])
else:
    h['mom'] = 0.9

if 'NEGSAMPLES' in input_config:
    h['neg'] = int(input_config['NEGSAMPLES'])
else:
    h['neg'] = 5

if 'BATCH' in input_config:
    h['batch'] = int(input_config['BATCH'])

h['amsgrad'] = False
if 'AMSGRAD' in input_config:
    h['amsgrad'] = (input_config['AMSGRAD'] == 'True')

h['adamw'] = False
if 'ADAMW' in input_config and input_config['ADAMW']:
    h['adamw'] = (input_config['ADAMW'] == 'True')

if 'EMBDIM' in input_config:
    h['emb_dim'] = int(input_config['EMBDIM'])

if "LAYERS" in input_config:
    h["layers"] = [int(x) for x in input_config['LAYERS'].split("-")]

if "NONLIN" in input_config:
    h["nonlin"] = input_config['NONLIN']

if "MAXSEQ" in input_config:
    h["max_seq"] = int(input_config["MAXSEQ"])

if "MINSEQ" in input_config:
    h["min_seq"] = int(input_config["MINSEQ"])

if "STEPSIZE" in input_config:
    h["step"] = int(input_config["STEPSIZE"])

if "FILTERSEQ" in input_config:
    h["fmax_seq"] = int(input_config["FILTERSEQ"])

betas = (h['mom'], 0.999)
use_cuda = True

tensorboard_base_dir = "sruns"
model_store_dir = "smodels"
n_iters = 50
# loss="adaptive_hinge"

log_loss_interval = 1000
log_eval_interval = 20000

model_alias = ",".join([k + "=" + str(v) for k, v in collections.OrderedDict(h).items()])
model_alias = "pool_" + suffix + "_" + model_alias

# train_data_path = "s3a://tubi-playground-production/smistry/emb3/train-aug-28-phase1"
train_data_path = "/home/ec2-user/emb3/data/train-aug-28-phase" + suffix

logging.basicConfig(filename="slogs/" + model_alias + '.log',
                    filemode='w',
                    format='%(asctime)s - %(message)s',
                    level=logging.INFO)

logger = logging.getLogger()
logger.setLevel(logging.DEBUG)

max_sequence_length = h["max_seq"]
min_sequence_length = h["min_seq"]
step_size = h["step"]


original_train_data = pd.read_parquet(train_data_path)
logger.info("Data is downloaded")
# train_data = original_train_data
uvs = original_train_data.groupby("uindex")["vindex"].agg(list)
train_data = original_train_data[original_train_data.uindex.isin(uvs[uvs.apply(lambda x: len(x)) <= h["fmax_seq"]].index)]

logger.info("Filtered train data..")

train_data["vindex"] = train_data["vindex"] + 1
num_items = len(original_train_data["vindex"].unique()) + 2
interactions = Interactions(train_data["uindex"].to_numpy(),
                            train_data["vindex"].to_numpy(),
                            train_data["pct_cvt"].to_numpy(),
                            train_data["latest_watch_time"].to_numpy(),
                            num_users=len(original_train_data["uindex"].unique()),
                            num_items=num_items)


# if "1500K" in suffix:
#     logger.info("Increasing step size and max_sequence_length")
#     step_size = 2
#     min_sequence_length = 2
#     max_sequence_length = 50

train_seq = interactions.to_sequence(max_sequence_length=max_sequence_length,
                                     min_sequence_length=min_sequence_length,
                                     step_size=step_size)

logger.info("Data is loaded and converted to sequences..")

writer = SummaryWriter(log_dir='{}/{}'.format(tensorboard_base_dir, model_alias))
writer.add_text('alias', model_alias, 0)
writer.add_text('hyperparameters', str(h), 0)


def notify_loss_completion(epoch_id, batch_id, loss, net, model):
    # print("notify_loss_completion")
    writer.add_scalar("Batch/loss", loss, batch_id)
    logging.info('[Epoch {}] Batch {}, Loss {}'.format(epoch_id, batch_id, loss))


def notify_batch_eval_completion(epoch_id, batch_id, loss, net, model):
    # print("notify_batch_eval_completion")
    m = 1


def notify_epoch_completion(epoch_num, total_loss, net, model):
    # print("notify_epoch_completion")
    writer.add_scalar("Epoch/loss", total_loss, epoch_num)
    pairs_ndcg = nn_pairs_ndcg_score(net)
    writer.add_scalar("Epoch/pairs_ndcg", pairs_ndcg, epoch_num)
    #     hit_ratio, ndcg = evaluate_hit_ratio_and_ndcg(model)
    #     writer.add_scalar("Epoch/HR", hit_ratio, epoch_num)
    #     writer.add_scalar("Epoch/NDCG", ndcg, epoch_num)
    hit_ratio, ndcg = -1, -1
    logging.info('******** [Epoch {}]  Embs NDCG {:.4f}, Hit Ratio: {:.4f}, NDCG: {:.4f}'.format(epoch_num,
                                                                                                 pairs_ndcg,
                                                                                                 hit_ratio,
                                                                                                 ndcg))
    torch.save(net, model_store_dir + "/" + model_alias + "-" + str(epoch_num))
    net.train()


if "BASE_DIR" not in os.environ:
    os.environ["BASE_DIR"] = "/home/ec2-user/emb3"

random_state = np.random.RandomState(100)

net = PoolNet(num_items,
              embedding_dim=h['emb_dim'],
              layers=h["layers"],
              nonlinearity=h['nonlin'])

model = ImplicitSequenceModel(loss=h['loss'],
                              representation=net,
                              batch_size=h['batch'],
                              learning_rate=h['lr'],
                              l2=h['l2'],
                              n_iter=h['n_iter'],
                              embedding_dim=h['emb_dim'],
                              use_cuda=use_cuda,
                              random_state=random_state,
                              notify_loss_completion=notify_loss_completion,
                              notify_batch_eval_completion=notify_batch_eval_completion,
                              notify_epoch_completion=notify_epoch_completion,
                              log_loss_interval=5000,
                              log_eval_interval=20000,
                              amsgrad=h['amsgrad'],
                              adamw=h['adamw'],
                              betas=betas,
                              num_negative_samples=h['neg'])

logger.info("Model is initialized, now fitting..")
#model.fit(train_seq)
# -

sequences = train_seq.sequences.astype(np.int64)

sequences_tensor = gpu(torch.from_numpy(sequences),
                                       True)

from spotlight.torch_utils import *

net = gpu(net, True)

net

self = model
for minibatch_num, batch_sequence in enumerate(minibatch(sequences_tensor,
                                                             batch_size=64)):
    net.train()
    net.train()
    sequence_var = batch_sequence

    user_representation, _ = net.user_representation(
        sequence_var
    )
    
    print(user_representation.shape)
    print(sequence_var.shape)
    positive_prediction = net(user_representation,
                                    sequence_var)

    if self._loss == 'adaptive_hinge':
        negative_prediction = self._get_multiple_negative_predictions(
            sequence_var.size(),
            user_representation,
            n=self._num_negative_samples)
    else:
        negative_prediction = self._get_negative_prediction(sequence_var.size(),
                                                            user_representation)
    
    print("Done")
    self._optimizer.zero_grad()

    loss = self._loss_func(positive_prediction,
                           negative_prediction,
                           mask=(sequence_var != PADDING_IDX))
    loss_val = loss.item()
    epoch_loss += loss_val
    interval_loss += loss_val

    loss.backward()

model

# +
suffix = "1500K-with-3-vids"
train_data_path = "/home/ec2-user/emb3/data/train-aug-28-phase" + suffix
original_train_data = pd.read_parquet(train_data_path)
validate_neg_flatten_vids = pd.read_parquet(os.environ['BASE_DIR'] + "/data/validate-neg-flatten-aug-28-phase" + suffix)
validate_pos_flatten_vids = pd.read_parquet(os.environ['BASE_DIR'] + "/data/validate-pos-flatten-aug-28-phase" + suffix)
validation_train_data = original_train_data[original_train_data["uindex"].isin(validate_pos_flatten_vids.uindex.tolist())]
validation_train_data["vindex"] = validation_train_data["vindex"] + 1
interactions = Interactions(validation_train_data["uindex"].to_numpy(),
            validation_train_data["vindex"].to_numpy(),
            validation_train_data["pct_cvt"].to_numpy(),
            validation_train_data["latest_watch_time"].to_numpy(),
            num_users=len(original_train_data["uindex"].unique()),
            num_items=len(original_train_data["vindex"].unique()) + 2)

max_sequence_len=100
sequences = interactions.to_sequence(max_sequence_length=max_sequence_len, step_size=max_sequence_len)
# -

from spotlight.torch_utils import *

use_cuda= True
sqs = sequences.sequences[:5]
sequences_tensor = gpu(torch.from_numpy(sqs).long(),use_cuda)



num_items=20000



net = PoolNet(num_items,
              embedding_dim=h['emb_dim'],
              layers=h["layers"],
              nonlinearity=h['nonlin'])
net = gpu(net, True)

net

user_rp = net.user_representation(sequences_tensor)[1]
vids=[1,3, 4,5,7]
target_embedding = net.item_embeddings(gpu(torch.from_numpy(np.array(vids).reshape(-1, 1)).long(),use_cuda)).permute(0, 2, 1).squeeze()

targets = gpu(torch.from_numpy(np.array(vids).reshape(-1, 1)), True)
net(user_rp, targets)

use_cuda= True
sqs = sequences.sequences[:5]
sequences_tensor = gpu(torch.from_numpy(sqs).long(),use_cuda)
user_rp = mm.user_representation(sequences_tensor)[1]

StarcoderdataJupyterScriptsDedupFiltered