#NOTA IMPORTANTE DEBE CORRERSE CON PYTHON 3 SINO DA ERRORES EN LOS CARACTERES ESPECIALES
import os
import json  #Libreria para leer JSON
from xlwt import Workbook #Libreria par escritura en Excel

#Define la ejecucion del programa
if __name__ == '__main__':
    listadoArchivos = os.listdir('.') #Obtiene el listado de archivos en un directorio
    listadoLimpio=[]
    metricas=[]

    #Selecciona solo los archivos presentes en el directorio que terminan en .json
    for each in listadoArchivos: #Para todos los archivos en el directorio
        if ".json" in each and ".py" not in each and ".xls" not in each:   #Si contienen .json en algun lado
            listadoLimpio.append(each)  #Agregar a la lista de archivos a procesar

    print(listadoLimpio)

    print("-- TABLAS --")

    for element in listadoLimpio: #Para cada archivo en la lista de archivos
        t=1 #Persistencia del cursor a través de archivos
        y=0#Persistencia del cursor a través de archivos
        cont = 0;#Cursor para definir en que fila del excel estoy escribiendo
        libro = Workbook() #Crear un libro de excel
        cursor=0;
        hoja1 = libro.add_sheet('Tablas') #Agregar una hoja al libro
        hoja1.write(0,1,'Tablas del Cubo')

        hoja1.write(0,2,'Tablas de la Bodega')
        hoja1.write(0,3,'Query')

        hoja1.col(0).width=1000
        hoja1.col(1).width=10000
        hoja1.col(2).width=10000
        hoja1.col(3).width=50000

        file= open(element)
        data=json.load(file)
        for table in data['model']['tables']:
            nombreTabla= table['name']
            hoja1.write(t,y,t)
            hoja1.write(t,y+1,nombreTabla)
            for anno in table['annotations']:
                if anno ['name'].find('_TM_ExtProp_DbTableName')!=-1:
                    hoja1.write(t,y+2,anno['value'])
            for  p in table['partitions']:
                hoja1.write(t,y+3,p['source']['query'])
            t=t+1
        hoja2= libro.add_sheet('Metricas')
        hoja2.write(0,1,'Metricas')
        hoja2.col(0).width=1000
        hoja2.col(1).width=6000
        hoja2.col(2).width=50000
        nombreMetrica=[]
        q=1
        g=0
        u=1

        for tab in data['model']['tables']:
            if  tab ['name'][0].find("M")== -1:  #Son Dimensiones
                print (tab['name'])
            else: #Son Metricaas
                hoja2.write(q,g,q)
                hoja2.write(q,g+1,tab['name'])
                for m in tab['measures']:
                    expr= m['name'],":=",m['expression']
                    #hoja2.write(u,y+2,str(expr)) #con caracteres especiales
                    hoja2.write(u,y+2,expr) #normal
                    u=u+1
                q=u

        jerarquias=[]
        hoja3= libro.add_sheet('Jerarquias')
        hoja3.write(0,1,'Tabla de la Jerarquia')
        hoja3.write(0,2,'Nombre de la Jerarquia')
        hoja3.write(0,3,'Niveles de la Jerarquia')
        hoja3.col(0).width=1000
        hoja3.col(1).width=5000
        hoja3.col(2).width=5000
        hoja3.col(3).width=7000
        cv=1
        ly=3
        lx=2
        cg=1

        for tx in data['model']['tables']:
            if 'hierarchies' in tx :
                for h in tx['hierarchies']:
                    hoja3.write(cv,y,cv)
                    hoja3.write(cv,y+1,tx['name'])
                    hoja3.write(cv,y+2,h['name'])
                    for l in h['levels']:
                        hoja3.write(cg,y+3,l['name'])
                        cg=cg+1
                    cv=cg
        hoja4= libro.add_sheet('Roles')
        hoja4.write(0,1,'Nombre del Rol')
        hoja4.write(0,2,'Permisos')
        hoja4.write(0,3,'Miembros')
        hoja4.col(0).width=1000
        hoja4.col(1).width=15000
        hoja4.col(2).width=4000
        hoja4.col(3).width=15000
        hr=1
        y=0
        cr=1
        for r in data['model']['roles']:
            hoja4.write(hr,y,hr)
            hoja4.write(hr,y+1,r['name'])
            hoja4.write(hr,y+2,r['modelPermission'])
            if 'members' in r:
                for m in r['members']:
                    hoja4.write(cr,y+3,m['memberName'])
                    cr=cr+1
            hr=hr+cr
        libro.save(element.replace('.json','.xls'))


import flask
import os
from flask import send_from_directory

app = flask.Flask(__name__)

#!/usr/bin/env python
# coding: utf-8

# ![header%20ipynb.png](attachment:header%20ipynb.png)

# # Hands-On
# ---
#
# Hands-On ini digunakan pada kegiatan Microcredential Associate Data Scientist 2021

# # Tugas Mandiri Pertemuan 10
# ---
# Pertemuan 10 (sepuluh) pada Microcredential Associate Data Scientist 2021 menyampaikan materi mengenai Membangun Model (Dasar Regresi dan Regresi Linier). silakan Anda kerjakan Latihan 1 s/d 20. Output yang anda lihat merupakan panduan yang dapat Anda ikuti dalam penulisan code :)

# # Latihan (1)
#
# ### Melakukan import library yang dibutuhkan

# In[1]:

# import library pandas
import pandas as pd

# Import library numpy
import numpy as np

# Import library matplotlib dan seaborn untuk visualisasi
import matplotlib.pyplot as plt
import seaborn as sns

# Import Module LinearRegression digunakan untuk memanggil algoritma Linear Regression.
from sklearn.linear_model import LinearRegression

# import Module train_test_split digunakan untuk membagi data kita menjadi training dan testing set.
from sklearn.model_selection import train_test_split

# import modul mean_absolute_error dari library sklearn
from sklearn.metrics import mean_absolute_error

#import math agar program dapat menggunakan semua fungsi yang ada pada modul math.(ex:sqrt)
import math

# me-non aktifkan peringatan pada python
import warnings
warnings.simplefilter("ignore")

@app.route('/favicon.ico')
def favicon():
    return send_from_directory(os.path.joint(app.root.path,'static'),  'favicon.ico', mimetype='image/favicon.png')

