1gktfkct.h11~

# reference: https://colab.research.google.com/github/kLabUM/pystorms/blob/master/tutorials/Scenario_Gamma.ipynb

import sys
#from modpods import topo_from_pystorms
#sys.path.append("G:/My Drive/modpods")
import modpods
import pystorms
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import control as ct
import dill as pickle
use_blind = False

'''
# uncontrolled
env = pystorms.scenarios.gamma()
done = False

while not done:
    # Query the current state of the simulation
    state = env.state()
    
    # Initialize actions to have each asset open
    actions = np.ones(11)
    
    # Set the actions and progress the simulation
    done = env.step(actions)
    
# Calculate the performance measure for the uncontrolled simulation
uncontrolled_perf = sum(env.data_log["performance_measure"])

print("The calculated performance for the uncontrolled case of Scenario gamma is:")
print("{}.".format(uncontrolled_perf))

basin_max_depths = [5., 10., 10., 10.]

plt.figure(figsize=(15,6))
plt.subplot(1,2,1)
plt.plot(np.asarray(env.data_log['depthN']['1'])/basin_max_depths[0], label='Basin 1')
plt.plot(np.asarray(env.data_log['depthN']['2'])/basin_max_depths[1], label='Basin 2')
plt.plot(np.asarray(env.data_log['depthN']['3'])/basin_max_depths[2], label='Basin 3')
plt.plot(np.asarray(env.data_log['depthN']['4'])/basin_max_depths[3], label='Basin 4')
plt.xlabel('Simulation Timestep')
plt.ylabel('Filling Degree')
plt.legend()

plt.subplot(1,2,2)
plt.plot(env.data_log['flow']['O1'], label='Basin 1')
plt.plot(env.data_log['flow']['O2'], label='Basin 2')
plt.plot(env.data_log['flow']['O3'], label='Basin 3')
plt.plot(env.data_log['flow']['O4'], label='Basin 4')
plt.xlabel('Simulation Timestep')
plt.ylabel('Basin Outflow (cfs)')
plt.legend()
plt.show()

uncontrolled_data = env.data_log
'''
def controller_efd(state, max_depths):
    # Initialize the action space so that we can compute the new settings
    new_settings = np.ones(len(state))
    # Set equal filling degree parameters
    c = 1.5
    theta = 0.25
    
    # Assign the current depth in each basin
    depths = state
    
    # Compute the filling degrees
    fd = depths/max_depths
    # Compute the average filling degree across each controlled basin
    fd_average = sum(fd)/len(fd)
    
    # Update each valve setting based on the relative fullness of each basin
    for i in range(0,len(fd)):
        
        # If a basin is very full compared to the average, we should open its
        # valve to release some water
        if fd[i] > fd_average:
            new_settings[i] = c*(fd[i]-fd_average)
        
        # If a basin's filling degree is close to the average (within some value
        # theta), its setting can be close to that average
        elif fd_average-fd[i] <= theta:
            new_settings[i] = fd_average
            
        # If a basin is very empty compared to the average, we can close its
        # valve to store more water at that location, prioritizing releasing at
        # the other locations
        else:
            new_settings[i] = 0.
        
        # Make sure the settings are in bounds [0,1]
        new_settings[i] = min(new_settings[i], 1.)
        new_settings[i] = max(new_settings[i], 0.)

    return new_settings



env = pystorms.scenarios.gamma()
done = False

# Specify the maximum depths for each basin we are controlling
basin_max_depths = [5., 10., 10., 10.]

while not done:
    # Query the current state of the simulation
    state = env.state()
    # Isolate only the states that we need (the 4 downstream basin depths)
    states_relevant = state[0:4]
    
    # Pass the current, relevant states and the maximum basin 
    # depths into our equal filling degree logic
    actions_efd = controller_efd(states_relevant, basin_max_depths)
    # Specify that the other 7 valves in the network should be 
    # open since we are not controlling them here
    actions_uncontrolled = np.ones(7)
    # Join the two above action arrays
    actions = np.concatenate((actions_efd, actions_uncontrolled), axis=0)
    
    # Set the actions and progress the simulation
    done = env.step(actions)
    
