from pandas import Series, DataFrame
from minpower.config import user_config
from minpower import powersystems, schedule, solve
from minpower.generators import Generator
from minpower.optimization import value, OptimizationError

import nose
from nose.tools import istest, with_setup, raises, set_trace
from nose.tools import eq_ as assert_equal


import logging

logging.basicConfig(level=logging.CRITICAL, format="%(levelname)s: %(message)s")


default_config = user_config.copy()


def get_duals():
    user_config.duals = True


def reset_config():
    # note - can't set a new user_config, because it would be a local variable
    #   instead, we update it with the defaults
    #   this only works if we aren't adding anything extra to the user_config
    user_config.update(default_config)
    try:
        assert user_config == default_config
    except:
        print(
            (
                DataFrame(
                    dict(config=Series(user_config), default=Series(default_config))
                )
            )
        )
        raise


singletime = schedule.just_one_time()


gen_costs = dict(cheap=10, mid=20, expensive=30)


def assertAlmostEqual(x, y):
    assert round(x - y, 5) == 0


# def make_single_bus(generators,loads):
# singlebus=powersystems.Bus()
# singlebus.generators=generators
# singlebus.loads=loads
# return [singlebus]


def make_cheap_gen(**kwargs):
    return Generator(
        name="cheap gen", costcurveequation="{}P".format(gen_costs["cheap"]), **kwargs
    )


def make_mid_gen(**kwargs):
    return Generator(
        name="middle-range gen",
        costcurveequation="{}P".format(gen_costs["mid"]),
        **kwargs
    )


def make_expensive_gen(**kwargs):
    if "costcurveequation" not in kwargs:
        kwargs["costcurveequation"] = "{}P".format(gen_costs["expensive"])
    return Generator(name="expensive gen", **kwargs)


def make_loads_times(Pd=200, Pdt=None, **kwargs):
    if Pdt is None:
        times = singletime
        sched = Series(Pd, index=times)
        loads = [powersystems.Load(schedule=sched, **kwargs)]
    else:
        times = schedule.make_times_basic(N=len(Pdt))
        sched = Series(Pdt, index=times)
        loads = [powersystems.Load(schedule=sched, **kwargs)]

    return dict(loads=loads, times=times)


def solve_problem(
    generators, loads=None, times=None, gen_init=None, lines=None, do_reset_config=True
):

    if lines is None:
        lines = []

    if len(times) > 0:
        for g, gen in enumerate(generators):
            gen.index = g
            if gen_init is None:
                gen.set_initial_condition()
            else:
                gen.set_initial_condition(**gen_init[g])

    power_system = powersystems.PowerSystem(generators, loads, lines)
    solve.create_problem(power_system, times)

    try:
        instance = power_system.solve()
    except OptimizationError:
        # re-do stage, with load shedding allowed
        logging.critical("stage infeasible, re-running with load shedding.")
        power_system.allow_shedding(times)
        try:
            instance = power_system.solve()
        except OptimizationError:
            scheduled, committed = power_system.debug_infeasibe(times)
            power_system.write_model("infeasible.lp")
            raise OptimizationError("failed to solve, even with load shedding.")

    if do_reset_config:
        reset_config()

    if power_system.solved:
        power_system.update_variables()
    else:
        power_system.write_model("problem.lp")
        raise OptimizationError("infeasible problem, wrote to problem.lp")
    return power_system, times
