# import packages
import pandas as pd
import ixmp
import message_ix
import os
import zipfile
import shutil

# define model, scenario and zipfile name
model = 'SSP2_test'
scenario = 'baseline_DEFAULT_2020'
zipfilename = 'MESSAGEix-GLOBIOM_SSP2.zip'

# connect to platform and load scenario
mp = ixmp.Platform()
scen = message_ix.Scenario(mp, model=model, scenario=scenario)

# create temporary directory to write csv files into
tmppath = 'tmp'
if not os.path.exists(tmppath):
    os.makedirs(tmppath)
os.chdir(tmppath)
    
# check which sets exist
sets = scen.set_list()
# check which parameters exist
parameters = scen.par_list()    

# create zipfile and open for writing
message_zipfile = zipfile.ZipFile(zipfilename, mode='w', compression=zipfile.ZIP_DEFLATED)

# loop over sets, write to csv and to zip archive
for s in sets:
    # read set into dataframe
    df = scen.set(s)
    # write set dataframe to csv file 
    df.to_csv('set_%s.csv' % s, index = False)
    message_zipfile.write("set_%s.csv" % s)

# loop over parameters, write to csv and to zip archive
for p in parameters:
    # read parameter into dataframe
    df = scen.par(p)
    # write parameter dataframe to csv file 
    df.to_csv('par_%s.csv' % p, index = False)
    message_zipfile.write("par_%s.csv" % p)
    
# close zipfile
message_zipfile.close()

# move zipfile to parent directory 
os.rename(zipfilename, '../' + zipfilename)
os.chdir('..')

# delete temporary directory and files
shutil.rmtree(tmppath)