# Calculate the performance measure for the uncontrolled simulation
equalfilling_perf = sum(env.data_log["performance_measure"])

print("The calculated performance for the equal filling degree case of Scenario gamma is:")
print("{}.".format(equalfilling_perf))
'''
plt.figure(figsize=(15,6))
plt.subplot(1,2,1)
plt.plot(np.asarray(env.data_log['depthN']['1'])/basin_max_depths[0], label='Basin 1')
plt.plot(np.asarray(env.data_log['depthN']['2'])/basin_max_depths[1], label='Basin 2')
plt.plot(np.asarray(env.data_log['depthN']['3'])/basin_max_depths[2], label='Basin 3')
plt.plot(np.asarray(env.data_log['depthN']['4'])/basin_max_depths[3], label='Basin 4')
plt.xlabel('Simulation Timestep')
plt.ylabel('Filling Degree')
plt.legend()

plt.subplot(1,2,2)
plt.plot(env.data_log['flow']['O1'], label='Basin 1')
plt.plot(env.data_log['flow']['O2'], label='Basin 2')
plt.plot(env.data_log['flow']['O3'], label='Basin 3')
plt.plot(env.data_log['flow']['O4'], label='Basin 4')
plt.xlabel('Simulation Timestep')
plt.ylabel('Basin Outflow (cfs)')
plt.legend()
'''
ef_data = env.data_log

# get the responses into one dataframe with zero pads in between 
#uncontrolled_flows = pd.DataFrame.from_dict(uncontrolled_data['flow'])
#uncontrolled_depthN = pd.DataFrame.from_dict(uncontrolled_data['depthN'])
#uncontrolled_response = pd.concat([uncontrolled_flows, uncontrolled_depthN], axis=1)
#print(uncontrolled_response)
ef_flows = pd.DataFrame.from_dict(ef_data['flow'])
ef_flows.columns = env.config['action_space']
ef_depthN = pd.DataFrame.from_dict(ef_data['depthN'])
ef_depthN.columns = env.config['states']
ef_response = pd.concat([ef_flows, ef_depthN], axis=1)
ef_response.index = env.data_log['simulation_time']
print(ef_response)
# for the columns of ef_response which do not contain the string "O" (i.e. the depths)
# if the current column name is "X", make it "(X, depthN)"
# this is because the modpods library expects the depths to be named "X, depthN"
# where X is the name of the corresponding flow
for col in ef_response.columns:
    if "O" in col: # the orifices
        # if there's a number in that column name that's greater than 4, drop this column
        # this is because we're only controlling the first 4 orifices and these measurements are redundant to the storage node depths if the orifices are always open
        if int(col[1:]) > 4:
            ef_response.drop(columns=col, inplace=True)    

        ef_response.rename(columns={col: (col, "flow")}, inplace=True)
        

    
print(ef_response)



# for debugging resample to a coarser time step (native resolution is about one minute but not consistent)
# need a consistent time step for modpods
ef_response = ef_response.resample('5T',axis='index').mean().copy(deep=True)
print(ef_response)
# we'll only use the ef response to infer the topology and dynamics
# for this experiment assume all of flow O1-O11 and depth 1-11 are observable
# but only O1-O4 are controllable 
independent_columns = ef_response.columns[0:4] # orifices O1 through O4
dependent_columns = ef_response.drop(columns=independent_columns).columns


# learn the topology from the data
# this will be the "blind" plant model
if use_blind: # don't have this on all the time because it's very expensive
    blind_topo = modpods.infer_causative_topology(ef_response, dependent_columns = dependent_columns,
                                                  independent_columns = independent_columns, verbose=True)

    print(blind_topo.causative_topo)
    print(blind_topo.total_graph)


