import FINE as fn
import pandas as pd


def test_TSAmultiStage(minimal_test_esM):
    """
    Get the minimal test system, and check if the Error-Bounding-Approach works for it
    """

    # modify the minimal LP and change it to a MILP
    esM = minimal_test_esM

    # get the components with capacity variables
    electrolyzers = esM.getComponent("Electrolyzers")
    pressureTank = esM.getComponent("Pressure tank")
    pipelines = esM.getComponent("Pipelines")

    # set binary variables and define bigM
    electrolyzers.hasIsBuiltBinaryVariable = True
    pressureTank.hasIsBuiltBinaryVariable = True
    pipelines.hasIsBuiltBinaryVariable = True

    electrolyzers.investIfBuilt = pd.Series(2e5, index=esM.locations)
    pressureTank.investIfBuilt = pd.Series(1e5, index=esM.locations)
    pipelines.investIfBuilt.loc[:] = 100

    electrolyzers.bigM = 30e4
    pressureTank.bigM = 30e6
    pipelines.bigM = 30e3

    electrolyzers.investIfBuilt = pd.Series(2e5, index=esM.locations)
    pressureTank.investIfBuilt = pd.Series(1e5, index=esM.locations)

    electrolyzers.bigM = 30e4
    pressureTank.bigM = 30e6

    # optimize with 2 Stage Approach
    fn.optimizeTSAmultiStage(
        esM,
        relaxIsBuiltBinary=True,
        solver="glpk",
        numberOfTypicalPeriods=2,
        numberOfTimeStepsPerPeriod=1,
    )

    # get gap
    gap = esM.gap

    assert gap > 0.1078 and gap < 0.1079
