import time

import numpy as np
import pandas as pd
import pytest

import tsam.timeseriesaggregation as tsam
from conftest import TESTDATA_CSV

pytestmark = pytest.mark.filterwarnings("ignore::tsam.exceptions.LegacyAPIWarning")


@pytest.mark.filterwarnings("ignore:Segmentation is turned off:UserWarning")
def test_properties():
    hoursPerPeriod = 24

    noSegments = 8

    noTypicalPeriods = 8

    raw = pd.read_csv(TESTDATA_CSV, index_col=0)

    starttime = time.time()

    aggregation1 = tsam.TimeSeriesAggregation(
        raw,
        noTypicalPeriods=noTypicalPeriods,
        hoursPerPeriod=hoursPerPeriod,
        clusterMethod="hierarchical",
        segmentation=True,
        noSegments=noSegments,
    )

    print("Clustering took " + str(time.time() - starttime))

    np.testing.assert_array_almost_equal(
        aggregation1.stepIdx, np.arange(noSegments), decimal=4
    )

    starttime = time.time()

    aggregation2 = tsam.TimeSeriesAggregation(
        raw,
        noTypicalPeriods=noTypicalPeriods,
        hoursPerPeriod=hoursPerPeriod,
        clusterMethod="hierarchical",
    )

    print("Clustering took " + str(time.time() - starttime))

    np.testing.assert_array_almost_equal(
        aggregation2.stepIdx, np.arange(hoursPerPeriod), decimal=4
    )

    starttime = time.time()

    aggregation3 = tsam.TimeSeriesAggregation(
        raw,
        noTypicalPeriods=noTypicalPeriods,
        hoursPerPeriod=hoursPerPeriod,
        clusterMethod="hierarchical",
    )

    print("Clustering took " + str(time.time() - starttime))

    np.testing.assert_array_almost_equal(
        aggregation3.clusterPeriodIdx, np.arange(noTypicalPeriods), decimal=4
    )

    starttime = time.time()

    aggregation4 = tsam.TimeSeriesAggregation(
        raw,
        noTypicalPeriods=noTypicalPeriods,
        hoursPerPeriod=hoursPerPeriod,
        clusterMethod="hierarchical",
        segmentation=True,
        noSegments=noSegments,
    )

    print("Clustering took " + str(time.time() - starttime))

    appearances = np.unique(aggregation4.clusterOrder, return_counts=True)[1].tolist()

    occurrenceDict = {i: j for i, j in enumerate(appearances)}

    # make sure that the clusterPeriodNoOccur equals the number of appearances in the clusterOrder
    np.testing.assert_array_almost_equal(
        list(aggregation4.clusterPeriodNoOccur.values()),
        list(occurrenceDict.values()),
        decimal=4,
    )

    starttime = time.time()

    aggregation5 = tsam.TimeSeriesAggregation(
        raw,
        noTypicalPeriods=noTypicalPeriods,
        hoursPerPeriod=hoursPerPeriod,
        clusterMethod="hierarchical",
        segmentation=True,
        noSegments=noSegments,
    )

    print("Clustering took " + str(time.time() - starttime))

    # make sure that the values of the clusterPeriodDict equal those from the typicalPeriods-dataframe
    np.testing.assert_array_almost_equal(
        pd.DataFrame.from_dict(data=aggregation5.clusterPeriodDict).values,
        aggregation5.createTypicalPeriods().values,
        decimal=4,
    )

    starttime = time.time()

    aggregation6 = tsam.TimeSeriesAggregation(
        raw,
        noTypicalPeriods=noTypicalPeriods,
        hoursPerPeriod=hoursPerPeriod,
        clusterMethod="hierarchical",
    )

    print("Clustering took " + str(time.time() - starttime))

    # make sure that the sum of all segment durations in each period equals the hours per period
    for i in range(noTypicalPeriods):
        print(i)
        print(pd.DataFrame.from_dict(aggregation6.segmentDurationDict).loc[(i,), :])
        # print(
        #     pd.DataFrame.from_dict(aggregation6.segmentDurationDict)
        #     .iloc[
        #         pd.DataFrame.from_dict(
        #             aggregation6.segmentDurationDict
        #         ).index.get_level_values(0)
        #     ]
        # )
        print("\n")
        np.testing.assert_array_almost_equal(
            pd.DataFrame.from_dict(aggregation6.segmentDurationDict)
            .loc[(i,), :]
            .sum()
            .iloc[0],
            hoursPerPeriod,
            decimal=4,
        )
        print("")

    starttime = time.time()

    aggregation7 = tsam.TimeSeriesAggregation(
        raw,
        noTypicalPeriods=noTypicalPeriods,
        hoursPerPeriod=hoursPerPeriod,
        clusterMethod="hierarchical",
        segmentation=True,
        noSegments=noSegments,
    )

    print("Clustering took " + str(time.time() - starttime))

    # make sure that the sum of all segment durations in each period equals the hours per period
    for i in range(noTypicalPeriods):
        np.testing.assert_array_almost_equal(
            pd.DataFrame.from_dict(aggregation7.segmentDurationDict)
            .loc[i]
            .sum()
            .iloc[0],
            hoursPerPeriod,
            decimal=4,
        )

    starttime = time.time()

    aggregation8 = tsam.TimeSeriesAggregation(
        raw,
        noTypicalPeriods=noTypicalPeriods,
        hoursPerPeriod=hoursPerPeriod,
        clusterMethod="hierarchical",
        segmentation=True,
        noSegments=noSegments,
    )

    print("Clustering took " + str(time.time() - starttime))

    indexTable = aggregation8.indexMatching()

    # make sure that the PeriodNum column equals the clusterOrder
    np.testing.assert_array_almost_equal(
        indexTable.loc[::24, "PeriodNum"].values, aggregation8.clusterOrder, decimal=4
    )

    # make sure that the TimeStep indices equal the number of hoursPerPeriod arranged as array
    np.testing.assert_array_almost_equal(
        pd.unique(indexTable.loc[:, "TimeStep"]),
        np.arange(hoursPerPeriod, dtype="int64"),
        decimal=4,
    )

    # make sure that the SegmentIndex indices equal the number of noSegments arranged as array
    np.testing.assert_array_almost_equal(
        pd.unique(indexTable.loc[:, "SegmentIndex"]),
        np.arange(noSegments, dtype="int64"),
        decimal=4,
    )


if __name__ == "__main__":
    test_properties()
