import itertools
import pyomo.environ as py
import pandas as pd
import pylab as pl
import numpy as np
from matplotlib import pyplot as plt



# to ignore the warning for dividing by 0
np.seterr(divide='ignore', invalid='ignore')

############################################################First Step#############################################
# impedance matrix
# importing the data from the excel for the bus
# <editor-fold desc="Reading the data from Excel file">
data = pd.ExcelFile('opf_data_time.xlsx')
databus = data.parse('bus')
datademand = data.parse('demand')
datacoeff = data.parse('coefficent')
datagen = data.parse('gen')
datatime = data.parse('time')
# </editor-fold>
dataimpedence = data.parse('impedance')
Imp_matrix = np.zeros((5, 5)).astype(complex)
for y in range(len(dataimpedence)):
    Imp_matrix[dataimpedence.loc[y, "From"] - 1, dataimpedence.loc[y, "To"] - 1] = dataimpedence.loc[y, "R"] + 1j * \
                                                                                   dataimpedence.loc[y, "X"]
    Imp_matrix[dataimpedence.loc[y, "To"] - 1, dataimpedence.loc[y, "From"] - 1] = Imp_matrix[
        dataimpedence.loc[y, "From"] - 1, dataimpedence.loc[y, "To"] - 1]

print(Imp_matrix)

# shunt admitance/ line charging
shunt_matrix = np.zeros((5, 5)).astype(complex)
for y in range(len(dataimpedence)):
    shunt_matrix[dataimpedence.loc[y, "From"] - 1, dataimpedence.loc[y, "To"] - 1] = 1j * dataimpedence.loc[y, "B"]
    shunt_matrix[dataimpedence.loc[y, "To"] - 1, dataimpedence.loc[y, "From"] - 1] = shunt_matrix[
        dataimpedence.loc[y, "From"] - 1, dataimpedence.loc[y, "To"] - 1]
print(shunt_matrix)

# converting the impedance to admittance
yb = -1 / Imp_matrix
yb[np.isnan(yb)] = 0

# positive of yb
pyb = -(yb)

# sum the row of the shunt admittance to preparid for the diaginal element
shunt_matrix_sum = np.sum(shunt_matrix, axis=1)

# ist i make the diagonal matrix to zero and then add row wise
pyb[np.diag_indices_from(pyb)] = 0
y_offdiagonal = pyb
y_offdiagonal_sum = np.sum(y_offdiagonal, axis=1)
# print(y_offdiagonal_sum)

# now adding the y_offdiagonal_sum with shunt_matrix_sum for diagonal matrix
yii = np.add(shunt_matrix_sum, y_offdiagonal_sum)  # get the diaginal matrix
# print(yii)


np.fill_diagonal(yb, yii)
ybus = yb
print('Y bus: ', ybus)

# calculating the magnitude and angle of ybus matrix
y_mag = abs(ybus)
print('Y bus magnitude : ', y_mag)
y_rad = np.angle(ybus)
print('Y bus angle (Rad) : ', y_rad)
y_ang = np.degrees(y_rad)

# convert the ybus matrix into real (Conductance) and imaginary ( substance) matrix
ybus_real = ybus.real
print('Conductance: ', ybus_real)

ybus_imag = ybus.imag

print('Susseptance: ', ybus_imag)

#########################################################Second Step###################################################


# optimal power flow
model = py.ConcreteModel()
# Set
# model.b = py.Set(initialize=[0, 1, 2, 3, 4], doc='bus')
model.b = py.Set(initialize=databus.bus, doc='bus')
model.index_i = py.Set(initialize=databus.bus, doc='index i')
model.index_j = py.Set(initialize=databus.bus, doc='index j')
model.t = py.Set(initialize=datatime.time, doc='time')

# Parameter

# demand form the homes orginal
model.pdmand = py.Param(model.b, model.t, initialize=dict(
    zip(list(itertools.product(range(1, len(databus.bus.values) + 1), range(1, len(model.t.data()) + 1))),
        np.append(datademand.active_power.values, len(model.t.data())))),
                        doc='Active power demand')
model.qdmand = py.Param(model.b, model.t, initialize=dict(
    zip(list(itertools.product(range(1, len(databus.bus.values) + 1), range(1, len(model.t.data()) + 1))),
        np.append(datademand.reactive_power.values, len(model.t.data())))),
                        doc='Reactive power demand')
# Variable
model.v = py.Var(model.index_i, model.t, initialize=1, bounds=(0.9, 1.07))
model.v[1, :].fix(1.06)
model.v[2, :].fix(1.045)
model.v[3, :].fix(1.03)
model.theta = py.Var(model.index_i, model.t, initialize=0, bounds=(-np.pi, np.pi))
model.theta[1, :].fix(0)

# bounds for generator active power
model.plb = py.Param(model.b, initialize=dict(zip(databus.bus, datagen.pl)),
                     doc='active power lower bound')
model.pub = py.Param(model.b, initialize=dict(zip(databus.bus, datagen.pu)),
                     doc='active power upper bound')


def pgenbound(model, i, t):
    return model.plb[i], model.pub[i]


# model.pgen = py.Var(model.b, model.t, bounds=pgenbound)
model.pgen = py.Var(model.b, model.t,bounds=(0, None))
model.pgen[4, :].fix(0)
model.pgen[5, :].fix(0)

# bounds for generator reactive power
model.qlb = py.Param(model.b, initialize=dict(zip(databus.bus, datagen.ql)), doc='reactive power lower bound')
model.qub = py.Param(model.b, initialize=dict(zip(databus.bus, datagen.qu)), doc='reactive power upper bound')