@app.route('/')
@app.route('/home')
def home():
    # ### Load Dataset

    # In[2]:

    #Panggil file (load file bernama CarPrice_Assignment.csv) dan simpan dalam dataframe Lalu tampilkan 10 baris awal dataset dengan function head()
    data = pd.read_csv('CarPrice_Assignment.csv')
    dataset = pd.DataFrame(data)
    dataset.head(10)

    # # Latihan (2)
    #
    # ### Review Dataset

    # In[3]:

    # melihat jumlah baris dan jumlah kolom (bentuk data) pada data df dengan fungsi .shape
    dataset.shape

    # Data kita mempunyai 26 kolom dengan 205 baris.

    # In[4]:

    # Melihat Informasi lebih detail mengenai struktur DataFrame dapat dilihat menggunakan fungsi info()
    dataset.info()

    # In[5]:

    # melihat statistik data untuk data numeric seperti count, mean, standard deviation, maximum, mininum, dan quartile.
    dataset.describe()

    # In[6]:

    # cek nilai yang hilang / missing values di dalam data
    dataset.isna().sum()

    # Ternyata data kita tidak ada missing values.

    # Simple linear regression atau regresi linear sederhana merupakan jenis regresi yang paling sederhana karena hanya melibatkan satu variabel bebas atau variabel independen X.

    # # Visualisasi data untuk pemilihan fitur / variabel independen X
    # ---
    # 1. Variabel y atau variabel dependent adalah 'price'
    # 2. Lakukan Visualisasi dalam penerapannya agar dapat terlihat jelas / mempermudah dalam membaca data tsb
    # 3. Untuk dapat menentukan variabel X yaitu dapat melihat korelasi antar variabel dengan variabel y / kolom 'price'

    # # Latihan (3)
    #
    # ### untuk dapat menentukan lebih detail / akurat dalam pemilihan fitur dapat dilihat dari hubungan korelasi nya dengan function corr()

    # In[7]:

    dataset.corr()

    # tampaknya enginesize, boreratio, horsepower, wheelbase memiliki korelasi yang signifikan dengan harga/price.

    # # Latihan (4)
    #
    # ### Buat Visualisasi scater plot dari kolom:
    #
    # 'enginesize', 'boreratio', 'stroke','compressionratio', 'horsepower', 'peakrpm', 'wheelbase', 'citympg', 'highwaympg'

    # In[8]:

    def pp(x,y,z):
        sns.pairplot(dataset, x_vars=[x, y, z], y_vars='price', size=4, aspect=1, diag_kind=None)
        plt.show()

    pp('enginesize', 'boreratio', 'stroke')
    pp('compressionratio', 'horsepower', 'peakrpm')
    pp('wheelbase', 'citympg', 'highwaympg')

    # # Latihan (5)
    #
    # ### Buat Visualisasi Heatmap dari kolom:
    #
    # 'enginesize', 'boreratio', 'stroke','compressionratio', 'horsepower', 'peakrpm', 'wheelbase', 'citympg', 'highwaympg'

    # In[9]:

    plt.figure(figsize = (8,8))
    data_fitur = dataset[['enginesize', 'boreratio', 'stroke','compressionratio', 'horsepower', 'peakrpm', 'wheelbase', 'citympg', 'highwaympg']]
    sns.heatmap(data_fitur.corr(),annot=False,fmt="f").set_title("Korelasi Heatmap Calon Variabel X")
    plt.show()

    # Dari hasil visualisasi diatas bahwa fitur/kolom enginesize memiliki korelasi yang tinggi terhadap kolom price / variabel dependent sehingga kita mengambil fitur/kolom enginesize untuk di training
    # * Independent variabel(x) adalah enginesize.
    # * Dependent variabel(y) adalah price.

    # # Latihan (6)
    #
    # ### Buat Visualisasi Scatter Plot antara calon variabel X(enginesize) dan y(price):

    # In[10]:

    plt.scatter(dataset['enginesize'], dataset['price'])
    plt.xlabel('enginesize')
    plt.ylabel('price')
    plt.title('Scatter Plot enginesize vs Price')
    plt.show()

    # Scatter plot menunjukkan dengan jelas hubungan antarvariabel serta sebarannya di dataset. Selain itu, dengan scatter plot juga kita dapat mengindikasikan bahwa variabel enginesize dan price memiliki hubungan linear.

    #     Catatan : korelasi 0.874145 adalah nilai yang cukup tinggi, artinya nilai price benar-benar sangat dipengaruhi oleh nilai enginesize, karena korelasi tinggi maka algoritma Regresi Linier ini cocok digunakan untuk data tersebut.

    # # Latihan (7)
    #
    # ### definisi variabel X(enginesize) dan y(price):

    # In[11]:

    # Prepare data
    # Pertama, buat variabel x dan y.
    x = dataset['enginesize'].values.reshape(-1,1)
    y = dataset['price'].values.reshape(-1,1)

    # Formula Regresi Linear
    # ![gambar.png](attachment:gambar.png)
    #
    #
    # "Jika kita melihat formula regresi linear di atas, kita pasti ingat rumus persamaan garis yang pernah dipelajari di bangku sekolah, yaitu y = mx + c, dimana m merupakan gradien atau kemiringan garis dan c merupakan konstanta."

    # * from scratch
    # * y = ax + b atau y = w1x + w0 atau y = mx + c
    # * x = input
    # * y = output
    # * b atau w0 = intercept / bias
    # * a atau w1 = slope / gradient / coefficient

    # # Latihan (8)
    #
    # ### definisi variabel nilai mean/rata-rata X(enginesize) dan nilai mean/rata-rata y(price):

    # In[12]:

    x_mean = np.mean(x)
    y_mean = np.mean(y)
    print('nilai mean var x: ', x_mean,'\n'
        'nilai mean var y: ', y_mean)

    # # Latihan (9)
    #
    # ### carilah nilai koefisien korelasi nya dengan rumus dibawah:

    # ![gambar.png](attachment:gambar.png)

    # In[13]:

    atas = sum((x - x_mean)*(y - y_mean))
    bawah = math.sqrt((sum((x - x_mean)**2)) * (sum((y - y_mean)**2)))
    correlation = atas/bawah
    print('Nilai Correlation Coefficient: ', correlation)

    #
    # ### carilah nilai parameter theta 1 dan theta 0 dengan rumus dibawah:

    #
    # ![gambar.png](attachment:gambar.png)

    #     theta_1 = ((111-104.11) * (13495-13276.71)) + ... + ((114-104.11) * (22625-13276.71)) / ((111-104.11)^2 + ... + (114-104.11)^2)

    # # Latihan (10)
    #
    # ### carilah nilai theta_1 atau nilai slope

    # In[14]:

    # slope
    # Slope adalah tingkat kemiringan garis, intercept
    # adalah jarak titik y pada garis dari titik 0
    variance = sum((x - x_mean)**2)
    covariance = sum((x - x_mean) * (y - y_mean))
    theta_1 = covariance/variance
    print('Nilai theta_1: ',theta_1)

    # # Latihan (11)
    #
    # ### carilah nilai theta_0 atau nilai intercept

    # In[15]:

    # intercept
    theta_0 = y_mean - (theta_1 * x_mean)
    print('Nilai theta_1: ',theta_0)

    # ### Maka persamaan garis :
    #
    # #     **y = 167.69x - 8005.44**
    #
    # Jadi persamaan garis diatas dapat digunakan untuk melakukan prediksi apabila kita memiliki data enginesize yang baru, price dapat diperkirakan dengan rumus tersebut, masukkan nilai enginesize baru ke x, maka perkiraan nilai y (price) akan didapat.

    # # Latihan (12)
    #
    # ### carilah nilai prediksi secara manual dan buatlah visualisasi scater plot nya

    # In[16]:

    # prediction manual
    y_pred = theta_0 + (theta_1 * 130)

    print(y_pred)

    # In[17]:

    # visualisasi prediksi dengan scatter plot
    y_pred = theta_0 + (theta_1 * x)

    plt.scatter(dataset['enginesize'], dataset['price'])
    plt.plot(x, y_pred, c='r')
    plt.xlabel('enginesize')
    plt.ylabel('Price')
    plt.title('Plot enginesize vs Price')

    # Linier Regression digunakan untuk Prediksi dengan mencari pola garis terbaik antara variable independent dan dependen
    #
    # Pros:
    #
    #     Mudah diimplementasikan
    #     Digunakan untuk memprediksi nilai numerik/ continous /data jenis interval dan ratio
    #
    # Cons :
    #
    #     Cenderung mudah Overfitting
    #     Tidak dapat digunakan bila relasi antara variabel independen dan dependen tidak linier atau korelasi variabel rendah

    # # Linier Regression dengan menggunakan library sklearn

    # 1. Pertama yang kita lakukan adalah split data, Train/test split adalah salah satu metode yang dapat digunakan untuk mengevaluasi performa model machine learning. Metode evaluasi model ini membagi dataset menjadi dua bagian yakni bagian yang digunakan untuk training data dan untuk testing data dengan proporsi tertentu. Train data digunakan untuk fit model machine learning, sedangkan test data digunakan untuk mengevaluasi hasil fit model tersebut.
    #
    # ![image.png](attachment:image.png)
    #
    #
    # Python memiliki library yang dapat mengimplementasikan train/test split dengan mudah yaitu Scikit-Learn. Untuk menggunakannya, kita perlu mengimport Scikit-Learn terlebih dahulu, kemudian setelah itu kita dapat menggunakan fungsi train_test_split().

    # # Latihan (13)
    #
    # ### split data train dan test dengan function  train_test_split() dengan train_size=0.8, test_size=0.2 dan random_state=100

    # In[18]:

    X_train, X_test, y_train, y_test = train_test_split(x, y, train_size = 0.8, test_size = 0.2, random_state = 100)

    # * X_train: Untuk menampung data source yang akan dilatih.
    # * X_test: Untuk menampung data target yang akan dilatih.
    # * y_train: Untuk menampung data source yang akan digunakan untuk testing.
    # * y_test: Untuk menampung data target yang akan digunakan untuk testing.
    #
    # X dan y adalah nama variabel yang digunakan saat mendefinisikan data source dan data target. Parameter *test_size* digunakan untuk mendefinisikan ukuran data testing. Dalam contoh di atas, test_size=0.2 berarti data yang digunakan sebagai data testing adalah sebesar 20% dari keseluruhan dataset.
    #
    # Perlu diketahui bahwa metode ini akan membagi train set dan test set secara random atau acak. Jadi, jika kita mengulang proses running, maka tentunya hasil yang didapat akan berubah-ubah. Untuk mengatasinya, kita dapat menggunakan parameter *random_state*

    # # Latihan (14)
    #
    # ### buat object variabel linier regression

    # In[19]:

    regressor = LinearRegression()

    # # Latihan (15)
    #
    # ### training the model menggunakan training data yang sudah displit sebelumnya.

    # In[20]:

    regressor.fit(X_train, y_train)

    # # Latihan (16)
    #
    # ### cari tau nilai slope/koefisien (m) dan intercept (b), dengan menggunakan function dari library sklearn -> LinierRegression

    # In[21]:

    print(regressor.coef_)
    print(regressor.intercept_)

    # Dari nilai m dan b diatas, kalau dimasukan ke dalam rumus persamaan menjadi:
    # #     **y = 168.17x - 8037.06**

    # # Latihan (17)
    #
    # ### cari tahu accuracy score dari model kita menggunakan testing data yang sudah displit sebelumnya. Dan nilai korelasinya

    # In[22]:

    regressor.score(X_test, y_test)

    # Model kita mendapatkan accuracy score sebesar 80.68%

    # In[23]:

    print('Correlation: ', math.sqrt(regressor.score(X_test,y_test)))

    # # Latihan (18)
    #
    # ### visualisasi Regression Line menggunakan data testing.

    # In[24]:

    y_prediksi = regressor.predict(X_test)

    plt.scatter(X_test,y_test)
    plt.plot(X_test, y_prediksi, c='r')
    plt.xlabel('enginesize')
    plt.ylabel('Price')
    plt.title('Plot enginesize vs Price')

    # Garis merah merupakan Regression Line dari model yang telah dibuat sebelumnya.

    # # Latihan (19)
    #
    # ### Setelah kita yakin dengan model yang dibuat, selanjutnya adalah prediksi dari harga mobil dengan enginesize 100, 150, dan 200.

    # In[25]:

    #Prediksi harga mobil dengan enginesize 130.

    print('nilai prediksi harga dengan enginesize 100 : ',regressor.predict([[100]]))
    print('nilai prediksi harga dengan enginesize 150 : ',regressor.predict([[150]]))
    print('nilai prediksi harga dengan enginesize 200 : ',regressor.predict([[200]]))

    # In[26]:

    np_table = np.concatenate((X_test,y_test,y_prediksi), axis=1)
    new_dataframe = pd.DataFrame(data=np_table, columns=['x_test','y_test','y_predict'])

    # In[27]:

    new_dataframe

    # Semakin tinggi nilai error, semakin besar errornya

    # # Latihan (20)
    #
    # ### Cetak nilai Mean Absolute Error, Mean Squared Error, dan Root Mean Squared Error

    # In[28]:

    from sklearn import metrics
    print('Mean Absolute Error:', metrics.mean_absolute_error(y_test, y_prediksi))
    print('Mean Squared Error:', metrics.mean_squared_error(y_test, y_prediksi))
    print('Root Mean Squared Error:', np.sqrt(metrics.mean_squared_error(y_test, y_prediksi)))

    # In[29]:

    plt.title('Comparison of Y values in test and the Predicted values')
    plt.ylabel('Test Set')
    plt.xlabel('Predicted values')
    plt.plot(y_prediksi, '.', y_test, 'x')
    plt.show()

    return "Hasil Eksekusi"

if __name__ == "__main__":
    app.secret_key = 'ItIsASecret'
    app.debug = True
    app.run()


from __future__ import absolute_import
from __future__ import print_function

import os
import sys
import optparse
import random
import numpy as np
import xml.etree.ElementTree as ET

if 'SUMO_HOME' in os.environ:
    tools = os.path.join(os.environ['SUMO_HOME'], 'tools')
    sys.path.append(tools)
else:
    sys.exit("please declare environment variable 'SUMO_HOME'")
from sumolib import checkBinary
import traci
import traci.constants as tc
import sumolib

SPEED = 5
depart_arr = [
    "-gneE38", "-gneE39", "-gneE40", "-gneE41", "-gneE42", "-gneE43",
    "-gneE44", "-gneE45", "-gneE46", "-gneE47", "-gneE48", "-gneE49"
]
arrival_arr = [
    "gneE38", "gneE39", "gneE40", "gneE41", "gneE42", "gneE43", "gneE44",
    "gneE45", "gneE46", "gneE47", "gneE48", "gneE49"
]
#Init

def init(is_gui):
    if is_gui:
        sumoBinary = os.path.join(os.environ['SUMO_HOME'], 'bin/sumo-gui')
    else:
        sumoBinary = os.path.join(os.environ['SUMO_HOME'], 'bin/sumo')

    sumoCmd = [sumoBinary, "-c", "data/deadlock.sumocfg"]
    tree = ET.parse('data/allRoute.rou.xml')
    root = tree.getroot()
    traci.start(sumoCmd)

def detect_gridlock(vehID, next_lane):
    try:
        locked_veh = traci.lane.getLastStepVehicleIDs(next_lane)[-1]
        locked_veh_speed = traci.vehicle.getSpeed(locked_veh)
        if locked_veh_speed == 0:
            return vehID, locked_veh
    except:
        pass

def control_vehicle(vehID, net):
    try:
        lanes = traci.vehicle.getBestLanes(vehID)[0][5]
        current_edge = traci.lane.getEdgeID(lanes[0])
        nextNodeID = net.getEdge(current_edge).getToNode().getID()
        vehicle_pos = traci.vehicle.getPosition(vehID)
        junction_pos = traci.junction.getPosition(nextNodeID)
        junction_vehicle_distance = traci.simulation.getDistance2D(
            vehicle_pos[0], vehicle_pos[1], junction_pos[0], junction_pos[1])

        if 1 < junction_vehicle_distance < 35:
            judge_num = (traci.lane.getLength(lanes[1]) *
                         (1 - traci.lane.getLastStepOccupancy(lanes[1])))

            if judge_num < 50:
                traci.vehicle.setSpeed(vehID, 0)
            else:
                traci.vehicle.setSpeed(vehID, SPEED)

        else:
            traci.vehicle.setSpeed(vehID, SPEED)

    except:
        pass

#Start
net = sumolib.net.readNet('data/deadlock.net.xml')
#Gridlock error
color_dic = {
    1: (255, 0, 0),
    2: (0, 255, 0),
    3: (0, 0, 255),
    4: (255, 255, 255)
}

