import pytest

import eppy_funcs as ef
import sampling
import objectives
from problem import EPProblem
from evaluator import EvaluatorEP, EvaluatorSR
from parameters import RangeParameter, CategoryParameter, expand_plist
import optimizer
import numpy as np

from sklearn import linear_model
from sklearn.model_selection import train_test_split


@pytest.fixture
def building():
    return ef.get_building()


@pytest.fixture
def problem():
    parameters = expand_plist(
        {
            "Mass NonRes Wall Insulation": {"Thickness": (0.01, 0.99)},
            "NonRes Fixed Assembly Window": {
                "U-Factor": (0.1, 5),
                "Solar Heat Gain Coefficient": (0.01, 0.99),
            },
        }
    )
    return EPProblem(parameters, ["Electricity:Facility"])


@pytest.fixture
def samples(problem):
    return sampling.dist_sampler(sampling.seeded_sampler, problem, 10)


def test_fit(building, problem, samples):
    """check to make sure linear regression works and is close to the actual value"""

    evaluator = EvaluatorEP(problem, building)
    outputs = evaluator.df_apply(samples)

    train_in, test_in, train_out, test_out = train_test_split(
        samples, outputs, test_size=0.2
    )

    reg = linear_model.LinearRegression()
    reg.fit(train_in, train_out)
    results = test_in.copy()
    results["energy use"] = test_out
    results["predicted"] = reg.predict(test_in)

    actual = results.iloc[0]["energy use"]
    predicted = results.iloc[0]["predicted"]
    assert np.isclose(
        actual - predicted, 8666133.340038776
    ), f"Unexpected difference of value: {actual - predicted}"

    # change this to 0 to see stdout and stderr
    assert 1


def test_optimisation_with_surrogate(building, problem, samples):
    """make sure we can run optimisation algorithms on the fit surrogates"""

    evaluator = EvaluatorEP(problem, building)
    outputs = evaluator.df_apply(samples)

    train_in, test_in, train_out, test_out = train_test_split(
        samples, outputs, test_size=0.2
    )

    reg = linear_model.LinearRegression()
    reg.fit(train_in, train_out)

    def evaluation_func(ind):
        return ((reg.predict([ind])[0][0],), ())

    surrogate = EvaluatorSR(evaluation_func, problem)

    s = optimizer.NSGAII(surrogate, 1000)

    optimal = s.iloc[0]["pareto-optimal"]
    non_optimal = s.iloc[1]["pareto-optimal"]

    assert optimal, f"Optimal ouput was displayed as not Pareto Optimal"
    assert not non_optimal, f"Non-optimal output was displayed as Pareto Optimal"

    # change this to 0 to see stdout and stderr
    assert 1