def qgenbound(model, i, t):
    return model.qlb[i], model.qub[i]


# model.qgen = py.Var(model.b, model.t, bounds=qgenbound)
model.qgen = py.Var(model.b, model.t)
model.qgen[4, :].fix(0)
model.qgen[5, :].fix(0)
print('model data is ok')


# Constraints
def active_power(model, i, t):
    return model.pgen[i, t] - model.pdmand[i, t]- sum(
        [y_mag[i - 1, j - 1] * (model.v[i, t]) * (model.v[j, t]) * py.cos(
            y_rad[i - 1, j - 1] + model.theta[j, t] - model.theta[i, t])
         for j in model.index_j]) ==0

model.Const_Active_power = py.Constraint(model.index_i, model.t, rule=active_power, doc='active power injection')


# the plus sign is because of the P-Qj = .....
def reactive_power(model, i, t):
    return model.qgen[i, t] - model.qdmand[i, t]+ sum(
        [y_mag[i - 1, j - 1] * (model.v[i, t]) * (model.v[j, t]) * py.sin(
            y_rad[i - 1, j - 1] + model.theta[j, t] - model.theta[i, t])
         for j in model.index_j])==0

model.Const_Reactive_power = py.Constraint(model.index_i, model.t, rule=reactive_power,
                                           doc='Reactive power injection')

print('model constraint are ok')

for k in model.index_i:
    for l in model.index_j:
        ybus_real[k - 1, l - 1] = (-(ybus_real[k - 1, l - 1]))


def objective_rule(model):
    return sum(
        sum(0.5 * (ybus_real[m - 1, n - 1]) * ((model.v[m, l]) ** 2 + (model.v[n, l]) ** 2 - 2 * (
                (model.v[m, l]) * (model.v[n, l])) * py.cos(model.theta[n, l] - model.theta[m, l])) for m in
            model.index_i for n in model.index_j if m != n) for l in model.t)


model.objective = py.Objective(rule=objective_rule, sense=py.minimize, doc='Definition of objective function')

# Solver

opt = py.SolverFactory('ipopt')
result = opt.solve(model)
model.write('model_opf.nl', io_options={'symbolic_solver_labels': True})

print('objective function = ', py.value(model.objective) * 100)

if (result.solver.status == py.SolverStatus.ok) and (
        result.solver.termination_condition == py.TerminationCondition.optimal):
    print('optimal solution')  # Do something when the solution in optimal and feasible
elif result.solver.termination_condition == py.TerminationCondition.infeasible:
    print('solution is not feasible')  # Do something when model in infeasible
else:
    # Something else is wrong
    print('Solver Status: ', result.solver.status)

# ploting:
pl_pdem = [[] for j in model.b]
pl_pgen = [[] for j in model.b]
time = [i for i in model.t]
pl_qdem = [[] for j in model.b]
pl_v = [[] for j in model.b]
pl_qgen = [[] for j in model.b]
pl_theta = [[] for j in model.b]

for j in model.b:
    for k, v in model.pgen.items():
        if k[0] == j:
            pl_pgen[j - 1].append(py.value(v))

for j in model.b:
    for k, v in model.qgen.items():
        if k[0] == j:
            pl_qgen[j - 1].append(py.value(v))

for j in model.b:
    for k, v in model.v.items():
        if k[0] == j:
            pl_v[j - 1].append(py.value(v))

for j in model.b:
    for k, v in model.theta.items():
        if k[0] == j:
            pl_theta[j - 1].append(py.value(v))

for j in model.b:
    for k, v in model.pdmand.items():
        if k[0] == j:
            pl_pdem[j - 1].append(py.value(v))

for j in model.b:
    for k, v in model.qdmand.items():
        if k[0] == j:
            pl_qdem[j - 1].append(py.value(v))

# width = 0.2
b = np.arange(len(model.b))
t = np.arange(len(model.t))
x = b  # number of buss
y = t  # time
z = pl_v  # votlage


#3d voltage profile
# X, Y = np.meshgrid(y, x)
# Z = np.reshape(z, X.shape)  # Z.shape must be equal to X.shape = Y.shape
# fig = plt.figure()
# # ax = fig.gca(projection='3d')
# # fig = plt.figure()
# ax = fig.add_subplot(projection='3d')
# ax.plot_surface(X, Y, Z, cmap=cm.coolwarm)
# ax.set_xlabel('time')
# ax.set_ylabel('bus')
# ax.set_zlabel('voltage')
# plt.show()



# all graphs
fig2, ax = plt.subplots(5, len(model.b), constrained_layout=True, figsize=(10, 10))
for k in range(len(model.b)):
    ax[0, k].plot(t, pl_v[k], color="blue")
    # ax[1, k].plot(t, pl_theta[k], color="blue")
    ax[1, k].bar(t, pl_pdem[k], color="green")
    ax[2, k].bar(t, pl_qdem[k], color="black")
    ax[3, k].bar(t, pl_pgen[k], color="red")
    ax[4, k].bar(t, pl_qgen[k], color="brown")

    ax[0, 0].set_ylabel('Voltage (PU)')
    # ax[1, 0].set_ylabel('Angle')
    ax[1, 0].set_ylabel('P_dem (PU)')
    ax[2, 0].set_ylabel('Q_dem (PU)')
    ax[3, 0].set_ylabel('P_gen (PU)')
    ax[4, 0].set_ylabel('Q_gen (PU)')

    ax[0, k].set_title(f"Bus {k + 1}")
    fig2.suptitle('Grid Data', fontsize=10)
pl.show()
# ##
#