# read the topology from the swmm file (this is much cheaper)
env.config['states'] = dependent_columns
env.config['action_space'] = independent_columns 
# the default is controlling all 11 orifices so we need to edit the environment
print("defining topology")
swmm_topo = modpods.topo_from_pystorms(env)
# the index of the causative topology should be the dependent columns
#swmm_topo.index = dependent_columns
# the columns of the causative topology should be the dependent columns plus the independent columns
#swmm_topo.columns = dependent_columns.append(independent_columns)

# show all columns when printing dataframes
pd.set_option('display.max_columns', None)
print(swmm_topo)

if use_blind:
    print("differences in topology")
    print(blind_topo.causative_topo.compare(swmm_topo))

# learn the dynamics now, or load a previously learned model
'''
# learn the dynamics from the efd response
print("learning dynamics")
lti_plant_approx_seeing = modpods.lti_system_gen(swmm_topo, ef_response, 
                                                 independent_columns= independent_columns,
                                                 dependent_columns = dependent_columns, max_iter = 0,
                                                 swmm=True,bibo_stable=True)
# pickle the plant approximation to load later
with open('G:/My Drive/modpods/swmm_lti_plant_approx_seeing.pickle', 'wb') as handle:
    pickle.dump(lti_plant_approx_seeing, handle)
'''

# load the plant approximation from a pickle
with open('G:/My Drive/modpods/swmm_lti_plant_approx_seeing.pickle', 'rb') as handle:
    print("loading previously trained model")
    lti_plant_approx_seeing = pickle.load(handle)
    
if use_blind:
    lti_plant_approx_blind = modpods.lti_system_gen(blind_topo, ef_response, 
                                                 independent_columns= independent_columns,
                                                 dependent_columns = dependent_columns)



# cast the columns of dataframes to strings for easier indexing
ef_response.columns = ef_response.columns.astype(str)
dependent_columns = [str(col) for col in dependent_columns]
independent_columns = [str(col) for col in independent_columns]
# reindex the ef_response to an integer step
ef_response.index = np.arange(0,len(ef_response),1)

# evaluate the plant approximation accuracy
# only plot the depths at 1, 2, 3, and 4
# the forcing is the flows at O1, O2, O3, and O4
approx_response = ct.forced_response(lti_plant_approx_seeing['system'], U=np.transpose(ef_response[independent_columns].values), T=ef_response.index.values)
approx_data = pd.DataFrame(index=ef_response.index.values)
approx_data[dependent_columns[0]] = approx_response.outputs[0][:]
approx_data[dependent_columns[1]] = approx_response.outputs[1][:]
approx_data[dependent_columns[2]] = approx_response.outputs[2][:]
approx_data[dependent_columns[3]] = approx_response.outputs[3][:]

output_columns = dependent_columns[0:4] # depth at 1,2,3,4

# create a vertical subplot of 3 axes
fig, axes = plt.subplots(4, 1, figsize=(10, 10))

for idx in range(len(output_columns)):
    axes[idx].plot(ef_response[output_columns[idx]],label='actual')
    axes[idx].plot(approx_data[output_columns[idx]],label='approx')
    if idx == 0:
        axes[idx].legend(fontsize='x-large',loc='best')
    axes[idx].set_ylabel(output_columns[idx],fontsize='large')
    if idx == len(output_columns)-1:
        axes[idx].set_xlabel("time",fontsize='x-large')
# label the left column of plots "training"
axes[0].set_title("outputs",fontsize='xx-large')

plt.savefig("G:/My Drive/modpods/test_lti_control_of_swmm_plant_approx.png")
plt.savefig("G:/My Drive/modpods/test_lti_control_of_swmm_plant_approx.svg")
#plt.show()

# same plot, but just the first few timesteps (this is the accuracy that matters for feedback control)
# create a vertical subplot of 3 axes
fig, axes = plt.subplots(4, 1, figsize=(10, 10))

