import os
from pyomo.environ import (
    Var,
    Constraint,
    Param,
    Set,
    Objective,
    Binary,
    ConcreteModel,
    SolverFactory,
    RangeSet
)

def _createPyomoModel(nVar):
    """
    Construct a simple pyomo model with 1 variable, 1 constraint
    The model is fixed to remove the solver time and quickly generate the gdx file

    Args:
        nVar (int): Number of 

    Returns:
        None
    """

    m = ConcreteModel()

    m.dummySet1 = Set(initialize=RangeSet(nVar))

    m.dummyVar1 = Var(
        m.dummySet1,
        within=Binary,
    )

    @m.Constraint(m.dummySet1)
    def dummyConstraint2(m, i, j):
        return m.dummyVar1[i] >= 1
 
    @m.Objective()
    def dummyObj(m):
        return sum(m.dummyVar1[i] for i in m.dummySet1)
    
    m.dummyVar1.fix(1)

    # reference: https://www.gams.com/latest/docs/UG_OptionStatement.html
    opt=SolverFactory('gams')
    add_options = ['option iterLim=1;'] # iteration limit
    add_options = ['option resLim=1;'] # solver time

    status = opt.solve(m,
                tee=False,
                keepfiles=True,
                tmpdir=os.getcwd()+f'/{nVar}',
                add_options=add_options)
        
def main():
    num_vars = [10, 100, 1000, 10000, 100000, 1000000, 5000000, 10000000]
    for n in num_vars:
        _createPyomoModel(n)
        
main()

import struct
import sys
import time
import os
import re

import numpy as np
import matplotlib.pyplot as plt
from gams.core import gdx

def gams_like(file_path):
    model_soln = {}

    system_directory = "/Library/Frameworks/GAMS.framework/Versions/49/Resources/"
    gdxHandle = gdx.new_gdxHandle_tp()
    gdx.gdxCreateD(gdxHandle, system_directory, gdx.GMS_SSSIZE)
    gdx.gdxOpenRead(gdxHandle, file_path)
    ret, symCount, _ = gdx.gdxSystemInfo(gdxHandle)

    # only setting SVs
    specVals = gdx.doubleArray(gdx.GMS_SVIDX_MAX)
    rc = gdx.gdxGetSpecialValues(gdxHandle, specVals)

    specVals[gdx.GMS_SVIDX_EPS] = sys.float_info.min
    specVals[gdx.GMS_SVIDX_UNDEF] = float("nan")
    specVals[gdx.GMS_SVIDX_PINF] = float("inf")
    specVals[gdx.GMS_SVIDX_MINF] = float("-inf")
    specVals[gdx.GMS_SVIDX_NA] = struct.unpack(">d", bytes.fromhex("fffffffffffffffe"))[
        0
    ]
    gdx.gdxSetSpecialValues(gdxHandle, specVals)
    # end set SVs

    start = time.time()
    ret, symCount, _ = gdx.gdxSystemInfo(gdxHandle)
    for n in range(1, symCount + 1):
        ret, sym_name, typ, subtype = gdx.gdxSymbolInfo(gdxHandle, n)
        if not ret:
            raise RuntimeError("GAMS GDX failure (gdxSymbolInfo).")

        ret, nrRecs = gdx.gdxDataReadRawStart(gdxHandle, n)
        if not ret:
            raise RuntimeError("GAMS GDX failure (gdxDataReadRawStart).")

        ret, keys, vals, dimfirst = gdx.gdxDataReadRaw(gdxHandle)
        if not ret:
            raise RuntimeError("GAMS GDX failure (gdxDataReadRaw).")

        ret = gdx.gdxDataReadDone(gdxHandle)
        if not ret:
            raise RuntimeError("GAMS GDX failure (gdxDataReadDone).")

        model_soln[sym_name] = (vals[0], vals[1])
    finish = time.time()

    gdx.gdxClose(gdxHandle)
    gdx.gdxFree(gdxHandle)
    gdx.gdxLibraryUnload()

    return finish - start

def _parse_special_values(value):
    if value == 1.0e300 or value == 2.0e300:
        return float("nan")
    if value == 3.0e300:
        return float("inf")
    if value == 4.0e300:
        return -float("inf")
    if value == 5.0e300:
        return sys.float_info.epsilon
    return value

