import pandas as pd
import numpy as np

from psst.network.graph import (
    NetworkModel, NetworkViewBase, NetworkView
)

from psst.case import read_matpower
case = read_matpower('../cases/case118.m')

"""
Explanation: How NetworkModel Works
End of explanation
"""


# Create the model from a PSSTCase, optionally passing a sel_bus
m = NetworkModel(case, sel_bus='Bus1')

"""
Explanation: I. Creating a NetworkModel
End of explanation
"""


display(m.case)     # saves the case
display(m.network)  # creates a PSSTNetwork
display(m.G)        # stores the networkX graph (an attribute of the PSSTNetwork)
display(m.model)    # builds/solves the model

# Creates df of x,y positions for each node (bus, load, gen), based off self.network.positions
m.all_pos.head(n=10)

# Creates a df of start and end x,y positions for each edge, based off self.G.edges()
m.all_edges.head(n=10)

"""
Explanation: In the __init__, the NetworkModel...
End of explanation
"""


# `sel_bus` is a single bus, upon which the visualization is initially centered.
# It can be changed programatically, or via the dropdown menu.

m.sel_bus

# At first, it is the only bus in view_buses.
# More buses get added to view_buses as they are clicked.

m.view_buses

"""
Explanation: The sel_bus and view_buses attributes
End of explanation
"""


# Create the view from the model
# (It can, alternatively, be created from a case.)

v = NetworkView(model=m)

v

"""
Explanation: II. Creating a NetworkView from the model
End of explanation
"""


# The subsetting that occurs is all based on `view_buses`
m.view_buses

"""
Explanation: III. Generating the x,y data for the view

Whenever the view_buses list get changed, it triggers the callback _callback_view_change
This function first calls subset_positions and subset_edges
Then, the subsetted DataFrames get segregated into seperate ones for bus, gen, and load
Finally, the x,y coordinates are extracted into a format the NetworkView can use.
End of explanation
"""


# Subset positions creates self.pos
m.pos

"""
Explanation: The subset_positions() call
End of explanation
"""


# Subset edges creates self.edges
m.edges

"""
Explanation: The function looks like this:
python
def subset_positions(self):
    """Subset self.all_pos to include only nodes adjacent to those in view_buses list."""
    nodes = [list(self.G.adj[item].keys()) for item in self.view_buses]  # get list of nodes adj to selected buses
    nodes = set(itertools.chain.from_iterable(nodes))  # chain lists together, eliminate duplicates w/ set
    nodes.update(self.view_buses)  # Add the view_buses themselves to the set
    return self.all_pos.loc[nodes]  # Subset df of all positions to include only desired nodes.
The subset_edges() call
End of explanation
"""


m.view_buses = ['Bus2','Bus3']

edge_list = m.G.edges(nbunch=m.view_buses)  # get edges of view_buses as list of tuples
edge_list

edges_fwd = m.all_edges.loc[edge_list]  # query all_pos with edge_list
edges_fwd

edge_list_rev = [tuple(reversed(tup)) for tup in edge_list]  # reverse order of each tuple
edge_list_rev

edges_rev = m.all_edges.loc[edge_list_rev]  # query all_pos again, with reversed edge_list
edges_rev

edges = edges_fwd.append(edges_rev).dropna(subset=['start_x'])  # combine results, dropping false hits
edges

"""
Explanation: The function looks like this:
python
def subset_edges(self):
    """Subset all_edges, with G.edges() info, based on view_buses list."""
    edge_list = self.G.edges(nbunch=self.view_buses)  # get edges of view_buses as list of tuples
    edges_fwd = self.all_edges.loc[edge_list]  # query all_pos with edge_list
    edge_list_rev = [tuple(reversed(tup)) for tup in edge_list]  # reverse order of each tuple
    edges_rev = self.all_edges.loc[edge_list_rev]  # query all_pos again, with reversed edge_list
    edges = edges_fwd.append(edges_rev).dropna(subset=['start_x'])  # combine results, dropping false hits
    return edges
 If you want a closer look...
End of explanation
"""


print("x_vals: ", m.bus_x_vals)
print("y_vals: ", m.bus_y_vals)
print("names:  ", m.bus_names)

"""
Explanation: Segregating DataFrames and extracting data

The DataFrames are segregated into bus, case, and load, using the names in case.bus, case.gen, and case.load
x,y data is extracted, ready to be plotted by NetworkView


Extracting bus data looks like this:
python
bus_pos = self.pos[self.pos.index.isin(self.case.bus_name)]
self.bus_x_vals = bus_pos['x']
self.bus_y_vals = bus_pos['y']
self.bus_names = list(bus_pos.index)
(Similar for the other nodes)
End of explanation
"""


print("bus_x_edges:")
print(m.bus_x_edges)

print("\nbus_y_edges:")
print(m.bus_y_edges)

"""
Explanation: Extracting branch data looks like this:
```python
edges = self.edges.reset_index()
_df = edges.loc[edges.start.isin(self.case.bus_name) & edges.end.isin(self.case.bus_name)]
self.bus_x_edges = [tuple(edge) for edge in _df[['start_x', 'end_x']].values]
self.bus_y_edges = [tuple(edge) for edge in _df[['start_y', 'end_y']].values]
```
(Similar for the other edges)
End of explanation
"""