for idx in range(len(output_columns)):
    axes[idx].plot(ef_response[output_columns[idx]][:10],label='actual')
    axes[idx].plot(approx_data[output_columns[idx]][:10],label='approx')
    if idx == 0:
        axes[idx].legend(fontsize='x-large',loc='best')
    axes[idx].set_ylabel(output_columns[idx],fontsize='large')
    if idx == len(output_columns)-1:
        axes[idx].set_xlabel("time",fontsize='x-large')
# label the left column of plots "training"
axes[0].set_title("outputs",fontsize='xx-large')

plt.savefig("G:/My Drive/modpods/test_lti_control_of_swmm_plant_approx_first10.png")
plt.savefig("G:/My Drive/modpods/test_lti_control_of_swmm_plant_approx_first10.svg")
#plt.show()


# define the cost function
Q = np.eye(len(lti_plant_approx_seeing['A'].columns)) / 10e12 # we don't want to penalize the transition states as their magnitude doesn't have directly tractable physical meaning
# bryson's rule based on the maxiumum depth of each basin
# note: the swmm file gamma.inp actually specifies the maximum depth of storage node 1 as 10 feet, 
# but i'll be consistent with the configuration of the efd controller for consistent comparison
basin_max_depths_all = [5.0, 10.0, 10.0, 10.0, 10.0,20.0, 10.0, 10.0,10.0, 13.72, 14.96]
for asset_index in range(len(dependent_columns)):
    Q[lti_plant_approx_seeing['A'].columns.get_loc(dependent_columns[asset_index]),lti_plant_approx_seeing['A'].columns.get_loc(dependent_columns[asset_index])] = 1 / ((basin_max_depths_all[asset_index])**2 )

# threshold on flows at 0.11 m^3 / s which is 3.9 cfs
R = np.eye(len(lti_plant_approx_seeing['B'].columns)) / (3.9**2) # bryson's rule on maximum allowable flow
# define the system
# sys_response_to_control = ct.ss(lti_plant_approx_seeing['A'],lti_plant_approx_seeing['B'],lti_plant_approx_seeing['C'],0) # just for defining the controller gain 
# (not necessary in this case, but would be if you've got disturbances)
# find the state feedback gain for the linear quadratic regulator
K,S,E = ct.lqr(lti_plant_approx_seeing['system'],Q,R) # one row of K should be zeros to reflect that u1 is not used as a control but is the disturbance

# define the observer gain
L,P,E = ct.lqe(lti_plant_approx_seeing['system'],np.eye(len(lti_plant_approx_seeing['B'].columns)),np.eye(len(lti_plant_approx_seeing['C'].index)) ) # unit covariance on process noise and measurement error 