def simulation(num_of_vehicles, with_program):
    init(True)
    for i in range(1, num_of_vehicles):
        traci.route.add(
            f"{i}_1",
            ["gneE9", "-gneE1", "-gneE0", "-gneE3", "-gneE26", "gneE38"])
        traci.route.add(
            f"{i}_2",
            ["gneE15", "-gneE3", "-gneE2", "-gneE1", "-gneE24", "gneE45"])
        traci.route.add(
            f"{i}_3",
            ["gneE6", "-gneE2", "-gneE1", "-gneE0", "gneE27", "gneE41"])
        traci.route.add(
            f"{i}_4",
            ["gneE12", "-gneE0", "-gneE3", "-gneE2", "-gneE25", "gneE49"])
        for k in range(1, 5):
            traci.vehicle.addLegacy(f"{k}_{i}", f"{i}_{k}", 0)
            traci.vehicle.setMaxSpeed(f"{k}_{i}", SPEED)
            traci.vehicle.setColor(f"{k}_{i}", color_dic[k])
    while traci.simulation.getMinExpectedNumber() > 0:
        if with_program:
            locked_vehicles = {}
            leaders = []
            gridlock = []
            control_arr = []
            current_time = traci.simulation.getTime()
            all_vehicles = traci.vehicle.getIDList()

            for vehicle in all_vehicles:
                leader = (traci.vehicle.getLeader(vehicle))
                vehicle_lane = traci.vehicle.getLaneID(vehicle)
                try:
                    if vehicle_lane[0] != ":":
                        if leader == None:
                            leaders.append(vehicle)
                        else:
                            leader_lane = traci.vehicle.getLaneID(leader[0])
                            if vehicle_lane != leader_lane:
                                leaders.append(vehicle)
                except:
                    pass
            for leader in leaders:
                leader_lanes = traci.vehicle.getBestLanes(leader)[0][5]
                control_arr.append(leader)

                if len(leader_lanes) > 1:
                    locked_veh = detect_gridlock(leader, leader_lanes[1])
                    if locked_veh:
                        locked_vehicles[locked_veh[0]] = locked_veh[1]

            for init_veh in (locked_vehicles):
                if len(gridlock) > 0:
                    if init_veh not in gridlock[0]:
                        temp = []
                        index = init_veh
                        for k in range(15):
                            try:
                                next_vehicle = locked_vehicles[index]
                                temp.append(next_vehicle)
                                index = next_vehicle
                                if next_vehicle == init_veh:
                                    gridlock.append(temp)
                                    break

                            except:
                                break
                elif len(gridlock) == 0:
                    temp = []
                    index = init_veh
                    for k in range(15):
                        try:
                            next_vehicle = locked_vehicles[index]
                            temp.append(next_vehicle)
                            index = next_vehicle
                            if next_vehicle == init_veh:
                                gridlock.append(temp)
                                break

                        except:
                            break

            for v in control_arr:
                control_vehicle(v, net)
            for vehicles in gridlock:
                for v in vehicles:
                    traci.vehicle.setSpeed(v, SPEED)
        traci.simulationStep()
        if traci.simulation.getTime() == 5000:
            traci.close()
    traci.close()

num_of_vehicles = int(input("Num of vehicles :"))
with_program = int(input("With control program (0=false, 1=true):"))
simulation(num_of_vehicles, bool(with_program))


from typing import Dict, List

import torch
import torch.nn as nn
import torch.nn.functional as F

from .common import build_target_normalization, pad_avail, pad_points
from .open_loop_model import VectorizedModel
from .global_graph import MultimodalMultiheadAttentionGlobalHead