def pyomo_like(file_path):
    system_directory = "/Library/Frameworks/GAMS.framework/Versions/49/Resources/"
    gdxHandle = gdx.new_gdxHandle_tp()
    gdx.gdxCreateD(gdxHandle, system_directory, gdx.GMS_SSSIZE)
    gdx.gdxOpenRead(gdxHandle, file_path)
    ret, symCount, _ = gdx.gdxSystemInfo(gdxHandle)

    i = 0
    model_soln = {}
    start = time.time()
    while True:
        i += 1
        ret = gdx.gdxDataReadRawStart(gdxHandle, i)
        if not ret[0]:
            break

        ret = gdx.gdxDataReadRaw(gdxHandle)
        if not ret[0] or len(ret[2]) < 2:
            raise RuntimeError("GAMS GDX failure (gdxDataReadRaw).")
        level = _parse_special_values(ret[2][0])
        dual = _parse_special_values(ret[2][1])

        ret = gdx.gdxSymbolInfo(gdxHandle, i)
        if not ret[0]:
            break
        if len(ret) < 2:
            raise RuntimeError("GAMS GDX failure (gdxSymbolInfo).")
        model_soln[ret[1]] = (level, dual)

    finish = time.time()

    gdx.gdxDataReadDone(gdxHandle)
    gdx.gdxClose(gdxHandle)
    gdx.gdxFree(gdxHandle)
    gdx.gdxLibraryUnload()

    return finish - start

def gdx_extractor(base_dir = os.getcwd()):
    gdx_file = []
    gms_size = []

    # Iterate through each subdirectory
    for entry in os.listdir(base_dir):
        subdir_path = os.path.join(base_dir, entry)
        model_size = re.findall('\d+',subdir_path.split('/')[-1])
        
        if os.path.isdir(subdir_path):

            # search for file name
            for file_name in os.listdir(subdir_path):    
                if file_name.endswith("p.gdx"):

                    file_path = os.path.join(subdir_path, file_name)
                    
                    gdx_file.append(file_path)
                    gms_size.append(model_size)

    # sort the gdx file in increasing number of variables
    temp = {}
    for i in range(len(gdx_file)):
        temp[float(gms_size[i][0])] = gdx_file[i]

    gdx_file = dict(sorted(temp.items()))
    return gdx_file

gdxFiles = gdx_extractor()

repeats = 4
tictoc_gams = {}
tictoc_pyomo = {}

for gmsSize, gdxPath in gdxFiles.items():
    print(gmsSize, gdxPath)
    tictoc_gams[gmsSize] = []
    tictoc_pyomo[gmsSize] = []

    # the process is repeated N times to reduce hardware noise
    for n in range(repeats):
        tictoc_gams[gmsSize].append(gams_like(gdxPath))
        tictoc_pyomo[gmsSize].append(pyomo_like(gdxPath))


plt.figure(figsize=(8, 6))

# Convert dictionary to sorted arrays
x_Gvals = np.array(sorted(tictoc_gams.keys()))
y_Gvals = np.array([np.mean(tictoc_gams[k]) for k in x_Gvals])
y_Gerrs = np.array([np.std(tictoc_gams[k]) for k in x_Gvals])

x_Pvals = np.array(sorted(tictoc_pyomo.keys()))
y_Pvals = np.array([np.mean(tictoc_pyomo[k]) for k in x_Pvals])
y_Perrs = np.array([np.std(tictoc_pyomo[k]) for k in x_Pvals])

x_log = [i+1 for i in range(len(x_Gvals))]
perf_improved = float(np.abs(y_Gvals[-1] - y_Pvals[-1]) / y_Pvals[-1]) * 100

# Plot with error bars
plt.errorbar(x_log, y_Gvals, yerr=y_Gerrs, fmt='o-', capsize=5, label='gams')
plt.errorbar(x_log, y_Pvals, yerr=y_Perrs, fmt='*-', capsize=5, label='pyomo')

plt.xticks(x_log, ['1.e+01', '1.e+02', '1.e+03', '1.e+04', '1.e+05', '1.e+06', '5.e+06', '1.e+07'])

plt.xlabel('Number of Variables')
plt.ylabel('Time (second)')
plt.title(f'Mean Reading Time of GAMS vs. Pyomo - {perf_improved:.2f}% Improved')
plt.legend()
plt.tight_layout()

plt.savefig('performance_comparison.png', dpi = 1000)

plt.show()