import sys
sys.path.append('../XGBoost_Prediction_Model/')

import warnings
warnings.filterwarnings("ignore")
import Predict
import torch
import numpy as np
import os
from os.path import isfile, isdir, join
from Magazine_Optimization import *
import cv2 as cv
import time

mypath = '../XGBoost_Prediction_Model/Magazine_Optimization_Demo/Magazines'
# results = torch.load('../XGBoost_Prediction_Model/Magazine_Optimization_Demo/results')
results = torch.load('../XGBoost_Prediction_Model/Magazine_Optimization_Demo/results_CNN')

for f in os.listdir(mypath):
    if isdir(join(mypath, f)):
        print('Currently processing Magazine '+f+'......')
        path_temp = join(mypath, f)
        dir_list = []
        for sub_f in os.listdir(path_temp):
            if isdir(join(path_temp, sub_f)):
                sub_path_temp = join(path_temp, sub_f)
                if (sub_f.split('_')[0]) == 'Jpg':
                    dir_list = os.listdir(sub_path_temp)
                    dir_list.sort()
                    for i in range(len(dir_list)):
                        dir_list[i] = join(sub_path_temp,dir_list[i])
                else:
                    Slots = torch.load(join(sub_path_temp,'Slots')).astype('int32')
        strategies = results[f]['AG'][0].split('; ')[:-1]
        newpath = '../XGBoost_Prediction_Model/Magazine_Optimization_Demo/Optimal_Magazines_CNN/'+f
        if not os.path.exists(newpath):
            os.makedirs(newpath)
        for item in strategies:
            ad_id = int(item.split(' ')[2])
            ctpg_id = int(item.split(' ')[5])
            ad_loc = 1-Slots[ad_id]
            ctpg_loc = Slots[ctpg_id]
            ad = cv.imread(dir_list[ad_id])[:,(ad_loc*640):((ad_loc+1)*640),:]
            ctpg = cv.imread(dir_list[ctpg_id])[:,(ctpg_loc*640):((ctpg_loc+1)*640),:]
            img = np.concatenate((ad,ctpg),axis=1)
            cv.imwrite(newpath+'/Newly_Assigned_Page'+str(ad_id)+'.jpg', img)