class VectorizedMultiModalPlanningModel(VectorizedModel):
    """Vector-based multimodal planning model."""

    def __init__(
        self,
        history_num_frames_ego: int,
        history_num_frames_agents: int,
        future_num_frames: int,
        num_outputs: int,
        num_modes: int,
        weights_scaling: List[float],
        criterion: nn.Module,
        global_head_dropout: float,
        disable_other_agents: bool,
        disable_map: bool,
        disable_lane_boundaries: bool,
        coef_alpha: float = 0.5,
    ) -> None:
        """Initializes the multimodal planning model.

        :param history_num_frames_ego: number of history ego frames to include
        :param history_num_frames_agents: number of history agent frames to include
        :param future_num_frames: number of future frames to predict
        :param num_outputs: number of output dimensions, by default is 3: x, y, heading
        :param num_modes: number of modes in predicted outputs
        :param weights_scaling: target weights for loss calculation
        :param criterion: loss function to use
        :param global_head_dropout: float in range [0,1] for the dropout in the MHA global head. Set to 0 to disable it
        :param disable_other_agents: ignore agents
        :param disable_map: ignore map
        :param disable_lane_boundaries: ignore lane boundaries
        :param coef_alpha: hyper-parameter used to trade-off between trajectory distance loss and classification
        cross-entropy loss
        """
        num_targets = (future_num_frames * num_outputs + 1) * num_modes
        super().__init__(
            history_num_frames_ego,
            history_num_frames_agents,
            num_targets,
            weights_scaling,
            criterion,
            global_head_dropout,
            disable_other_agents,
            disable_map,
            disable_lane_boundaries,
        )
        self.num_modes = num_modes
        self.num_timestamps = future_num_frames
        self.num_outputs = num_outputs
        self.coef_alpha = coef_alpha

        self.normalize_targets = True
        num_timestamps = self.num_timestamps

        if self.normalize_targets:
            scale = build_target_normalization(num_timestamps)
            self.register_buffer("xy_scale", scale)

        # 重写父类的global_head，将输出[future_num_frames,num_outputs]合并成一个维度num_targets
        self.global_head = MultimodalMultiheadAttentionGlobalHead(
            self._d_global, num_targets, dropout=self._global_head_dropout
        )

    def forward(self, data_batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        ## Load and prepare vectors for the model call, split into map and agents

        # ==== LANES ====
        # batch size x num lanes x num vectors x num features
        polyline_keys = ["lanes_mid", "crosswalks"]
        if not self.disable_lane_boundaries:
            polyline_keys += ["lanes"]
        avail_keys = [f"{k}_availabilities" for k in polyline_keys]

        max_num_vectors = max([data_batch[key].shape[-2] for key in polyline_keys])
        map_polys = torch.cat([pad_points(data_batch[key], max_num_vectors) for key in polyline_keys], dim=1)
        map_polys[..., -1].fill_(0)
        # batch size x num lanes x num vectors
        map_availabilities = torch.cat([pad_avail(data_batch[key], max_num_vectors) for key in avail_keys], dim=1)

        # ==== AGENTS ====
        # batch_size x (1 + M) x seq len x self._vector_length
        agents_polys = torch.cat(
            [data_batch["agent_trajectory_polyline"].unsqueeze(1), data_batch["other_agents_polyline"]], dim=1
        )
        # batch_size x (1 + M) x num vectors x self._vector_length
        agents_polys = pad_points(agents_polys, max_num_vectors)

        # batch_size x (1 + M) x seq len
        agents_availabilities = torch.cat(
            [
                data_batch["agent_polyline_availability"].unsqueeze(1),
                data_batch["other_agents_polyline_availability"],
            ],
            dim=1,
        )
        # batch_size x (1 + M) x num vectors
        agents_availabilities = pad_avail(agents_availabilities, max_num_vectors)

        # batch_size x (1 + M) x num features
        type_embedding = self.type_embedding(data_batch).transpose(0, 1)
        lane_bdry_len = data_batch["lanes"].shape[1]

        ## call the model with these features
        #[batch_size, (future_num_frames * num_outputs + 1) * num_modes]
        outputs_all, attns = self.model_call(
            agents_polys, map_polys, agents_availabilities, map_availabilities, type_embedding, lane_bdry_len
        )
        batch_size = len(outputs_all)

        # [batch_size, num_timestamps * num_outputs * num_modes]
        outputs = outputs_all[:, :-self.num_modes]
        # [batch_size, num_modes], classification flags
        outputs_nll = outputs_all[:, -self.num_modes:]

        # .view(-1, self.num_timesteps, self.num_outputs)

        # calculate loss or return predicted position for inference
        if self.training:
            if self.criterion is None:
                raise NotImplementedError("Loss function is undefined.")

            # [batch_size, num_timestamps * num_modes, num_outputs]
            outputs = outputs.view(
                -1, self.num_timestamps * self.num_modes, self.num_outputs
            )

            # [batch_size, num_timestamps, 2]
            xy = data_batch["target_positions"]
            # [batch_size, num_timestamps, 1]
            yaw = data_batch["target_yaws"]

            if self.normalize_targets:
                xy /= self.xy_scale

            # [batch_size, num_timestamps, num_outputs]
            targets = torch.cat((xy, yaw), dim=-1)
            # [batch_size, num_timestamps]
            target_weights = data_batch["target_availabilities"] > 0.5
            # [batch_size, num_timestamps, num_outputs]
            target_weights = target_weights.unsqueeze(-1) * self.weights_scaling
            # [batch_size, num_timestamps * num_modes, num_outputs]
            losses = self.criterion(
                outputs, targets.repeat(1, self.num_modes, 1)
            ) * target_weights.repeat(1, self.num_modes, 1)
            # [batch_size, num_modes]
            cost_dist = losses.view(batch_size, self.num_modes, -1).mean(dim=-1)
            # [batch_size,],最接近target轨迹的output轨迹的编号
            assignment = cost_dist.argmin(dim=-1)
            # 最接近轨迹的距离loss
            loss_dist = cost_dist[torch.arange(batch_size, device=outputs.device), assignment].mean()
            # 每个output轨迹是接近target轨迹的分类概率loss
            loss_nll = F.cross_entropy(outputs_nll, assignment)
            train_dict = {"loss": loss_dist + self.coef_alpha * loss_nll}
            return train_dict
        else:
            outputs = outputs.view(batch_size, self.num_modes, self.num_timestamps, self.num_outputs)
            outputs_selected = outputs[torch.arange(batch_size, device=outputs.device), outputs_nll.argmax(-1)]
            pred_positions, pred_yaws = outputs_selected[..., :2], outputs_selected[..., 2:3]
            pred_pos_all = outputs[..., :2]
            pred_yaw_all = outputs[..., 2:3]
            if self.normalize_targets:
                pred_positions *= self.xy_scale
                pred_pos_all *= self.xy_scale

            eval_dict = {
                "positions": pred_positions,
                "yaws": pred_yaws,
                "positions_all": pred_pos_all,
                "yaws_all": pred_yaw_all,
            }
            if attns is not None:
                eval_dict["attention_weights"] = attns
            return eval_dict


# %%
import pandas as pd
from matplotlib import pyplot as plt
# %matplotlib inline

# %%
df = pd.read_csv("insurance_data.csv")
df.head()

# %%
plt.scatter(df.age, df.bought_insurance, marker="+", color="red")

# %%
df.shape

# %%
from sklearn.model_selection import train_test_split

# %%
X_train, X_test, y_train, y_test = train_test_split(df[['age']], df.bought_insurance, test_size=0.10)

# %%
X_test

# %%
from sklearn.linear_model import LogisticRegression

# %%
model = LogisticRegression()

# %%
model.fit(X_train, y_train)

# %%
model.predict(X_test)

# %%
model.score(X_test, y_test)

# %%
model.predict_proba(X_test)

# %% [markdown]
# model.predict throws a ValueError when handling scalar, therefore I converted the value to a 2D array using double parenthesis [[]].

# %%
model.predict([[25]])


from collections import defaultdict

def get_coordinate_info():
	text_input = open('input.txt', "r")
	all_inputs = []
	coordinate_info = []

	for line in text_input:
		all_inputs.append(line.rstrip())

	for line in all_inputs:
		line_input = line.replace(' ', '').replace('#', '').replace('@', ':').replace('x', ':').replace(',', ':')
		line_input = line_input.split(":")
		coordinate_info.append(line_input)

	return coordinate_info

def create_coordinate_mapping():
	coordinates = get_coordinate_info()
	fabric_squares = defaultdict(list)

	for coordinate in coordinates:
		claim_number, x, y, w, h = coordinate
		x_squares = [(i + int(x), int(y)) for i in range(int(w))]
		claimed_squares = [(int(x), int(y) + i) for x, y in x_squares for i in range(int(h))]
		for square in claimed_squares:
			fabric_squares[square].append(claim_number)

	return fabric_squares

# Part 1

def get_overlapping_squares():
	fabric_squares = create_coordinate_mapping()

	total_overlapping = 0
	for key, val in fabric_squares.items():
		if len(val) > 1:
			total_overlapping += 1

	return total_overlapping

# Part 2
def get_non_overlapping_claim():
	coordinates = get_coordinate_info()
	claim_ids = set([coordinate[0] for coordinate in coordinates])
	fabric_squares = create_coordinate_mapping()

	multiple_claims = set()
	for coordinate, claims in fabric_squares.items():
		if len(claims) > 1:
			for claim in claims:
				multiple_claims.add(claim)

	return claim_ids - multiple_claims


from __future__ import absolute_import, division, print_function
import argparse
import base64
import os.path as op
import random, time, json
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler
from tqdm import tqdm

from oscar.utils.logger import setup_logger
from oscar.utils.tsv_file import TSVFile
from oscar.utils.tsv_file_ops import tsv_writer
from oscar.utils.misc import (mkdir, set_seed,
        load_from_yaml_file, find_file_path_in_yaml)
from oscar.utils.caption_evaluate import (evaluate_on_coco_caption, evaluate_on_reddit_caption,
        evaluate_on_nocaps, ScstRewardCriterion)
from oscar.utils.cbs import ConstraintFilter, ConstraintBoxesReader
from oscar.utils.cbs import FiniteStateMachineBuilder
from oscar.modeling.modeling_bert import BertForMultiModalDialog
from transformers.pytorch_transformers import BertTokenizer, BertConfig
from transformers.pytorch_transformers import AdamW, WarmupLinearSchedule, WarmupConstantSchedule
import torch.nn as nn
import sys
sys.setrecursionlimit(1000000)

total_num_boxes = 36

class CaptionTSVDataset(Dataset):
    def __init__(self, yaml_file, tokenizer=None, add_od_labels=True,
            max_img_seq_length=50, max_seq_length=70, max_seq_a_length=40, max_seq_c_length=40,
            is_train=True, mask_prob=0.15,mask_response_prob=0.7, max_masked_tokens=3, **kwargs):
        """Constructor.
        Args:
            yaml file with all required data (image feature, caption, labels, etc)
            tokenizer: tokenizer for text processing.
            add_od_labels: whether to add labels from yaml file to BERT.
            max_img_seq_length: max image sequence length.
            max_seq_length: max text sequence length.
            max_seq_a_length: max caption sequence length.
            is_train: train or test mode.
            mask_prob: probability to mask a input token.
            max_masked_tokens: maximum number of tokens to be masked in one sentence.
            kwargs: other arguments.
        """
        self.yaml_file = yaml_file
        self.cfg = load_from_yaml_file(yaml_file)
        self.root = op.dirname(yaml_file)
        self.tag_and_feat_file = find_file_path_in_yaml(self.cfg['feature'], self.root)
        # self.feat_file = find_file_path_in_yaml(self.cfg['feature'], self.root)
        self.dialog_file = find_file_path_in_yaml(self.cfg.get('dialog'), self.root)
        self.dialog2img_file = find_file_path_in_yaml(self.cfg.get('dialog2img'), self.root)

        assert op.isfile(self.tag_and_feat_file)
        # if add_od_labels: assert op.isfile(self.label_file)
        if is_train: assert op.isfile(self.dialog_file) and tokenizer is not None

        # self.label_tsv = None if not self.label_file else TSVFile(self.label_file)
        self.tag_and_feat_tsv = TSVFile(self.tag_and_feat_file, gen_lineidx=True)
        if self.dialog_file and op.isfile(self.dialog_file):
            self.dialogs = [line.decode().rstrip('\n').split('\t') for line in open(self.dialog_file,'rb')] # [[sent1,sent2,....sentn],...[sent1,sent2,....sentn]]
        if self.dialog2img_file and op.isfile(self.dialog2img_file):
            self.dialog2img = [line.decode().rstrip('\n').split('\t') for line in open(self.dialog2img_file,'rb')] # [[top1,top2,....topn],...[top1,top2,....topn]]

        # print('len of dialogs', len(self.dialogs))

        self.tokenizer = tokenizer
        self.tensorizer = CaptionTensorizer(self.tokenizer, max_img_seq_length,
                max_seq_length, max_seq_a_length, max_seq_c_length, mask_prob, mask_response_prob, max_masked_tokens,
                is_train=is_train)
        self.add_od_labels = add_od_labels
        self.is_train = is_train
        self.kwargs = kwargs
        # self.image_keys = self.prepare_image_keys()
        self.key2index = self.prepare_image_key_to_index()
        # self.key2captions = self.prepare_image_key_to_captions()

    def get_valid_tsv(self):
        return self.tag_and_feat_tsv

    # def prepare_image_keys(self):
    #     tsv = self.get_valid_tsv()
    #     return [tsv.seek(i)[0] for i in range(tsv.num_rows())]

    def get_dialog_file(self):
        return self.dialog_file

    def prepare_image_key_to_index(self):
        tsv = self.get_valid_tsv()
        # print(tsv.seek(0)[0])
        # print(type(tsv.seek(0)[0]))
        # print(tsv.seek(0)[1])
        return {tsv.seek(i)[0] : i for i in range(tsv.num_rows())}

    def get_image_features(self, img_idx):
        # feat_info = json.loads(self.tag_and_feat_tsv.seek(img_idx)[1])
        feat_info = eval(self.tag_and_feat_tsv.seek(img_idx)[1]) #transfer to dict
        features = np.frombuffer(base64.b64decode(feat_info['features']), np.float32
                ).reshape((total_num_boxes, -1))

        fea = torch.Tensor(features)

        return fea

    def get_od_labels(self, img_idx):
        od_labels = None
        if self.add_od_labels:
            label_info = eval(self.tag_and_feat_tsv.seek(img_idx)[1])
            # If filter the repeat token
            # od_labels = " ".join(list(set(label_info['objects_class'])))
            od_labels = " ".join(label_info['objects_class'])
        return od_labels

    def __getitem__(self, idx):

        dialogs = self.dialogs[idx]
        context = '[SEP]'.join(dialogs[:-1])
        # if add tag bias
        context += '[TAG]'
        #else when decoding
        # context += '[SEP]'
        response = dialogs[-1]

        img_key = op.splitext(self.dialog2img[idx][0])[0] # get the img idx from img path: get the xxx from xxx.jpg
        img_idx = self.key2index[img_key]

        features = self.get_image_features(img_idx)

        od_labels = self.get_od_labels(img_idx)

        text_b = context + od_labels

        example = self.tensorizer.tensorize_example(response, features, text_b=text_b)

        return example

    def __len__(self):
        return len(self.dialogs)

class CaptionTensorizer(object):
    def __init__(self, tokenizer, max_img_seq_length=50, max_seq_length=70,
            max_seq_a_length=40, max_seq_c_length=40, mask_prob=0.15, mask_response_prob=0.7, max_masked_tokens=3,
            is_train=True):
        """Constructor.
        Args:
            tokenizer: tokenizer for text processing.
            max_img_seq_length: max image sequence length.
            max_seq_length: max text sequence length.
            max_seq_a_length: max caption sequence length.
            is_train: train or test mode.
            mask_prob: probability to mask a input token.
            max_masked_tokens: maximum number of tokens to be masked in one sentence.
        """
        self.tokenizer = tokenizer
        self.is_train = is_train
        self.max_img_seq_len = max_img_seq_length
        self.max_seq_len = max_seq_length
        self.max_seq_a_len = max_seq_a_length
        self.max_seq_c_len = max_seq_c_length
        self.mask_prob = mask_prob
        self.mask_response_prob = mask_response_prob
        self.max_masked_tokens = max_masked_tokens
        self._triangle_mask = torch.tril(torch.ones((self.max_seq_len,
            self.max_seq_len), dtype=torch.long))

    def tensorize_example(self, text_a, img_feat, text_b=None,
            cls_token_segment_id=0, pad_token_segment_id=0,
            sequence_a_segment_id=0, sequence_b_segment_id=1, sequence_tag_segment_id=2):
        if self.is_train:
            tokens_a = self.tokenizer.tokenize(text_a)
            if len(tokens_a) > self.max_seq_a_len - 2:
                tokens_a = tokens_a[:(self.max_seq_a_len - 2)]
        else:
            # fake tokens to generate masks
            tokens_a = [self.tokenizer.mask_token] * (self.max_seq_a_len - 2)

            gth_tokens_a = self.tokenizer.tokenize(text_a)
            if len(tokens_a) > self.max_seq_a_len - 2:
                tokens_a = tokens_a[:(self.max_seq_a_len - 2)]

            if len(gth_tokens_a) > self.max_seq_a_len:
                gth_tokens_a = gth_tokens_a[:(self.max_seq_a_len)]

            padding_gth_len = self.max_seq_a_len - len(gth_tokens_a)
            gth_tokens_a += [self.tokenizer.pad_token] * padding_gth_len
            gth_ids = self.tokenizer.convert_tokens_to_ids(gth_tokens_a)
            gth_ids = torch.tensor(gth_ids, dtype=torch.long)

        tokens = [self.tokenizer.cls_token] + tokens_a + [self.tokenizer.sep_token]
        segment_ids = [cls_token_segment_id] + [sequence_a_segment_id] * (len(tokens) - 1)
        seq_a_len = len(tokens)
        if text_b:
            # pad text_a to keep it in fixed length for better inference.
            padding_a_len = self.max_seq_a_len - seq_a_len
            tokens += [self.tokenizer.pad_token] * padding_a_len
            segment_ids += ([pad_token_segment_id] * padding_a_len)

            tokens_b = self.tokenizer.tokenize(text_b)
            if len(tokens_b) > self.max_seq_len - len(tokens) - 1:
                tokens_b = tokens_b[: (self.max_seq_len - len(tokens) - 1)]
            tokens += tokens_b + [self.tokenizer.sep_token]

            # tag_start = [i for i,x in enumerate(tokens) if x=="[SEP]"][-2]
            tmp_idx = [i for i,x in enumerate(tokens) if x=="[TAG]"]
            if len(tmp_idx) is not 0:
                tag_start = tmp_idx[-1]
                tokens[tag_start] = "[SEP]"
                tag_end = (tag_start + total_num_boxes) if (tag_start + total_num_boxes) < (self.max_seq_len-1) else (self.max_seq_len-1)
            else: # hacked
                # tag_start = (self.max_seq_len -1 ) -1
                tag_start = (len(tokens) -1 ) -5
                tag_end = tag_start + 1

            segment_ids += [sequence_b_segment_id] * (len(tokens_b) - total_num_boxes)
            segment_ids += [sequence_tag_segment_id] * (total_num_boxes + 1)

        seq_len = len(tokens)
        if self.is_train:

            masked_pos = torch.zeros(self.max_seq_len, dtype=torch.int)
            masked_tag_pos = torch.zeros(self.max_seq_len, dtype=torch.int)

        #-----------------------------mask response part-----------------------------

            # randomly mask words for prediction, ignore [CLS]
            candidate_masked_idx = list(range(1, seq_a_len)) # only mask text_a
            random.shuffle(candidate_masked_idx)
            num_masked = min(max(round(self.mask_response_prob * seq_a_len), 1), self.max_masked_tokens)
            num_masked = int(num_masked)
            masked_idx = candidate_masked_idx[:num_masked]
            masked_idx = sorted(masked_idx)
            masked_token = [tokens[i] for i in masked_idx]
            for pos in masked_idx:
                if random.random() <= 0.8:
                    # 80% chance to be a ['MASK'] token
                    tokens[pos] = self.tokenizer.mask_token
                elif random.random() <= 0.5:
                    # 10% chance to be a random word ((1-0.8)*0.5)
                    from random import randint
                    i = randint(0, len(self.tokenizer.vocab))
                    self.tokenizer._convert_id_to_token(i)
                    tokens[pos] = self.tokenizer._convert_id_to_token(i)
                else:
                    # 10% chance to remain the same (1-0.8-0.1)
                    pass

            masked_pos[masked_idx] = 1
            # # pad masked tokens to the same length
            # if num_masked < self.max_masked_tokens:
            #     masked_token = masked_token + ([self.tokenizer.pad_token] *
            #             (self.max_masked_tokens - num_masked))
            # masked_ids = self.tokenizer.convert_tokens_to_ids(masked_token)

        #-----------------------------mask context part-----------------------------
        # else:
            candidate_masked_idx_text_b = list(range(self.max_seq_a_len, tag_start )) #   mask context
            random.shuffle(candidate_masked_idx_text_b)
            num_masked_tb = min(max(round(self.mask_prob * (tag_start - self.max_seq_a_len)), 1), self.max_masked_tokens)
            num_masked_tb = int(num_masked_tb)
            masked_idx_tb = candidate_masked_idx_text_b[:num_masked_tb]
            masked_idx_tb = sorted(masked_idx_tb)
            masked_token_tb = [tokens[i] for i in masked_idx_tb]
            for pos in masked_idx_tb:
                if random.random() <= 0.8:
                    # 80% chance to be a ['MASK'] token
                    tokens[pos] = self.tokenizer.mask_token
                elif random.random() <= 0.5:
                    # 10% chance to be a random word ((1-0.8)*0.5)
                    from random import randint
                    i = randint(0, len(self.tokenizer.vocab))
                    self.tokenizer._convert_id_to_token(i)
                    tokens[pos] = self.tokenizer._convert_id_to_token(i)
                else:
                    # 10% chance to remain the same (1-0.8-0.1)
                    pass

            #  MRM with MLM in the same time
            masked_pos[masked_idx_tb] = 1
            # pad masked tokens to the same length
            total_num_masked = (num_masked  + num_masked_tb)
            if total_num_masked < self.max_masked_tokens:
                masked_token = masked_token  + masked_token_tb + ([self.tokenizer.pad_token] *
                        (self.max_masked_tokens - total_num_masked))
            masked_ids = self.tokenizer.convert_tokens_to_ids(masked_token)

            #-----------------------------mask object tag  part-----------------------------
        # else:
            candidate_masked_idx_tag = list(range(tag_start, tag_end )) #   mask text_b
            candidate_masked_idx_tag = [ min(i,len(tokens)-1) for i in candidate_masked_idx_tag]
            random.shuffle(candidate_masked_idx_tag)
            num_masked_tag = min(max(round(self.mask_prob * (tag_end - tag_start)), 1), self.max_masked_tokens)
            num_masked_tag = int(num_masked_tag)
            masked_idx_tag = candidate_masked_idx_tag[:num_masked_tag]
            masked_idx_tag = sorted(masked_idx_tag)
            masked_token_tag = [tokens[i] for i in masked_idx_tag]
            for pos in masked_idx_tag:
                if random.random() <= 0.8:
                    # 80% chance to be a ['MASK'] token
                    tokens[pos] = self.tokenizer.mask_token
                elif random.random() <= 0.5:
                    # 10% chance to be a random word ((1-0.8)*0.5)
                    from random import randint
                    i = randint(0, len(self.tokenizer.vocab))
                    self.tokenizer._convert_id_to_token(i)
                    tokens[pos] = self.tokenizer._convert_id_to_token(i)
                else:
                    # 10% chance to remain the same (1-0.8-0.1)
                    pass

            # way1: if conditioned Masked LM
            masked_tag_pos[masked_idx_tag] = 1
            # pad masked tokens to the same length
            if num_masked_tag < self.max_masked_tokens:
                masked_tag_token = masked_token_tag + ([self.tokenizer.pad_token] *
                        (self.max_masked_tokens - num_masked_tag))
            masked_tag_ids = self.tokenizer.convert_tokens_to_ids(masked_tag_token)

        else:
            masked_pos = torch.ones(self.max_seq_len, dtype=torch.int)

        # pad on the right for image captioning
        padding_len = self.max_seq_len - seq_len
        tokens = tokens + ([self.tokenizer.pad_token] * padding_len)
        segment_ids += ([pad_token_segment_id] * padding_len)
        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)

        # image features
        img_len = img_feat.shape[0]
        if img_len > self.max_img_seq_len:
            img_feat = img_feat[0 : self.max_img_seq_len, ]
            img_len = img_feat.shape[0]
        else:
            padding_matrix = torch.zeros((self.max_img_seq_len - img_len,
                                          img_feat.shape[1]))
            img_feat = torch.cat((img_feat, padding_matrix), 0)

        # prepare attention mask:
        # note that there is no attention from caption to image
        # because otherwise it will violate the triangle attention
        # for caption as caption will have full attention on image.
        max_len = self.max_seq_len + self.max_img_seq_len
        attention_mask = torch.zeros((max_len, max_len), dtype=torch.long)
        # C: caption, L: label, R: image region
        c_start, c_end = 0, seq_a_len
        l_start, l_end = self.max_seq_a_len, seq_len
        r_start, r_end = self.max_seq_len, self.max_seq_len + img_len
        # triangle mask for caption to caption C-C
        attention_mask[c_start : c_end, c_start : c_end].copy_(self._triangle_mask[0 : seq_a_len, 0 : seq_a_len])
        # full attention for L-L, R-R
        attention_mask[l_start : l_end, l_start : l_end] = 1
        attention_mask[r_start : r_end, r_start : r_end] = 1
        # full attention for C-L, C-R
        attention_mask[c_start : c_end, l_start : l_end] = 1
        attention_mask[c_start : c_end, r_start : r_end] = 1
        # full attention for L-R, R-L:
        attention_mask[l_start : l_end, r_start : r_end] = 1
        attention_mask[r_start : r_end, l_start : l_end] = 1
        # no attention for R-C, L-C

        input_ids = torch.tensor(input_ids, dtype=torch.long)
        segment_ids = torch.tensor(segment_ids, dtype=torch.long)

        # used for get tag position
        tag_pos = torch.zeros(self.max_seq_len, dtype=torch.int)
        tag_idx = list(range( tag_start,  tag_end))
        tag_pos[tag_idx] = 1

        if self.is_train:
            masked_ids = torch.tensor(masked_ids, dtype=torch.long)
            masked_tag_ids = torch.tensor(masked_tag_ids, dtype=torch.long)
            return (input_ids, attention_mask, segment_ids, img_feat, masked_pos, masked_ids, tag_pos, masked_tag_pos, masked_tag_ids)
        return (input_ids, attention_mask, segment_ids, img_feat, masked_pos, gth_ids)

