# ---------------------------------------------------------------------------------------------------------------------
# region import & abstract model

import pyomo.environ as pyo
from Master import data

ex_model = pyo.AbstractModel()

# endregion
# ---------------------------------------------------------------------------------------------------------------------
# region parameters

ex_model.M = pyo.Param()  # bigM

# 3D (c_rate)

ex_model.bp_num_x1 = pyo.Param(within=pyo.NonNegativeIntegers)  # number of bp's in x
ex_model.bp_num_y1 = pyo.Param(within=pyo.NonNegativeIntegers)  # number of bp's in y

ex_model.bp_index_x1 = pyo.RangeSet(1, ex_model.bp_num_x1 - 1)  # index number of bp's in x-1 for counting
ex_model.bp_index_y1 = pyo.RangeSet(1, ex_model.bp_num_y1 - 1)  # index number of bp's in y-1 for counting

ex_model.bp_index_v_x1 = pyo.RangeSet(1, ex_model.bp_num_x1)  # index number of bp's in x for values
ex_model.bp_index_v_y1 = pyo.RangeSet(1, ex_model.bp_num_y1)  # index number of bp's in y for values

ex_model.bp_index_z1 = pyo.RangeSet(1, (ex_model.bp_num_x1 - 1) * ex_model.bp_num_y1)  # index number for z values

# output of func.
ex_model.x1 = pyo.Param(ex_model.bp_index_v_x1)  # running over i
ex_model.y1 = pyo.Param(ex_model.bp_index_v_y1)  # running over j
ex_model.z1 = pyo.Param(ex_model.bp_index_z1)  # running over i and j

# 3D (SOC)

ex_model.bp_num_x2 = pyo.Param(within=pyo.NonNegativeIntegers)  # number of bp's in x
ex_model.bp_num_y2 = pyo.Param(within=pyo.NonNegativeIntegers)  # number of bp's in y

ex_model.bp_index_x2 = pyo.RangeSet(1, ex_model.bp_num_x2 - 1)  # index number of bp's in x-1 for counting
ex_model.bp_index_y2 = pyo.RangeSet(1, ex_model.bp_num_y2 - 1)  # index number of bp's in y-1 for counting

ex_model.bp_index_v_x2 = pyo.RangeSet(1, ex_model.bp_num_x2)  # index number of bp's in x for values
ex_model.bp_index_v_y2 = pyo.RangeSet(1, ex_model.bp_num_y2)  # index number of bp's in y for values

ex_model.bp_index_z2 = pyo.RangeSet(1, (ex_model.bp_num_x2 - 1) * ex_model.bp_num_y2)  # index number for z values

# output of func.
ex_model.x2 = pyo.Param(ex_model.bp_index_v_x2)  # running over i
ex_model.y2 = pyo.Param(ex_model.bp_index_v_y2)  # running over j
ex_model.z2 = pyo.Param(ex_model.bp_index_z2)  # running over i and j

# endregion
# ---------------------------------------------------------------------------------------------------------------------
# region variables

# 3D (c_rate)

ex_model.S_ddd1 = pyo.Var(ex_model.steps, ex_model.bp_index_x1, ex_model.bp_index_y1,
                          domain=pyo.Binary)  # 3 times indexed

# 3D (SOC)

ex_model.S_ddd2 = pyo.Var(ex_model.steps, ex_model.bp_index_x2, ex_model.bp_index_y2,
                          domain=pyo.Binary)  # 3 times indexed

# system variables

# NL-variables
ex_model.c_rate = pyo.Var(ex_model.steps, domain=pyo.NonNegativeReals, bounds=(0, 0.25))  # c_rate (crg)
ex_model.P_max = pyo.Var(ex_model.steps, domain=pyo.NonNegativeReals, bounds=(0, 250))  # P_max (crg)
ex_model.E_max = pyo.Var(domain=pyo.NonNegativeReals, bounds=(1, 1000))  # target sizing
ex_model.E = pyo.Var(ex_model.steps_ESS, domain=pyo.NonNegativeReals, bounds=(0, 1000))  # current SoC [kWh]
ex_model.SOC = pyo.Var(ex_model.steps, domain=pyo.NonNegativeReals, bounds=(0, 1))  # current SoC [%/100]


# endregion
# ---------------------------------------------------------------------------------------------------------------------
# region constraints

# 3D pwla constraints
# y is the pwl variable and x the discrete (global) variable

# 3D (c_rate)
# c = f(P_max(t), E_max)


def area_assign1_1(ex_model, t, i, j):
    return ex_model.E_max + ex_model.S_ddd1[t, i, j] * ex_model.M <= ex_model.x1[i + 1] + ex_model.M - 0.001


ex_model.ddd_1_1 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x1, ex_model.bp_index_y1, rule=area_assign1_1)


def area_assign1_2(ex_model, t, i, j):
    return ex_model.x1[i] + ex_model.S_ddd1[t, i, j] * ex_model.M <= ex_model.E_max + ex_model.M


ex_model.ddd_1_2 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x1, ex_model.bp_index_y1, rule=area_assign1_2)


def area_assign1_3(ex_model, t, i, j):
    return ex_model.P_max[t] + ex_model.S_ddd1[t, i, j] * ex_model.M <= ex_model.y1[j + 1] + ex_model.M


