import sys
import os 
from tempfile import TemporaryDirectory
import numpy as np
import tensorflow.compat.v1 as tf
tf.get_logger().setLevel('ERROR') # only show error messages

from recommenders.utils.timer import Timer
from recommenders.utils.constants import SEED
from recommenders.models.deeprec.deeprec_utils import (
    prepare_hparams
)
from recommenders.datasets.amazon_reviews import download_and_extract, data_preprocessing, _create_vocab
from recommenders.datasets.download_utils import maybe_download


from recommenders.models.deeprec.models.sequential.sli_rec import SLI_RECModel as SeqModel 
# from recommenders.models.deeprec.models.sequential.asvd import A2SVDModel as SeqModel
# from recommenders.models.deeprec.models.sequential.caser import CaserModel as SeqModel
# from recommenders.models.deeprec.models.sequential.gru4rec import GRU4RecModel as SeqModel
# from recommenders.models.deeprec.models.sequential.sum import SUMModel as SeqModel

#from recommenders.models.deeprec.models.sequential.nextitnet import NextItNetModel

from recommenders.models.deeprec.io.sequential_iterator import SequentialIterator
#from recommenders.models.deeprec.io.nextitnet_iterator import NextItNetIterator

print("System version: {}".format(sys.version))
print("Tensorflow version: {}".format(tf.__version__))

yaml_file = '/home/jialia/wsdm/src/recommenders/examples/wsdm2022/sli_rec_B.yaml'  
  
RANDOM_SEED = SEED  # Set None for non-deterministic result

# data_path = os.path.join("tests", "resources", "deeprec", "slirec")
# data_path = '/home/jialia/wsdm/seq_datasets/B_full_feature_v2'
data_path = sys.argv[1]
print(os.path.abspath(data_path))  ## the path where I enter the cmd

# for test
train_file = os.path.join(data_path, r'train_instances.txt')
valid_file = os.path.join(data_path, r'valid_instances.txt')
test_file = os.path.join(data_path, r'valid.tsv')
pred_file = os.path.join(data_path, r'inter_test.tsv')
final_pred_file = os.path.join(data_path, r'final_test.tsv')
user_vocab = os.path.join(data_path, r'user_vocab.pkl')
item_vocab = os.path.join(data_path, r'item_vocab.pkl')
cate_vocab = os.path.join(data_path, r'category_vocab.pkl')
output_file = os.path.join(data_path, r'inter_test_output.txt')
submit_file = os.path.join(data_path, r'final_test_output.txt')


train_num_ngs = 9 # number of negative instances with a positive instance for training
valid_num_ngs = 9 # number of negative instances with a positive instance for validation
test_num_ngs = 9 # number of negative instances with a positive instance for testing


_create_vocab(
    [train_file, valid_file],
    user_vocab, item_vocab, cate_vocab
)


### NOTE:  
### remember to use `_create_vocab(train_file, user_vocab, item_vocab, cate_vocab)` to generate the user_vocab, item_vocab and cate_vocab files, if you are using your own dataset rather than using our demo Amazon dataset.
hparams = prepare_hparams(yaml_file, 
                        #   user_dropout=False,
                          embed_l2=0.,  
                          layer_l2=0., 
                          enable_BN=True, ##-- True
                          learning_rate=0.001,  # set to 0.01 if batch normalization is disable  else 0.001
                          epochs=100000,
                          EARLY_STOP=40000,
                          batch_size=400,
                          show_step=5000,
                          MODEL_DIR=os.path.join(data_path, "model/"),
                          SUMMARIES_DIR=os.path.join(data_path, "summary/"),
                          user_vocab=user_vocab,
                          item_vocab=item_vocab,
                          cate_vocab=cate_vocab,
                          need_sample=False,
                          train_num_ngs=train_num_ngs, # provides the number of negative instances for each positive instance for loss computation.
                          loss='log_loss', #'log_loss', 'softmax'
                          max_seq_length=50,
                          cont_feat_len=85,
                          use_cont_feat=False,
                          init_item_emb=False, 
                          shuffle=True
            )
print(hparams.values)

input_creator = SequentialIterator


model = SeqModel(hparams, input_creator, seed=RANDOM_SEED)

# model.load_model(os.path.join(data_path, "model_20220118_20k_0.8923", 'step_20000'))

with Timer() as train_time:
    model = model.fit(train_file, valid_file, valid_num_ngs=9, eval_metric='auc')
print('Time cost for training is {0:.2f} mins'.format(train_time.interval/60.0))

### model = model.fit(test_file, test_file, valid_num_ngs=9, eval_metric='auc') ##-- quick test

model.load_model(os.path.join(data_path, "model", 'best_model'))


res_syn = model.run_eval(test_file, num_ngs=9)
print(res_syn)
model.predict(pred_file, output_file)
model.predict(final_pred_file, submit_file)
# print('Job finished. B, continue training = 20k, seq=50')
# print('Job finished. B_v2, epoch=50k, seq=100')
## ASVD: 0.867497
## GRU:  0.877529
## SLi-Rec: 0.892736 
## B_v4: 0.8937

print("Job:B_full_feature_v2, with BN, no cont feat, seq=50, shuffle=True")