def build_dataset(yaml_file, tokenizer, args, is_train=True):
    if not op.isfile(yaml_file):
        yaml_file = op.join(args.data_dir, yaml_file)
        assert op.isfile(yaml_file)

    if is_train:
        return CaptionTSVDataset(yaml_file, tokenizer=tokenizer,
            add_od_labels=args.add_od_labels, max_img_seq_length=args.max_img_seq_length,
            max_seq_length=args.max_seq_length, max_seq_a_length=args.max_seq_a_length,max_seq_c_length=args.max_seq_c_length,
            is_train=True, mask_prob=args.mask_prob, mask_response_prob=args.mask_response_prob, max_masked_tokens=args.max_masked_tokens)
    if args.use_cbs:
        dataset_class = CaptionTSVDatasetWithConstraints
    else:
        dataset_class = CaptionTSVDataset
    return dataset_class(yaml_file, tokenizer=tokenizer,
            add_od_labels=args.add_od_labels, max_img_seq_length=args.max_img_seq_length,
            max_seq_length=args.max_seq_length, max_seq_a_length=args.max_gen_length,
            is_train=False)

def save_checkpoint(model, tokenizer, args, epoch, global_step):
    checkpoint_dir = op.join(args.output_dir, 'checkpoint-{}-{}'.format(
        epoch, global_step))
    mkdir(checkpoint_dir)
    model_to_save = model.module if hasattr(model, 'module') else model
    save_num = 0
    while (save_num < 10):
        try:
            model_to_save.save_pretrained(checkpoint_dir)
            torch.save(args, op.join(checkpoint_dir, 'training_args.bin'))
            tokenizer.save_pretrained(checkpoint_dir)
            logger.info("Save checkpoint to {}".format(checkpoint_dir))
            break
        except:
            save_num += 1
    if save_num == 10:
        logger.info("Failed to save checkpoint after 10 trails.")
    return checkpoint_dir

def compute_score_with_logits(logits, labels):
    logits = torch.max(logits, -1)[1].data # argmax
    scores = logits == labels
    return scores