'''
# define the observer based compensator (per freudenberg 560 course notes 2.4)
obc_A = lti_plant_approx_seeing['A'].values-lti_plant_approx_seeing['B'].values@K - L@lti_plant_approx_seeing['C'].values
# ingests measurements, returns control actions
obc = ct.ss(obc_A, L, -K, 0, inputs=list(lti_plant_approx_seeing['C'].index),outputs=list(lti_plant_approx_seeing['B'].columns)) # negate K to give back commanded flows which are positive


# need to separately define the observer and controller because we can't close the loop in the typical way
# the observer takes the control input and measured output as inputs and outputs the estimated full state
#observer_input = np.concatenate((lti_plant_approx_seeing['B'].values,L),axis = 1)
#observer = ct.ss(lti_plant_approx_seeing['A'] - L@lti_plant_approx_seeing['C'].values, observer_input, np.eye(len(lti_plant_approx_seeing['A']) ) , 0 )

# the controller takes in an estiamte of the state and returns a control command
# this is just, u = -K @ xhat, not necessary to define a state space model for that as there's no evolution

# can't form the closed loop system because the plant is not a state space system, but rather a software model
# the state estimate and control actions will be computed iteratively as the simulation is stepped through
'''
env = pystorms.scenarios.gamma()
done = False
u = np.zeros((len(lti_plant_approx_seeing['B'].columns),1) ) # start with all orifices completely closed
xhat = np.zeros((len(lti_plant_approx_seeing['A'].columns),1)) # initial state estimate
while not done:
    # Query the current state of the simulation
    observables = env.state()

    # update the observer based on these measurements (xhat_dot = (A-LC) xhat + B u + L y_m)
    state_evolution = (lti_plant_approx_seeing['A'] - L@lti_plant_approx_seeing['C'].values) @ xhat
    impact_of_control = lti_plant_approx_seeing['B'].values @ u
    output_updating = (L @ np.transpose(observables)).reshape((-1,1)) # provided as row vector, need a column vector. also need to reshape to 2d array with 1 column
    xhat_dot = state_evolution + impact_of_control + output_updating
    xhat += xhat_dot # update the state estimate
    yhat = lti_plant_approx_seeing['C'] @ xhat # just for reference, could be useful for plotting later
    u = -K @ xhat # calculate control command

    # convert control command (flow) into orifice open percentage
    # per the EPA-SWMM user manual volume ii hydraulics, orifices (section 6.2, page 107) - https://nepis.epa.gov/Exe/ZyPDF.cgi/P100S9AS.PDF?Dockey=P100S9AS.PDF 
    # all orifices in gamma are "bottom"
    Cd = 0.65 # happens to be the same for all of them
    Ao = 1 # area is one square foot. again, happens to be the same for all of them. 
    g = 32.2 # ft / s^2
    # the expression for discharge is found using Torricelli's equation: Q = Cd * (Ao*open_pct) sqrt(2*g*H_e)
    # H_e is the effective head in feet, which is just the depth in the basin as the orifices are "bottom"
    # to get the action command as a percent open, we solve as: open_pct = Q_desired / (Cd * Ao * sqrt(2*g*H_e))
    u_open_pct = u*-1
    for idx in range(len(u)):
        u_open_pct.iloc[idx,0] = u.iloc[idx,0] / (Cd*Ao * np.sqrt(2*g*observables[idx]))

   
    # Specify that the other 7 valves in the network should be 
    # open since we are not controlling them here
    actions_uncontrolled = np.ones(7)
    # Join the two above action arrays
    actions = np.concatenate((u_open_pct.flatten(), actions_uncontrolled), axis=0)
    
    # Set the actions and progress the simulation
    done = env.step(actions)

    
# Calculate the performance measure for the uncontrolled simulation
obc_perf = sum(env.data_log["performance_measure"])

print("The calculated performance for the observer based compensator case of Scenario gamma is:")
print("{}.".format(obc_perf))

plt.figure(figsize=(15,6))
plt.subplot(1,2,1)
plt.plot(np.asarray(env.data_log['depthN']['1'])/basin_max_depths[0], label='Basin 1')
plt.plot(np.asarray(env.data_log['depthN']['2'])/basin_max_depths[1], label='Basin 2')
plt.plot(np.asarray(env.data_log['depthN']['3'])/basin_max_depths[2], label='Basin 3')
plt.plot(np.asarray(env.data_log['depthN']['4'])/basin_max_depths[3], label='Basin 4')
plt.xlabel('Simulation Timestep')
plt.ylabel('Filling Degree')
plt.legend()

plt.subplot(1,2,2)
plt.plot(env.data_log['flow']['O1'], label='Basin 1')
plt.plot(env.data_log['flow']['O2'], label='Basin 2')
plt.plot(env.data_log['flow']['O3'], label='Basin 3')
plt.plot(env.data_log['flow']['O4'], label='Basin 4')
plt.xlabel('Simulation Timestep')
plt.ylabel('Basin Outflow (cfs)')
plt.legend()


print("done")