## B_full_feature_v2 no cont_feat, with BN
##5k: 0.8778
##10k: 0.8827
##20k: 0.8848
##25k: 0.8824
##35k: 0.8878
##40k: 0.8903
##45k: 0.8876
##50k: 0.8925
##55k: 0.8903
##60k: 0.8894
##65k: 0.8904
##70k: 0.8814
##75k: 0.8896
##80k: 0.8871
##85k: 0.8920

## with shuffle:
##5k: 0.8793
##10k: 0.8884
##15k: 0.8898
##20k: 0.8923
##25k: 0.8908
##30k: 0.8895
##35k: 0.8888
##40k: 0.8913
##45k: 0.8909
##50k: 0.8876
##65k: 0.8881
[((41, 13, 41, 60), 'os.path.join', 'os.path.join', ({(41, 26, 41, 35): 'data_path', (41, 37, 41, 59): '"""train_instances.txt"""'}, {}), "(data_path, 'train_instances.txt')", False, 'import os\n'), ((42, 13, 42, 60), 'os.path.join', 'os.path.join', ({(42, 26, 42, 35): 'data_path', (42, 37, 42, 59): '"""valid_instances.txt"""'}, {}), "(data_path, 'valid_instances.txt')", False, 'import os\n'), ((43, 12, 43, 49), 'os.path.join', 'os.path.join', ({(43, 25, 43, 34): 'data_path', (43, 36, 43, 48): '"""valid.tsv"""'}, {}), "(data_path, 'valid.tsv')", False, 'import os\n'), ((44, 12, 44, 54), 'os.path.join', 'os.path.join', ({(44, 25, 44, 34): 'data_path', (44, 36, 44, 53): '"""inter_test.tsv"""'}, {}), "(data_path, 'inter_test.tsv')", False, 'import os\n'), ((45, 18, 45, 60), 'os.path.join', 'os.path.join', ({(45, 31, 45, 40): 'data_path', (45, 42, 45, 59): '"""final_test.tsv"""'}, {}), "(data_path, 'final_test.tsv')", False, 'import os\n'), ((46, 13, 46, 55), 'os.path.join', 'os.path.join', ({(46, 26, 46, 35): 'data_path', (46, 37, 46, 54): '"""user_vocab.pkl"""'}, {}), "(data_path, 'user_vocab.pkl')", False, 'import os\n'), ((47, 13, 47, 55), 'os.path.join', 'os.path.join', ({(47, 26, 47, 35): 'data_path', (47, 37, 47, 54): '"""item_vocab.pkl"""'}, {}), "(data_path, 'item_vocab.pkl')", False, 'import os\n'), ((48, 13, 48, 59), 'os.path.join', 'os.path.join', ({(48, 26, 48, 35): 'data_path', (48, 37, 48, 58): '"""category_vocab.pkl"""'}, {}), "(data_path, 'category_vocab.pkl')", False, 'import os\n'), ((49, 14, 49, 63), 'os.path.join', 'os.path.join', ({(49, 27, 49, 36): 'data_path', (49, 38, 49, 62): '"""inter_test_output.txt"""'}, {}), "(data_path, 'inter_test_output.txt')", False, 'import os\n'), ((50, 14, 50, 63), 'os.path.join', 'os.path.join', ({(50, 27, 50, 36): 'data_path', (50, 38, 50, 62): '"""final_test_output.txt"""'}, {}), "(data_path, 'final_test_output.txt')", False, 'import os\n'), ((58, 0, 61, 1), 'recommenders.datasets.amazon_reviews._create_vocab', '_create_vocab', ({(59, 4, 59, 28): '[train_file, valid_file]', (60, 4, 60, 14): 'user_vocab', (60, 16, 60, 26): 'item_vocab', (60, 28, 60, 38): 'cate_vocab'}, {}), '([train_file, valid_file], user_vocab, item_vocab, cate_vocab)', False, 'from recommenders.datasets.amazon_reviews import download_and_extract, data_preprocessing, _create_vocab\n'), ((95, 8, 95, 58), 'recommenders.models.deeprec.models.sequential.sli_rec.SLI_RECModel', 'SeqModel', (), '', True, 'from recommenders.models.deeprec.models.sequential.sli_rec import SLI_RECModel as SeqModel\n'), ((38, 6, 38, 32), 'os.path.abspath', 'os.path.abspath', ({(38, 22, 38, 31): 'data_path'}, {}), '(data_path)', False, 'import os\n'), ((99, 5, 99, 12), 'recommenders.utils.timer.Timer', 'Timer', ({}, {}), '()', False, 'from recommenders.utils.timer import Timer\n'), ((105, 17, 105, 63), 'os.path.join', 'os.path.join', ({(105, 30, 105, 39): 'data_path', (105, 41, 105, 48): '"""model"""', (105, 50, 105, 62): '"""best_model"""'}, {}), "(data_path, 'model', 'best_model')", False, 'import os\n'), ((6, 0, 6, 15), 'tensorflow.compat.v1.get_logger', 'tf.get_logger', ({}, {}), '()', True, 'import tensorflow.compat.v1 as tf\n'), ((76, 36, 76, 69), 'os.path.join', 'os.path.join', ({(76, 49, 76, 58): 'data_path', (76, 60, 76, 68): '"""model/"""'}, {}), "(data_path, 'model/')", False, 'import os\n'), ((77, 40, 77, 75), 'os.path.join', 'os.path.join', ({(77, 53, 77, 62): 'data_path', (77, 64, 77, 74): '"""summary/"""'}, {}), "(data_path, 'summary/')", False, 'import os\n')]