ex_model.ddd_1_3 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x1, ex_model.bp_index_y1, rule=area_assign1_3)


def area_assign1_4(ex_model, t, i, j):
    return ex_model.y1[j] + ex_model.S_ddd1[t, i, j] * ex_model.M <= ex_model.P_max[t] + ex_model.M


ex_model.ddd_1_4 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x1, ex_model.bp_index_y1, rule=area_assign1_4)


def area_assign1_5(ex_model, t, i, j):
    return ex_model.c_rate[t] + ex_model.S_ddd1[t, i, j] * ex_model.M <= (
            ex_model.z1[(j + 1) + ex_model.bp_num_y1 * (i - 1)]
            - ex_model.z1[j + ex_model.bp_num_y1 * (i - 1)]) / (
                   ex_model.y1[j + 1] - ex_model.y1[j]) * \
           (ex_model.P_max[t] - ex_model.y1[j]) + ex_model.z1[j + ex_model.bp_num_y1 * (i - 1)] + ex_model.M


ex_model.ddd_1_5 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x1, ex_model.bp_index_y1, rule=area_assign1_5)


def area_assign1_6(ex_model, t, i, j):
    return (ex_model.z1[(j + 1) + ex_model.bp_num_y1 * (i - 1)] - ex_model.z1[j + ex_model.bp_num_y1 * (i - 1)]) / \
           (ex_model.y1[j + 1] - ex_model.y1[j]) * (ex_model.P_max[t] - ex_model.y1[j]) + \
           ex_model.S_ddd1[t, i, j] * ex_model.M + ex_model.z1[j + ex_model.bp_num_y1 * (i - 1)] <= ex_model.c_rate[
               t] + ex_model.M


ex_model.ddd_1_6 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x1, ex_model.bp_index_y1, rule=area_assign1_6)


def sum_S_ddd1(ex_model, t, i, j):
    return sum(ex_model.S_ddd1[t, i, j] for i, j in ex_model.bp_index_x1 * ex_model.bp_index_y1) == 1


ex_model.sum_ddd1 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x1, ex_model.bp_index_y1, rule=sum_S_ddd1)


# 3D (SOC)
# SOC = f(E(t), E_max)

def area_assign2_1(ex_model, t, i, j):
    return ex_model.E_max + ex_model.S_ddd2[t, i, j] * ex_model.M <= ex_model.x2[i + 1] + ex_model.M - 0.001


ex_model.ddd_2_1 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x2, ex_model.bp_index_y2, rule=area_assign2_1)


def area_assign2_2(ex_model, t, i, j):
    return ex_model.x2[i] + ex_model.S_ddd2[t, i, j] * ex_model.M <= ex_model.E_max + ex_model.M


ex_model.ddd_2_2 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x2, ex_model.bp_index_y2, rule=area_assign2_2)


def area_assign2_3(ex_model, t, i, j):
    return ex_model.E[t] + ex_model.S_ddd2[t, i, j] * ex_model.M <= ex_model.y2[j + 1] + ex_model.M


ex_model.ddd_2_3 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x2, ex_model.bp_index_y2, rule=area_assign2_3)


def area_assign2_4(ex_model, t, i, j):
    return ex_model.y2[j] + ex_model.S_ddd2[t, i, j] * ex_model.M <= ex_model.E[t] + ex_model.M


ex_model.ddd_2_4 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x2, ex_model.bp_index_y2, rule=area_assign2_4)


def area_assign2_5(ex_model, t, i, j):
    return ex_model.SOC[t] + ex_model.S_ddd2[t, i, j] * ex_model.M <= (
            ex_model.z2[(j + 1) + ex_model.bp_num_y2 * (i - 1)]
            - ex_model.z2[j + ex_model.bp_num_y2 * (i - 1)]) / (
                   ex_model.y2[j + 1] - ex_model.y2[j]) * \
           (ex_model.E[t] - ex_model.y2[j]) + ex_model.z2[j + ex_model.bp_num_y2 * (i - 1)] + ex_model.M


ex_model.ddd_2_5 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x2, ex_model.bp_index_y2, rule=area_assign2_5)


def area_assign2_6(ex_model, t, i, j):
    return (ex_model.z2[(j + 1) + ex_model.bp_num_y2 * (i - 1)] - ex_model.z2[j + ex_model.bp_num_y2 * (i - 1)]) / \
           (ex_model.y2[j + 1] - ex_model.y2[j]) * (ex_model.E[t] - ex_model.y2[j]) + \
           ex_model.S_ddd2[t, i, j] * ex_model.M + ex_model.z2[j + ex_model.bp_num_y2 * (i - 1)] <= ex_model.SOC[
               t] + ex_model.M


ex_model.ddd_2_6 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x2, ex_model.bp_index_y2, rule=area_assign2_6)


def sum_S_ddd2(ex_model, t, i, j):
    return sum(ex_model.S_ddd2[t, i, j] for i, j in ex_model.bp_index_x2 * ex_model.bp_index_y2) == 1


ex_model.sum_ddd2 = pyo.Constraint(ex_model.steps, ex_model.bp_index_x2, ex_model.bp_index_y2, rule=sum_S_ddd2)

#  endregion
# ---------------------------------------------------------------------------------------------------------------------