import numpy as np

def test_segmentation(minimal_test_esM):
    '''
    Get the minimal test system, and check that for different segment and period configurations the same solution is
    found.
    '''

    # First, the non-aggregated case is compared to the aggregation mode of the model, but without aggregated data.
    # For this, the mini system is first optimized without any aggregation at all.
    esM1 = minimal_test_esM
    esM1.optimize(solver='glpk')
    # Then, the four time steps of the model are represented by two 4380-hourly typical periods with two segments per
    # typical period, so effectively the data is not aggregated.
    esM2 = minimal_test_esM
    esM2.cluster(numberOfTypicalPeriods=2, numberOfTimeStepsPerPeriod=2, storeTSAinstance=False,
                 segmentation=True, numberOfSegmentsPerPeriod=2, clusterMethod='hierarchical',
                 sortValues=False, rescaleClusterPeriods=False)
    esM2.optimize(timeSeriesAggregation=True, solver='glpk')
    # It is now checked that both models, i.e. the one without aggregation at all and the one without aggregation, but
    # in aggregation mode, lead to the same result.
    assert esM1.pyM.Obj() == esM2.pyM.Obj()

    # For the following configurations, no storage is built and the demand is always the same. Further, the prices are
    # calculated by their centroid. Accordingly, both configurations should lead to the same objective value.
    # First, the mini system is clustered to one period with four time steps that is further segmented to three segments
    # so that the first segment is twice as long as the first and the second segment.
    esM3 = minimal_test_esM
    esM3.cluster(numberOfTypicalPeriods=1, numberOfTimeStepsPerPeriod=4, storeTSAinstance=False,
                 segmentation=True, numberOfSegmentsPerPeriod=3, clusterMethod='hierarchical',
                 sortValues=False, rescaleClusterPeriods=False)
    esM3.optimize(timeSeriesAggregation=True, solver='glpk')
    # Then, the model is optimized again with two 4380-hourly periods that are segmented to one segment per period, i.e.
    # the model contains only two time steps in total with averaged values in each period.
    esM4 = minimal_test_esM
    esM4.cluster(numberOfTypicalPeriods=2, numberOfTimeStepsPerPeriod=2, storeTSAinstance=False,
                 segmentation=True, numberOfSegmentsPerPeriod=1, clusterMethod='hierarchical',
                 sortValues=False, rescaleClusterPeriods=False)
    esM4.optimize(timeSeriesAggregation=True, solver='glpk')
    # Here, it is checked that the results of the third and the fourth model run are identical because no storage is
    # chosen and because of the averaged data the costs should stay the same.
    # Note: The segmentation also averages the models' constraints, but in this specific example, the most restrictive
    # and thus size-determining constraints of the model are coincidentally not affected by the aggregation and the
    # optimal solutions of the third and fourth model are identical.
    assert esM3.pyM.Obj() == esM4.pyM.Obj()
