import sys
import glob
import numpy as np
import pandas as pd
from enum import Enum

class G_enum(Enum):
    TOTAL_RESIZES = 0
    TOTAL_GROUPS = 1
    SDR = 2
    ADR = 3
    DR = 4
    RED_METHOD = 5
    RED_STRATEGY = 6
    SPAWN_METHOD = 7
    SPAWN_STRATEGY = 8
    GROUPS = 9
    ITERS = 10
    T_SPAWN = 11
    T_SR = 12
    T_AR = 13
    T_MALLEABILITY = 14
    T_TOTAL = 15
    #Malleability specific
    NP = 0
    NC = 1


#columnsG = ["Total_Resizes", "Total_Groups", "SDR", "ADR", "DR", "Redistribution_Method", \
#            "Redistribution_Strategy", "Spawn_Method", "Spawn_Strategy", "Groups", \
#            "Iters", "T_spawn", "T_SR", "T_AR", "T_Malleability", "T_total"] #16


columnsM = ["NP", "NC", "SDR", "ADR", "DR", "Redistribution_Method", \
            "Redistribution_Strategy", "Spawn_Method", "Spawn_Strategy",  \
            "Iters", "T_spawn", "T_SR", "T_AR", "T_Malleability"] #15

def copy_resize(row, dataM_it, resize):
  basic_indexes = [G_enum.SDR.value, G_enum.ADR.value, G_enum.DR.value]
  array_actual_group = [G_enum.ITERS.value,  \
          G_enum.T_SPAWN.value, G_enum.T_SR.value, \
          G_enum.T_AR.value, G_enum.T_MALLEABILITY.value]
  array_next_group = [G_enum.RED_METHOD.value, G_enum.RED_STRATEGY.value, \
          G_enum.SPAWN_METHOD.value, G_enum.SPAWN_STRATEGY.value]

  dataM_it[G_enum.NP.value] = row[G_enum.GROUPS.value][resize]
  dataM_it[G_enum.NC.value] = row[G_enum.GROUPS.value][resize+1]

  for index in basic_indexes:
    dataM_it[index] = row[index]
    
  for index in array_actual_group:
    dataM_it[index-1] = row[index][resize] 

  for index in array_next_group:
    dataM_it[index] = row[index][resize+1]


#-----------------------------------------------

def create_resize_dataframe(dfG, dataM):
  it = -1
  for row_index in range(len(dfG)):
    row = dfG.iloc[row_index]
    resizes = row[G_enum.TOTAL_RESIZES.value]

    for resize in range(resizes):
      it += 1
      dataM.append( [None] * len(columnsM) )
      copy_resize(row, dataM[it], resize)

#-----------------------------------------------
if len(sys.argv) < 2:
    print("The files name is missing\nUsage: python3 CreateResizeDataframe.py input_file.pkl output_name")
    exit(1)

input_name = sys.argv[1]
if len(sys.argv) > 2:
  name = sys.argv[2]
else:
  name = "dataM"
print("File name will be: " + name + ".pkl")


dfG = pd.read_pickle(input_name)
dataM = []
create_resize_dataframe(dfG, dataM)

dfM = pd.DataFrame(dataM, columns=columnsM)
dfM.to_pickle(name + '.pkl')

print(dfG)
print(dfM)