def train(args, train_dataset, val_dataset, test_dataset, model, tokenizer):
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
    train_sampler = RandomSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler,
            batch_size=args.train_batch_size, num_workers=args.num_workers)

    if args.max_steps > 0:
        t_total = args.max_steps
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // \
                args.gradient_accumulation_steps) + 1
    else:
        t_total = len(train_dataloader) // args.gradient_accumulation_steps \
                * args.num_train_epochs

    # Prepare optimizer and scheduler
    no_decay = ['bias', 'LayerNorm.weight']
    grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not \
            any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in model.named_parameters() if \
            any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]
    optimizer = AdamW(grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    if args.scheduler == "constant":
        scheduler = WarmupConstantSchedule(
                optimizer, warmup_steps=args.warmup_steps)
    elif args.scheduler == "linear":
        scheduler = WarmupLinearSchedule(
                optimizer, warmup_steps=args.warmup_steps, t_total=t_total)
    else:
        raise ValueError("Unknown scheduler type: {}".format(args.scheduler))

    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info("  Batch size per GPU = %d", args.per_gpu_train_batch_size)
    logger.info("  Total train batch size (w. parallel, & accumulation) = %d",
                   args.train_batch_size * args.gradient_accumulation_steps)
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    if args.scst:
        scst_criterion = ScstRewardCriterion()
        logger.info("  SCST training...")

    global_step, global_loss, global_acc =0,  0.0, 0.0
    global_tag_loss, global_tag_acc = 0.0, 0.0
    masked_global_loss, masked_global_acc =  0.0, 0.0
    # lm_global_loss, lm_global_acc =  0.0, 0.0
    model.zero_grad()
    eval_log = []
    eval_testset_log = []
    best_Bleu_1 = 0
    best_Bleu_2 = 0
    best_Bleu_3 = 0
    best_Bleu_4 = 0
    best_ROUGE_L = 0
    best_METEOR = 0
    best_SPICE = 0

    for epoch in range(int(args.num_train_epochs)):
        # for step, (img_keys, batch) in enumerate(train_dataloader):
        for step, ( batch) in enumerate(train_dataloader):
            batch = tuple(t.to(args.device) for t in batch)

            if not args.scst:
                model.train()
                inputs = {'input_ids': batch[0], 'attention_mask': batch[1],
                    'token_type_ids': batch[2], 'img_feats': batch[3],
                    'masked_pos': batch[4], 'masked_ids': batch[5], 'tag_pos': batch[6],'masked_tag_pos': batch[7], 'masked_tag_ids': batch[8]
                    # 'masked_pos': batch[4], 'masked_ids': batch[5] , 'response_pos': batch[6], 'response_ids': batch[7]
                }

                outputs = model(**inputs)
                masked_loss, masked_logits, MTP_loss, tag_logits, tag_ids = outputs[:5]
                loss = masked_loss + MTP_loss
                masked_ids = inputs['masked_ids']
                masked_ids = masked_ids[masked_ids != 0]
                batch_score = compute_score_with_logits(masked_logits, masked_ids)
                batch_acc = torch.sum(batch_score.float()) / torch.sum(inputs['masked_pos'])
                batch_tag_score = compute_score_with_logits(tag_logits, tag_ids)
                batch_tag_acc = torch.sum(batch_tag_score.float()) / torch.sum(inputs['masked_tag_pos'])

            if args.n_gpu > 1:
                loss = loss.mean() # mean() to average on multi-gpu parallel training
                masked_loss = masked_loss.mean() # mean() to average on multi-gpu parallel training
                MTP_loss = MTP_loss.mean() # mean() to average on multi-gpu parallel training
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
            global_loss += loss.item()
            global_acc += batch_acc
            global_tag_loss += MTP_loss.item()
            global_tag_acc += batch_tag_acc
            masked_global_loss += masked_loss.item()
            masked_global_acc += batch_acc
            if (step + 1) % args.gradient_accumulation_steps == 0:
                global_step += 1
                scheduler.step()
                optimizer.step()
                model.zero_grad()
                if global_step % args.logging_steps == 0:
                    logger.info("Epoch: {}, global_step: {}, lr: {:.6f}, loss: {:.4f} ({:.4f}), MTM loss: {:.4f} ({:.4f}), MLM(MRM+MCM) loss: {:.4f} ({:.4f}), " \
                        "MLM(MRM+MCM) score: {:.4f} ({:.4f}), MTM score: {:.4f} ({:.4f})".format(epoch, global_step,
                        optimizer.param_groups[0]["lr"], loss, global_loss / global_step, MTP_loss, global_tag_loss / global_step, masked_loss, masked_global_loss / global_step,
                        batch_acc, global_acc / global_step, batch_tag_acc, global_tag_acc / global_step)
                    )

                if (args.save_steps > 0 and global_step % args.save_steps == 0) or \
                        global_step == t_total:
                    checkpoint_dir = save_checkpoint(model, tokenizer, args, epoch, global_step)
                    # evaluation
                    if args.evaluate_during_training:
                        logger.info("Perform evaluation at step: %d" % (global_step))
                        evaluate_file = evaluate(args, val_dataset, model, tokenizer,
                                checkpoint_dir)
                        with open(evaluate_file, 'r') as f:
                            res = json.load(f)
                        best_Bleu_1 = max(best_Bleu_1, res['Bleu_1'])
                        best_Bleu_2 = max(best_Bleu_2, res['Bleu_2'])
                        best_Bleu_3 = max(best_Bleu_3, res['Bleu_3'])
                        best_Bleu_4 = max(best_Bleu_4, res['Bleu_4'])
                        best_ROUGE_L = max(best_ROUGE_L, res['ROUGE_L'])
                        best_METEOR = max(best_METEOR, res['METEOR'])
                        best_SPICE = max(best_SPICE, res['SPICE'])
                        res['epoch'] = epoch
                        res['global_step'] = step
                        res['best_Bleu_1'] = best_Bleu_1
                        res['best_Bleu_2'] = best_Bleu_2
                        res['best_Bleu_3'] = best_Bleu_3
                        res['best_Bleu_4'] = best_Bleu_4
                        res['best_ROUGE_L'] = best_ROUGE_L
                        res['best_METEOR'] = best_METEOR
                        res['best_SPICE'] = best_SPICE
                        eval_log.append(res)
                        with open(args.output_dir + '/eval_logs.json', 'w') as f:
                            json.dump(eval_log, f)
                    if args.evaluate_testset_during_training:
                        logger.info("Perform evaluation testset at step: %d" % (global_step))
                        evaluate_file = evaluate(args, test_dataset, model, tokenizer,
                                checkpoint_dir)
                        with open(evaluate_file, 'r') as f:
                            res = json.load(f)
                        best_Bleu_1 = max(best_Bleu_1, res['Bleu_1'])
                        best_Bleu_2 = max(best_Bleu_2, res['Bleu_2'])
                        best_Bleu_3 = max(best_Bleu_3, res['Bleu_3'])
                        best_Bleu_4 = max(best_Bleu_4, res['Bleu_4'])
                        best_ROUGE_L = max(best_ROUGE_L, res['ROUGE_L'])
                        best_METEOR = max(best_METEOR, res['METEOR'])
                        best_SPICE = max(best_SPICE, res['SPICE'])
                        res['epoch'] = epoch
                        res['global_step'] = step
                        res['best_Bleu_1'] = best_Bleu_1
                        res['best_Bleu_2'] = best_Bleu_2
                        res['best_Bleu_3'] = best_Bleu_3
                        res['best_Bleu_4'] = best_Bleu_4
                        res['best_ROUGE_L'] = best_ROUGE_L
                        res['best_METEOR'] = best_METEOR
                        res['best_SPICE'] = best_SPICE

                        eval_testset_log.append(res)
                        with open(args.output_dir + '/eval_testset_logs.json', 'w') as f:
                            json.dump(eval_testset_log, f)
    return global_step, global_loss / global_step

def scst_train_iter(args, train_dataset, model, scst_criterion, img_keys, batch, tokenizer):
    cls_token_id, sep_token_id, pad_token_id, mask_token_id = tokenizer.convert_tokens_to_ids(
            [tokenizer.cls_token, tokenizer.sep_token, tokenizer.pad_token,
                tokenizer.mask_token]
            )
    inputs = {'is_decode': True,
        'input_ids': batch[0], 'attention_mask': batch[1],
        'token_type_ids': batch[2], 'img_feats': batch[3],
        'masked_pos': batch[4],
        'do_sample': False,
        'bos_token_id': cls_token_id,
        'pad_token_id': pad_token_id,
        'eos_token_ids': [sep_token_id, pad_token_id],
        'mask_token_id': mask_token_id,
        # for adding od labels
        'add_od_labels': args.add_od_labels, 'od_labels_start_posid': args.max_seq_a_length,

        # hyperparameters of beam search
        'max_length': args.max_seq_a_length,
        'num_beams': 1,
        "temperature": args.temperature,
        "top_k": args.top_k,
        "top_p": args.top_p,
        "repetition_penalty": args.repetition_penalty,
        "length_penalty": args.length_penalty,
        "num_return_sequences": 1,
        "num_keep_best": 1,
    }

    model.eval()
    with torch.no_grad():
        greedy_res_raw, _ = model(**inputs)
        greedy_res_raw.squeeze_(1)  # batch_size * max_len

    model.train()
    inputs['do_sample'] = True
    sample_res_raw, sample_logprobs = model(**inputs)
    sample_res_raw.squeeze_(1)
    sample_logprobs.squeeze_(1)
    assert sample_logprobs.requires_grad == True
    assert sample_res_raw.requires_grad == False

    def _ids_to_captions(all_ids):
        captions = []
        for ids in all_ids:
            c = tokenizer.decode(ids.tolist(), skip_special_tokens=True)
            captions.append(c)
        return captions

    greedy_res = _ids_to_captions(greedy_res_raw)
    sample_res = _ids_to_captions(sample_res_raw)
    gt_res = [train_dataset.get_captions_by_key(k) for k in img_keys]

    loss = scst_criterion(gt_res, greedy_res, sample_res, sample_logprobs)
    return loss

def get_predict_file(output_dir, yaml_file, args):
    cc = ['pred']
    # make sure it works with/without / in end of the path.
    data = op.basename(op.join(args.data_dir, '')[:-1])
    split = op.basename(yaml_file)
    assert split.endswith('.yaml')
    split = split[:-5]
    cc.append(data)
    cc.append(split)
    cc.append('beam{}'.format(args.num_beams))
    cc.append('sample{}'.format(args.do_sample))
    cc.append('top_p{}'.format(args.top_p))
    cc.append('top_k{}'.format(args.top_k))
    cc.append('temperature{}'.format(args.temperature))
    cc.append('repetition_penalty{}'.format(args.repetition_penalty))
    cc.append('max{}'.format(args.max_gen_length))
    if args.add_od_labels:
        cc.append('odlabels')
    if args.num_keep_best != 1:
        cc.append('best{}'.format(args.num_keep_best))
    if args.use_cbs:
        cc.append('cbs{}'.format(args.min_constraints_to_satisfy))
    if args.output_hidden_states:
        cc.append('hidden')
    return op.join(output_dir, '{}.tsv'.format('.'.join(cc)))

def get_evaluate_file(predict_file):
    assert predict_file.endswith('.tsv')
    fpath = op.splitext(predict_file)[0]
    return fpath + '.eval.json'

def get_evaluate_method(predict_file):
    if 'nocaps' in op.basename(predict_file):
        return 'nocaps'
    else:
        return 'coco'

def evaluate(args, val_dataset, model, tokenizer, output_dir):
    assert op.isdir(output_dir)
    predict_file = get_predict_file(output_dir, val_dataset.yaml_file, args)
    if op.isfile(predict_file):
        logger.info('Skip predict. {} already exists'.format(predict_file))
    else:
        test(args, val_dataset, model, tokenizer, predict_file)

    evaluate_file = get_evaluate_file(predict_file)
    if op.isfile(evaluate_file):
        logger.info('Skip evaluation. {} already exists'.format(evaluate_file))
        return evaluate_file

    eval_method = get_evaluate_method(predict_file)
    if eval_method == 'coco':
        # gt_file = val_dataset.get_caption_file_in_coco_format()
        # result = evaluate_on_coco_caption(predict_file, gt_file, outfile=evaluate_file)

        gt_file = val_dataset.get_dialog_file()
        result = evaluate_on_reddit_caption(predict_file, gt_file, outfile=evaluate_file)
        #for evaluating the DialoGPT result
        # result = evaluate_on_reddit_caption(predict_file+'.txt', gt_file, outfile=evaluate_file)
    else:
        split = 'val' if 'val' in op.basename(val_dataset.yaml_file) else 'test'
        result = evaluate_on_nocaps(split, predict_file,
                    data_dir=args.data_dir, evaluate_file=evaluate_file)
    logger.info("evaluation result: {}".format(str(result)))
    return evaluate_file

def test(args, test_dataset, model, tokenizer, predict_file):
    args.test_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
    test_sampler = SequentialSampler(test_dataset)
    cache_file = predict_file

    test_dataloader = DataLoader(test_dataset, sampler=test_sampler,
            batch_size=args.test_batch_size, num_workers=args.num_workers)

    cls_token_id, sep_token_id, pad_token_id, mask_token_id, period_token_id = \
        tokenizer.convert_tokens_to_ids( [tokenizer.cls_token,
            tokenizer.sep_token, tokenizer.pad_token, tokenizer.mask_token, '.']
        )
    model.eval()

    def gen_rows():
        time_meter = 0

        with torch.no_grad():

            for step, batch in tqdm(enumerate(test_dataloader), ncols=100,desc="Testing Forward: ", total=len(test_dataloader)):

                batch = tuple(t.to(args.device) for t in batch)
                inputs = {'is_decode': True,
                    'input_ids': batch[0], 'attention_mask': batch[1],
                    'token_type_ids': batch[2], 'img_feats': batch[3],
                    'masked_pos': batch[4],
                    'gth_ids': batch[5],
                    'do_sample': args.do_sample,
                    'bos_token_id': cls_token_id,
                    'pad_token_id': pad_token_id,
                    # 'eos_token_ids': [sep_token_id, pad_token_id, period_token_id],
                    'eos_token_ids': [sep_token_id, pad_token_id],
                    'mask_token_id': mask_token_id,
                    # for adding od labels
                    'add_od_labels': args.add_od_labels, 'od_labels_start_posid': args.max_seq_a_length,

                    # hyperparameters of beam search
                    'max_length': args.max_gen_length,
                    'num_beams': args.num_beams,
                    "temperature": args.temperature,
                    "top_k": args.top_k,
                    "top_p": args.top_p,
                    "repetition_penalty": args.repetition_penalty,
                    "length_penalty": args.length_penalty,
                    "num_return_sequences": args.num_return_sequences,
                    "num_keep_best": args.num_keep_best,
                }
                if args.use_cbs:
                    inputs.update({'use_cbs': True,
                        'fsm': batch[5],
                        'num_constraints': batch[6],
                        'min_constraints_to_satisfy': args.min_constraints_to_satisfy,
                    })
                tic = time.time()
                # captions, logprobs
                outputs = model(**inputs)
                time_meter += time.time() - tic
                all_caps = outputs[0]  # [batch_size, num_keep_best, max_len]
                all_confs = torch.exp(outputs[1])

                for  caps, confs in zip( all_caps, all_confs):
                    res = []
                    for cap, conf in zip(caps, confs):
                        cap = tokenizer.decode(cap.tolist(), skip_special_tokens=True)

                        cap = cap.split('.')[0]
                        res.append({'response': cap, 'conf': conf.item()})

                    # yield json.dumps(res)
                    yield res
                    # return res

        logger.info("Inference model computing time: {} seconds per batch".format(time_meter / (step+1)))

    tsv_writer(gen_rows(), cache_file)

    return predict_file

def restore_training_settings(args):
    assert not args.do_train
    assert args.do_test or args.do_eval
    # restore training settings, check hasattr for backward compatibility
    train_args = torch.load(op.join(args.eval_model_dir, 'training_args.bin'))
    if hasattr(train_args, 'max_seq_a_length'):
        max_od_labels_len = train_args.max_seq_length - train_args.max_seq_a_length
        max_seq_length = args.max_gen_length + max_od_labels_len
        args.max_seq_length = max_seq_length
        logger.warning('Override max_seq_length to {} = max_gen_length:{} + od_labels_len:{}'.format(
                max_seq_length, args.max_gen_length, max_od_labels_len))

    override_params = ['max_seq_a_length', 'do_lower_case', 'add_od_labels',
            'max_img_seq_length', 'img_feature_dim',
            'img_feature_type']
    for param in override_params:
        if hasattr(train_args, param):
            train_v = getattr(train_args, param)
            test_v = getattr(args, param)
            if train_v != test_v:
                logger.warning('Override {} with train args: {} -> {}'.format(param,
                    test_v, train_v))
                setattr(args, param, train_v)
    return args

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", default='/data/data_for_oscar/reddit_oi', type=str, required=False,
                        help="The input data dir with all required files.")
    parser.add_argument("--tag_vocab_file", default='/data/data_for_oscar/reddit_oi/objects_vocab.txt', type=str, required=False,
                        help="The files for saving the tag vocabulary.")
    parser.add_argument("--train_yaml", default='train.yaml', type=str, required=False,
                        help="yaml file for training.")
    parser.add_argument("--test_yaml", default='test.yaml', type=str, required=False,
                        help="yaml file for testing.")
    parser.add_argument("--val_yaml", default='val.yaml', type=str, required=False,
                        help="yaml file used for validation during training.")
    parser.add_argument("--model_name_or_path", default=None, type=str, required=False,
                        help="Path to pre-trained model or model type.")
    parser.add_argument("--output_dir", default='output/', type=str, required=False,
                        help="The output directory to save checkpoint and test results.")
    parser.add_argument("--loss_type", default='sfmx', type=str,
                        help="Loss function types: support kl, x2, sfmx")
    parser.add_argument("--config_name", default="", type=str,
                        help="Pretrained config name or path if not the same as model_name.")
    parser.add_argument("--tokenizer_name", default="", type=str,
                        help="Pretrained tokenizer name or path if not the same as model_name.")
    parser.add_argument("--max_seq_length", default=120, type=int,
                        help="The maximum total input sequence length after tokenization. "
                             "Sequences longer than this will be truncated, "
                             "sequences shorter will be padded.")
    parser.add_argument("--max_seq_a_length", default=40, type=int,
                        help="The maximum sequence length for caption.")
    parser.add_argument("--max_seq_c_length", default=80, type=int, #dump, useless args now
                        help="The maximum sequence length for dialog context.")
    parser.add_argument("--do_train", action='store_true', help="Whether to run training.")
    parser.add_argument("--do_sample", action='store_true', help="Whether to do sampling.")
    parser.add_argument("--do_test", action='store_true', help="Whether to run inference.")
    parser.add_argument("--do_eval", action='store_true', help="Whether to run evaluation.")
    parser.add_argument("--do_lower_case", action='store_true',
                        help="Set this flag if you are using an uncased model.")
    parser.add_argument("--mask_prob", default=0.15, type=float,
                        help= "Probability to mask input sentence during training.")
    parser.add_argument("--mask_response_prob", default=0.7, type=float,
                        help= "Probability to mask input sentence during training.")
    parser.add_argument("--max_masked_tokens", type=int, default=300,
                        help="The max number of masked tokens per sentence.")
    parser.add_argument("--add_od_labels", default=False, action='store_true',
                        help="Whether to add object detection labels or not")
    parser.add_argument("--drop_out", default=0.1, type=float, help="Drop out in BERT.")
    parser.add_argument("--max_img_seq_length", default=50, type=int,
                        help="The maximum total input image sequence length.")
    parser.add_argument("--img_feature_dim", default=2048, type=int,
                        help="The Image Feature Dimension.")
    parser.add_argument("--type_vocab_size", default=3, type=int,
                        help="segment embedding vocabulary size.")
    parser.add_argument("--img_feature_type", default='frcnn', type=str,
                        help="Image feature type.")
    parser.add_argument("--per_gpu_train_batch_size", default=64, type=int,
                        help="Batch size per GPU/CPU for training.")
    parser.add_argument("--per_gpu_eval_batch_size", default=6, type=int,
                        help="Batch size per GPU/CPU for evaluation.")
    parser.add_argument("--output_mode", default='classification', type=str,
                        help="output mode, support classification or regression.")
    parser.add_argument("--num_labels", default=2, type=int,
                        help="num_labels is 2 for classification and 1 for regression.")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before backward.")
    parser.add_argument("--learning_rate", default=3e-5, type=float, help="The initial lr.")
    parser.add_argument("--weight_decay", default=0.05, type=float, help="Weight deay.")
    parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
    parser.add_argument("--warmup_steps", default=0, type=int, help="Linear warmup.")
    parser.add_argument("--scheduler", default='linear', type=str, help="constant or linear or")
    parser.add_argument("--num_workers", default=4, type=int, help="Workers in dataloader.")
    parser.add_argument("--num_train_epochs", default=40, type=int,
                        help="Total number of training epochs to perform.")
    parser.add_argument("--max_steps", default=-1, type=int,
                        help="Total number of training steps. Override num_train_epochs.")
    parser.add_argument('--logging_steps', type=int, default=10, help="Log every X steps.")
    parser.add_argument('--save_steps', type=int, default=-1,
                        help="Save checkpoint every X steps. Will also perform evaluatin.")
    parser.add_argument("--evaluate_during_training", action='store_true',
                        help="Run evaluation during training at each save_steps.")
    parser.add_argument("--evaluate_testset_during_training", action='store_true',
                        help="Run evaluation testset during training at each save_steps.")
    parser.add_argument("--no_cuda", action='store_true', help="Avoid using CUDA.")
    parser.add_argument('--seed', type=int, default=88, help="random seed for initialization.")
    parser.add_argument('--scst', action='store_true', help='Self-critical sequence training')
    # for generation
    parser.add_argument("--eval_model_dir", type=str, default='',
                        help="Model directory for evaluation.")
    parser.add_argument('--max_gen_length', type=int, default=40,
                        help="max length of generated sentences")
    parser.add_argument('--output_hidden_states', action='store_true',
                        help="Turn on for fast decoding")
    parser.add_argument('--num_return_sequences', type=int, default=1,
                        help="repeating times per image")
    parser.add_argument('--num_beams', type=int, default=5, help="beam search width")
    parser.add_argument('--num_keep_best', type=int, default=1,
                        help="number of hypotheses to keep in beam search")
    parser.add_argument('--temperature', type=float, default=1,
                        help="temperature in softmax for sampling")
    parser.add_argument('--top_k', type=int, default=0,
                        help="filter distribution for sampling")
    parser.add_argument('--top_p', type=float, default=1,
                        help="filter distribution for sampling")
    parser.add_argument('--repetition_penalty', type=int, default=3,
                        help="repetition penalty from CTRL paper (https://arxiv.org/abs/1909.05858)")
    parser.add_argument('--length_penalty', type=float, default=0.3,#1,
                        help="beam search length penalty")
    # for Constrained Beam Search
    parser.add_argument('--use_cbs', action='store_true',
                        help='Use constrained beam search for decoding')
    parser.add_argument('--min_constraints_to_satisfy', type=int, default=2,
                        help="minimum number of constraints to satisfy")
    args = parser.parse_args()

    global logger

    args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
    args.n_gpu = torch.cuda.device_count()

    output_dir = args.output_dir
    mkdir(output_dir)

    logger = setup_logger("vlpretrain", output_dir, 0)
    logger.warning("Device: %s, n_gpu: %s", args.device, args.n_gpu)
    set_seed(args.seed, args.n_gpu)

    # Load pretrained model and tokenizer
    config_class, model_class, tokenizer_class = BertConfig, BertForMultiModalDialog, BertTokenizer
    if args.do_train:
        assert args.model_name_or_path is not None
        config = config_class.from_pretrained(args.config_name if args.config_name else \
                args.model_name_or_path, num_labels=args.num_labels, finetuning_task='image_captioning')
        if args.scst:
            # avoid using too much memory
            config.output_hidden_states = True
        tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name \
                else args.model_name_or_path, do_lower_case=args.do_lower_case)
        #---------------add by Jokie for add object tags to the vocab ------------------------
        config.img_feature_dim = args.img_feature_dim
        config.img_feature_type = args.img_feature_type
        config.hidden_dropout_prob = args.drop_out
        config.loss_type = args.loss_type
        config.type_vocab_size = args.type_vocab_size
        model = model_class.from_pretrained(args.model_name_or_path,
                from_tf=bool('.ckpt' in args.model_name_or_path), config=config)
        print(args.train_yaml)

        TAGS_TOKENS_LIST = [line.decode().rstrip('\n').split(',')[0] for line in open(args.tag_vocab_file,'rb')]
        print('len(TAGS_TOKENS_LIST):', len(TAGS_TOKENS_LIST))
        num_add_tokens = tokenizer.add_tokens(TAGS_TOKENS_LIST)
        SPECIAL_TOKENS = ['[TAG]']
        SPECIAL_TOKENS_DICT = {'additional_special_tokens':'[TAG]'}
        num_add_sep_tokens = tokenizer.add_special_tokens(SPECIAL_TOKENS_DICT)
        print('num_add_tokens:', num_add_tokens)
        print('num_add_sep_tokens:', num_add_sep_tokens)
        print('num_tokens:', len(tokenizer))
        model.resize_token_embeddings(len(tokenizer))
        model.bert._resize_token_embeddings(len(tokenizer))
        bert_embedding_weight = model.bert.embeddings.word_embeddings.weight
        print('after resized decoder_shape: ', bert_embedding_weight.size(1))
        print('after resized decoder_shape1: ', bert_embedding_weight.size(0))
        model.decoder = nn.Linear(bert_embedding_weight.size(1),
                            bert_embedding_weight.size(0), bias=False)
        model.vocab_size = bert_embedding_weight.size(0)

    else:
        checkpoint = args.eval_model_dir
        assert op.isdir(checkpoint)
        config = config_class.from_pretrained(checkpoint)
        config.output_hidden_states = args.output_hidden_states
        tokenizer = tokenizer_class.from_pretrained(checkpoint)
        logger.info("Evaluate the following checkpoint: %s", checkpoint)
        model = model_class.from_pretrained(checkpoint, config=config)
        print('num_tokens for decoding when testing:', len(tokenizer))

    model.to(args.device)
    logger.info("Training/evaluation parameters %s", args)
    if args.do_train:
        train_dataset = build_dataset(op.join(args.data_dir, args.train_yaml), tokenizer, args)
        val_dataset = build_dataset(op.join(args.data_dir, args.val_yaml),
                tokenizer, args, is_train=False)
        test_dataset = build_dataset(op.join(args.data_dir, args.test_yaml),
                tokenizer, args, is_train=False)
        global_step, avg_loss = train(args, train_dataset, val_dataset, test_dataset, model, tokenizer)
        logger.info("Training done: total_step = %s, avg loss = %s", global_step, avg_loss)

    # inference and evaluation
    if args.do_test or args.do_eval:
        args = restore_training_settings(args)
        test_dataset = build_dataset(op.join(args.data_dir, args.test_yaml),
                tokenizer, args, is_train=False)
        if args.n_gpu > 1:
            model = torch.nn.DataParallel(model)

        if not args.do_eval:
            predict_file = get_predict_file(checkpoint, test_dataset.yaml_file, args)
            test(args, test_dataset, model, tokenizer, predict_file)
            logger.info("Prediction results saved to: {}".format(predict_file))
        else:
            evaluate_file = evaluate(args, test_dataset, model, tokenizer,
                    checkpoint)
            logger.info("Evaluation results saved to: {}".format(evaluate_file))

