content: # This file is part of GridCal.
  #
  # GridCal is free software: you can redistribute it and/or modify
  # it under the terms of the GNU General Public License as published by
  # the Free Software Foundation, either version 3 of the License, or
  # (at your option) any later version.
  #
  # GridCal is distributed in the hope that it will be useful,
  # but WITHOUT ANY WARRANTY; without even the implied warranty of
  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  # GNU General Public License for more details.
  #
  # You should have received a copy of the GNU General Public License
  # along with GridCal.  If not, see <http://www.gnu.org/licenses/>.
  
  import numpy as np
  import pandas as pd
  import numba as nb
  import time
  from warnings import warn
  import scipy.sparse as sp
  from scipy.sparse import coo_matrix, csc_matrix
  from scipy.sparse import hstack as hs, vstack as vs
  from scipy.sparse.linalg import factorized, spsolve, inv
  from matplotlib import pyplot as plt
  from GridCal.Engine import *
  
  
  def SysMat(Y, Ys, pq, pvpq):
      """
      Computes the system Jacobian matrix in polar coordinates
      Args:
          Ybus: Admittance matrix
          V: Array of nodal voltages
          Ibus: Array of nodal current injections
          pq: Array with the indices of the PQ buses
          pvpq: Array with the indices of the PV and PQ buses
  
      Returns:
          The system Jacobian matrix
      """
      A11 = -Ys.imag[np.ix_(pvpq, pvpq)]
      A12 = Y.real[np.ix_(pvpq, pq)]
      A21 = -Ys.real[np.ix_(pq, pvpq)]
      A22 = -Y.imag[np.ix_(pq, pq)]
  
      Asys = sp.vstack([sp.hstack([A11, A12]),
                        sp.hstack([A21, A22])], format="csc")
  
      return Asys
  
  
  def compute_acptdf(Ybus, Yseries, Yf, Yt, Cf, V, pq, pv, distribute_slack):
      """
      Compute the AC-PTDF
      :param Ybus: admittance matrix
      :param Yf: Admittance matrix of the buses "from"
      :param Yt: Admittance matrix of the buses "to"
      :param Cf: Connectivity branch - bus "from"
      :param V: voltages array
      :param Ibus: array of currents
      :param pq: array of pq node indices
      :param pv: array of pv node indices
      :return: AC-PTDF matrix (branches, buses)
      """
      n = len(V)
      pvpq = np.r_[pv, pq]
      npq = len(pq)
  
      # compute the Jacobian
      J = SysMat(Ybus, Yseries, pq, pvpq)
  
      if distribute_slack:
          dP = np.ones((n, n)) * (-1 / (n - 1))
          for i in range(n):
              dP[i, i] = 1.0
      else:
          dP = np.eye(n, n)
  
      # compose the compatible array (the Q increments are considered zero
      dQ = np.zeros((npq, n))
      # dQ = np.eye(n, n)[pq, :]
      dS = np.r_[dP[pvpq, :], dQ]
  
      # solve the voltage increments
      dx = spsolve(J, dS)
  
      # compute branch derivatives
      If = Yf * V
      E = V / np.abs(V)
      Vdiag = sp.diags(V)
      Vdiag_conj = sp.diags(np.conj(V))
      Ediag = sp.diags(E)
      Ediag_conj = sp.diags(np.conj(E))
      If_diag_conj = sp.diags(np.conj(If))
  
      Yf_conj = Yf.copy()
      Yf_conj.data = np.conj(Yf_conj.data)
      Yt_conj = Yt.copy()
      Yt_conj.data = np.conj(Yt_conj.data)
  
      dSf_dVa = 1j * (If_diag_conj * Cf * Vdiag - sp.diags(Cf * V) * Yf_conj * Vdiag_conj)
      dSf_dVm = If_diag_conj * Cf * Ediag - sp.diags(Cf * V) * Yf_conj * Ediag_conj
  
      # compose the final AC-PTDF
      dPf_dVa = dSf_dVa.real[:, pvpq]
      dPf_dVm = dSf_dVm.real[:, pq]
      PTDF = sp.hstack((dPf_dVa, dPf_dVm)) * dx
  
      return PTDF
  
  
  def make_lodf(circuit: SnapshotCircuit, PTDF, correct_values=True):
      """
  
      :param circuit:
      :param PTDF: PTDF matrix in numpy array form
      :return:
      """
      nl = circuit.nbr
  
      # compute the connectivity matrix
      Cft = circuit.C_branch_bus_f - circuit.C_branch_bus_t
  
      H = PTDF * Cft.T
  
      # old code
      # h = sp.diags(H.diagonal())
      # LODF = H / (np.ones((nl, nl)) - h * np.ones(nl))
  
      # divide each row of H by the vector 1 - H.diagonal
      # LODF = H / (1 - H.diagonal())
      # replace possible nan and inf
      # LODF[LODF == -np.inf] = 0
      # LODF[LODF == np.inf] = 0
      # LODF = np.nan_to_num(LODF)
  
      # this loop avoids the divisions by zero
      # in those cases the LODF column should be zero
      LODF = np.zeros((nl, nl))
      div = 1 - H.diagonal()
      for j in range(H.shape[1]):
          if div[j] != 0:
              LODF[:, j] = H[:, j] / div[j]
  
      # replace the diagonal elements by -1
      # old code
      # LODF = LODF - sp.diags(LODF.diagonal()) - sp.eye(nl, nl), replaced by:
      for i in range(nl):
          LODF[i, i] = - 1.0
  
      if correct_values:
          i1, j1 = np.where(LODF > 1)
          for i, j in zip(i1, j1):
              LODF[i, j] = 1
  
          i2, j2 = np.where(LODF < -1)
          for i, j in zip(i2, j2):
              LODF[i, j] = -1
  
      return LODF
  
  
  def get_branch_time_series(circuit: TimeCircuit, PTDF):
      """
  
      :param grid:
      :return:
      """
  
      # option 2: call the power directly
      P = circuit.Sbus.real
      Pbr = np.dot(PTDF, P).T * circuit.Sbase
  
      return Pbr
  
  
  def multiple_failure_old(flows, LODF, beta, delta, alpha):
      """
  
      :param flows: array of all the pre-contingency flows
      :param LODF: Line Outage Distribution Factors Matrix
      :param beta: index of the first failed line
      :param delta: index of the second failed line
      :param alpha: index of the line where you want to see the effects
      :return: post contingency flow in the line alpha
      """
      # multiple contingency matrix
      M = np.ones((2, 2))
      M[0, 1] = -LODF[beta, delta]
      M[1, 0] = -LODF[delta, beta]
  
      # normal flows of the lines beta and delta
      F = flows[[beta, delta]]
  
      # contingency flows after failing the ines beta and delta
      Ff = np.linalg.solve(M, F)
  
      # flow delta in the line alpha after the multiple contingency of the lines beta and delta
      L = LODF[alpha, :][[beta, delta]]
      dFf_alpha = np.dot(L, Ff)
  
      return F[alpha] + dFf_alpha
  
  
  def multiple_failure(flows, LODF, failed_idx):
      """
      From the paper:
      Multiple Element Contingency Screening
          IEEE TRANSACTIONS ON POWER SYSTEMS, VOL. 26, NO. 3, AUGUST 2011
          C. Matthew Davis and Thomas J. Overbye
      :param flows: array of all the pre-contingency flows (the base flows)
      :param LODF: Line Outage Distribution Factors Matrix
      :param failed_idx: indices of the failed lines
      :return: all post contingency flows
      """
      # multiple contingency matrix
      M = -LODF[np.ix_(failed_idx, failed_idx)]
      for i in range(len(failed_idx)):
          M[i, i] = 1.0
  
      # normal flows of the failed lines indicated by failed_idx
      F = flows[failed_idx]
  
      # Affected flows after failing the lines indicated by failed_idx
      Ff = np.linalg.solve(M, F)
  
      # flow delta in the line alpha after the multiple contingency of the lines indicated by failed_idx
      L = LODF[:, failed_idx]
      dFf_alpha = np.dot(L, Ff)
  
      # return the final contingency flow as the base flow plus the contingency flow delta
      return flows + dFf_alpha
  
  
  def get_n_minus_1_flows(circuit: MultiCircuit):
  
      opt = PowerFlowOptions()
      branches = circuit.get_branches()
      m = circuit.get_branch_number()
      Pmat = np.zeros((m, m))  # monitored, contingency
  
      for c, branch in enumerate(branches):
  
          if branch.active:
              branch.active = False
  
              pf = PowerFlowDriver(circuit, opt)
              pf.run()
              Pmat[:, c] = pf.results.Sbranch.real
  
              branch.active = True
  
      return Pmat
  
  
  def check_lodf(grid: MultiCircuit):
  
      flows_n1_nr = get_n_minus_1_flows(grid)
  
      # assume 1 island
      nc = compile_snapshot_circuit(grid)
      islands = split_into_islands(nc)
      circuit = islands[0]
  
      pf_driver = PowerFlowDriver(grid, PowerFlowOptions())
      pf_driver.run()
  
      PTDF = compute_acptdf(Ybus=circuit.Ybus,
                            Yseries=circuit.Yseries,
                            Yf=circuit.Yf,
                            Yt=circuit.Yt,
                            Cf=circuit.C_branch_bus_f,
                            V=pf_driver.results.voltage,
                            pq=circuit.pq,
                            pv=circuit.pv,
                            distribute_slack=True)
      LODF = make_lodf(circuit, PTDF)
  
      Pbus = circuit.get_injections(False).real
      flows_n = np.dot(PTDF, Pbus)
  
      nl = circuit.nbr
      flows_n1 = np.zeros((nl, nl))
      for c in range(nl):  # branch that fails (contingency)
          # for m in range(nl):  # branch to monitor
          #     flows_n1[m, c] = flows_n[m] + LODF[m, c] * flows_n[c]
          flows_n1[:, c] = flows_n[:] + LODF[:, c] * flows_n[c]
  
      return flows_n, flows_n1_nr, flows_n1
  
  
  def test_ptdf(grid):
      """
      Sigma-distances test
      :param grid:
      :return:
      """
      nc = compile_snapshot_circuit(grid)
      islands = split_into_islands(nc)
      circuit = islands[0]  # pick the first island
  
      pf_driver = PowerFlowDriver(grid, PowerFlowOptions())
      pf_driver.run()
  
      PTDF = compute_acptdf(Ybus=circuit.Ybus,
                            Yseries=circuit.Yseries,
                            Yf=circuit.Yf,
                            Yt=circuit.Yt,
                            Cf=circuit.C_branch_bus_f,
                            V=pf_driver.results.voltage,
                            pq=circuit.pq,
                            pv=circuit.pv,
                            distribute_slack=False)
  
      print('PTDF:')
      print(PTDF)
  
  
  if __name__ == '__main__':
      from GridCal.Engine import FileOpen
      import pandas as pd
  
      np.set_printoptions(threshold=sys.maxsize, linewidth=200000000)
      # np.set_printoptions(linewidth=2000, suppress=True)
      pd.set_option('display.max_rows', 500)
      pd.set_option('display.max_columns', 500)
      pd.set_option('display.width', 1000)
  
      # fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/IEEE39_1W.gridcal'
      # fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/IEEE 14.xlsx'
      # fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/lynn5buspv.xlsx'
      # fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/IEEE 118.xlsx'
      fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/1354 Pegase.xlsx'
      # fname = 'helm_data1.gridcal'
      # fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/IEEE 14 PQ only.gridcal'
      # fname = 'IEEE 14 PQ only full.gridcal'
      # fname = '/home/santi/Descargas/matpower-fubm-master/data/case5.m'
      # fname = '/home/santi/Descargas/matpower-fubm-master/data/case30.m'
      # fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/PGOC_6bus.gridcal'
      grid_ = FileOpen(fname).open()
  
      test_ptdf(grid_)
      name = os.path.splitext(fname.split(os.sep)[-1])[0]
      method = 'ACPTDF (No Jacobian, V=Vpf)'
      nc_ = compile_snapshot_circuit(grid_)
      islands_ = split_into_islands(nc_)
      circuit_ = islands_[0]
  
      pf_driver_ = PowerFlowDriver(grid_, PowerFlowOptions())
      pf_driver_.run()
  
      H_ = compute_acptdf(Ybus=circuit_.Ybus,
                          Yseries=circuit_.Yseries,
                          Yf=circuit_.Yf,
                          Yt=circuit_.Yt,
                          Cf=circuit_.C_branch_bus_f,
                          V=pf_driver_.results.voltage,
                          pq=circuit_.pq,
                          pv=circuit_.pv,
                          distribute_slack=False)
  
      LODF_ = make_lodf(circuit_, H_)
  
      if H_.shape[0] < 50:
          print('PTDF:\n', H_)
          print('LODF:\n', LODF_)
  
      flows_n_, flows_n1_nr_, flows_n1_ = check_lodf(grid_)
  
      # in the case of the grid PGOC_6bus
      flows_multiple = multiple_failure(flows=flows_n_,
                                        LODF=LODF_,
                                        failed_idx=[1, 5])  # failed lines 2 and 6
  
      Pn1_nr_df = pd.DataFrame(data=flows_n1_nr_, index=nc_.branch_names, columns=nc_.branch_names)
      flows_n1_df = pd.DataFrame(data=flows_n1_, index=nc_.branch_names, columns=nc_.branch_names)
  
      # plot N-1
      fig = plt.figure(figsize=(12, 8))
      title = 'N-1 with ' + method + ' (' + name + ')'
      fig.suptitle(title)
      ax1 = fig.add_subplot(221)
      ax2 = fig.add_subplot(222)
      ax3 = fig.add_subplot(223)
  
      Pn1_nr_df.plot(ax=ax1, legend=False)
      flows_n1_df.plot(ax=ax2, legend=False)
      diff = Pn1_nr_df - flows_n1_df
      diff.plot(ax=ax3, legend=False)
  
      ax1.set_title('Newton-Raphson N-1 flows')
      ax2.set_title('PTDF N-1 flows')
      ax3.set_title('Difference')
      fig.savefig(title + '.png')
  
      # ------------------------------------------------------------------------------------------------------------------
      # Perform real time series
      # ------------------------------------------------------------------------------------------------------------------
      if grid_.time_profile is not None:
          grid_.ensure_profiles_exist()
          nc_ts = compile_time_circuit(grid_)
          islands_ts = split_time_circuit_into_islands(nc_ts)
          circuit_ts = islands_ts[0]
  
          pf_options = PowerFlowOptions()
          ts_driver = TimeSeries(grid=grid_, options=pf_options)
          ts_driver.run()
          Pbr_nr = ts_driver.results.Sbranch.real
          df_Pbr_nr = pd.DataFrame(data=Pbr_nr, columns=circuit_ts.branch_names, index=circuit_ts.time_array)
  
          # Compute the PTDF based flows
          Pbr_ptdf = get_branch_time_series(circuit=circuit_ts, PTDF=H_)
          df_Pbr_ptdf = pd.DataFrame(data=Pbr_ptdf, columns=circuit_ts.branch_names, index=circuit_ts.time_array)
  
          # plot
          fig = plt.figure(figsize=(12, 8))
          title = 'Flows with ' + method + ' (' + name + ')'
          fig.suptitle(title)
          ax1 = fig.add_subplot(221)
          ax2 = fig.add_subplot(222)
          ax3 = fig.add_subplot(223)
  
          df_Pbr_nr.plot(ax=ax1, legend=False)
          df_Pbr_ptdf.plot(ax=ax2, legend=False)
          diff = df_Pbr_nr - df_Pbr_ptdf
          diff.plot(ax=ax3, legend=False)
  
          ax1.set_title('Newton-Raphson flows')
          ax2.set_title('PTDF flows')
          ax3.set_title('Difference')
          fig.savefig(title + '.png')
  
      plt.show()
  
  