if __name__ == "__main__":
    main()


import json

# Open the files and load the JSON data
with open('StreamingHistory0.json', encoding='utf-8') as file:
    data0 = json.load(file)
with open('StreamingHistory1.json', encoding='utf-8') as file:
    data1 = json.load(file)
with open('StreamingHistory2.json', encoding='utf-8') as file:
    data2 = json.load(file)
with open('StreamingHistory3.json', encoding='utf-8') as file:
    data3 = json.load(file)

# Combine the data from all four files into a single list
data = data0 + data1 + data2 + data3

# Create a dictionary to store the total msPlayed for each artist
totals = {}

# Loop through each item in the data and add the msPlayed to the artist's total
for item in data:
    artist = item['artistName']
    ms_played = item['msPlayed']
    if artist in totals:
        totals[artist] += ms_played
    else:
        totals[artist] = ms_played

# Sort the artists in descending order by total msPlayed
sorted_artists = sorted(totals, key=totals.get, reverse=True)

# Write the results to a file
with open('outputfile.txt', 'w') as file:
    for artist in sorted_artists:
        file.write(f"{artist}: {totals[artist]}\n")


import paddle.fluid as fluid

from paddlerec.core.utils import envs
from paddlerec.core.model import Model as ModelBase

class Model(ModelBase):
    def __init__(self, config):
        ModelBase.__init__(self, config)

    def config_read(self, config_path):
        with open(config_path, "r") as fin:
            user_count = int(fin.readline().strip())
            item_count = int(fin.readline().strip())
            cat_count = int(fin.readline().strip())
        return user_count, item_count, cat_count

    def din_attention(self, hist, target_expand, mask):
        """activation weight"""

        hidden_size = hist.shape[-1]

        concat = fluid.layers.concat(
            [hist, target_expand, hist - target_expand, hist * target_expand],
            axis=2)
        atten_fc1 = fluid.layers.fc(name="atten_fc1",
                                    input=concat,
                                    size=80,
                                    act=self.act,
                                    num_flatten_dims=2)
        atten_fc2 = fluid.layers.fc(name="atten_fc2",
                                    input=atten_fc1,
                                    size=40,
                                    act=self.act,
                                    num_flatten_dims=2)
        atten_fc3 = fluid.layers.fc(name="atten_fc3",
                                    input=atten_fc2,
                                    size=1,
                                    num_flatten_dims=2)
        atten_fc3 += mask
        atten_fc3 = fluid.layers.transpose(x=atten_fc3, perm=[0, 2, 1])
        atten_fc3 = fluid.layers.scale(x=atten_fc3, scale=hidden_size**-0.5)
        weight = fluid.layers.softmax(atten_fc3)
        out = fluid.layers.matmul(weight, hist)
        out = fluid.layers.reshape(x=out, shape=[0, hidden_size])
        return out

    def train_net(self):
        seq_len = -1
        self.item_emb_size = envs.get_global_env("hyper_parameters.item_emb_size", 64, self._namespace)
        self.cat_emb_size = envs.get_global_env("hyper_parameters.cat_emb_size", 64, self._namespace)
        self.act = envs.get_global_env("hyper_parameters.act", "sigmoid", self._namespace)
        #item_emb_size = 64
        #cat_emb_size = 64
        self.is_sparse = envs.get_global_env("hyper_parameters.is_sparse", False, self._namespace)
        #significant for speeding up the training process
        self.config_path = envs.get_global_env("hyper_parameters.config_path", "data/config.txt", self._namespace)
        self.use_DataLoader = envs.get_global_env("hyper_parameters.use_DataLoader", False, self._namespace)
        user_count, item_count, cat_count = self.config_read(self.config_path)

        item_emb_attr = fluid.ParamAttr(name="item_emb")
        cat_emb_attr = fluid.ParamAttr(name="cat_emb")

        hist_item_seq = fluid.data(
            name="hist_item_seq", shape=[None, seq_len], dtype="int64")
        self._data_var.append(hist_item_seq)

        hist_cat_seq = fluid.data(
            name="hist_cat_seq", shape=[None, seq_len], dtype="int64")
        self._data_var.append(hist_cat_seq)

        target_item = fluid.data(name="target_item", shape=[None], dtype="int64")
        self._data_var.append(target_item)

        target_cat = fluid.data(name="target_cat", shape=[None], dtype="int64")
        self._data_var.append(target_cat)

        label = fluid.data(name="label", shape=[None, 1], dtype="float32")
        self._data_var.append(label)

        mask = fluid.data(name="mask", shape=[None, seq_len, 1], dtype="float32")
        self._data_var.append(mask)

        target_item_seq = fluid.data(
            name="target_item_seq", shape=[None, seq_len], dtype="int64")
        self._data_var.append(target_item_seq)

        target_cat_seq = fluid.data(
            name="target_cat_seq", shape=[None, seq_len], dtype="int64")
        self._data_var.append(target_cat_seq)

        if self.use_DataLoader:
            self._data_loader = fluid.io.DataLoader.from_generator(
                feed_list=self._data_var, capacity=10000, use_double_buffer=False, iterable=False)

        hist_item_emb = fluid.embedding(
            input=hist_item_seq,
            size=[item_count, self.item_emb_size],
            param_attr=item_emb_attr,
            is_sparse=self.is_sparse)

        hist_cat_emb = fluid.embedding(
            input=hist_cat_seq,
            size=[cat_count, self.cat_emb_size],
            param_attr=cat_emb_attr,
            is_sparse=self.is_sparse)

        target_item_emb = fluid.embedding(
            input=target_item,
            size=[item_count, self.item_emb_size],
            param_attr=item_emb_attr,
            is_sparse=self.is_sparse)

        target_cat_emb = fluid.embedding(
            input=target_cat,
            size=[cat_count, self.cat_emb_size],
            param_attr=cat_emb_attr,
            is_sparse=self.is_sparse)

        target_item_seq_emb = fluid.embedding(
            input=target_item_seq,
            size=[item_count, self.item_emb_size],
            param_attr=item_emb_attr,
            is_sparse=self.is_sparse)

        target_cat_seq_emb = fluid.embedding(
            input=target_cat_seq,
            size=[cat_count, self.cat_emb_size],
            param_attr=cat_emb_attr,
            is_sparse=self.is_sparse)

        item_b = fluid.embedding(
            input=target_item,
            size=[item_count, 1],
            param_attr=fluid.initializer.Constant(value=0.0))

        hist_seq_concat = fluid.layers.concat([hist_item_emb, hist_cat_emb], axis=2)
        target_seq_concat = fluid.layers.concat(
            [target_item_seq_emb, target_cat_seq_emb], axis=2)
        target_concat = fluid.layers.concat(
            [target_item_emb, target_cat_emb], axis=1)

        out = self.din_attention(hist_seq_concat, target_seq_concat, mask)
        out_fc = fluid.layers.fc(name="out_fc",
                                input=out,
                                size=self.item_emb_size + self.cat_emb_size,
                                num_flatten_dims=1)
        embedding_concat = fluid.layers.concat([out_fc, target_concat], axis=1)

        fc1 = fluid.layers.fc(name="fc1",
                            input=embedding_concat,
                            size=80,
                            act=self.act)
        fc2 = fluid.layers.fc(name="fc2", input=fc1, size=40, act=self.act)
        fc3 = fluid.layers.fc(name="fc3", input=fc2, size=1)
        logit = fc3 + item_b

        loss = fluid.layers.sigmoid_cross_entropy_with_logits(x=logit, label=label)

        avg_loss = fluid.layers.mean(loss)
        self._cost = avg_loss

        self.predict = fluid.layers.sigmoid(logit)
        predict_2d = fluid.layers.concat([1 - self.predict, self.predict], 1)
        label_int = fluid.layers.cast(label, 'int64')
        auc_var, batch_auc_var, _ = fluid.layers.auc(input=predict_2d,
                                                            label=label_int,
                                                            slide_steps=0)
        self._metrics["AUC"] = auc_var
        self._metrics["BATCH_AUC"] = batch_auc_var

    def optimizer(self):
        learning_rate = envs.get_global_env("hyper_parameters.learning_rate", None, self._namespace)
        optimizer = fluid.optimizer.Adam(learning_rate, lazy_mode=True)
        return optimizer

    def infer_net(self, parameter_list):
        self.deepfm_net()