testing.py

from sympy import ff
import hydra
from omegaconf import DictConfig
import torch
import json
import numpy as np
from hydra import initialize, compose
from src.algos.registry import get_model
import os 
import copy
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib as mpl
import time
from collections import defaultdict

RUN_TIME = time.strftime("%Y%m%d-%H%M%S")

def setup_sumo(cfg):
    from src.envs.sim.sumo_env import Scenario, AMoD, GNNParser
    cfg.simulator.cplexpath = cfg.model.cplexpath
    if not cfg.simulator.directory:
        cfg.simulator.directory = f"{cfg.model.name}/{cfg.simulator.city}"
    cfg = cfg.simulator
    scenario_path = 'src/envs/data'
    cfg.sumocfg_file = f'{scenario_path}/{cfg.city}/{cfg.sumocfg_file}'
    cfg.net_file = f'{scenario_path}/{cfg.city}/{cfg.net_file}'
    demand_file = f'src/envs/data/scenario_lux{cfg.num_regions}.json'
    aggregated_demand = not cfg.random_od

    scenario = Scenario(
        num_cluster=cfg.num_regions, json_file=demand_file, aggregated_demand=aggregated_demand,
        sumo_net_file=cfg.net_file, acc_init=cfg.acc_init, sd=cfg.seed, demand_ratio=cfg.demand_ratio,
        time_start=cfg.time_start, time_horizon=cfg.time_horizon, duration=cfg.duration,
        tstep=cfg.matching_tstep, max_waiting_time=cfg.max_waiting_time
    )
    env = AMoD(scenario, cfg=cfg, beta=cfg.beta)
    parser = GNNParser(env, T=cfg.time_horizon, json_file=demand_file)
    return env, parser
    
def setup_macro(cfg):
    from src.envs.sim.macro_env import Scenario, AMoD, GNNParser
    with open("src/envs/data/macro/calibrated_parameters.json", "r") as file:
        calibrated_params = json.load(file)
    cfg.simulator.cplexpath = cfg.model.cplexpath
    if not cfg.simulator.directory:
        cfg.simulator.directory = f"{cfg.model.name}/{cfg.simulator.city}"
    cfg = cfg.simulator
    city = cfg.city
    scenario = Scenario(
        json_file=f"src/envs/data/macro/scenario_{city}.json",
        demand_ratio=calibrated_params[city]["demand_ratio"],
        json_hr=calibrated_params[city]["json_hr"],
        sd=cfg.seed,
        json_tstep=calibrated_params[city]["test_tstep"],
        tf=cfg.max_steps,
    )
    env = AMoD(scenario, cfg = cfg, beta = calibrated_params[city]["beta"])
    parser = GNNParser(env, T=cfg.time_horizon, json_file=f"src/envs/data/macro/scenario_{city}.json")
    return env, parser

def setup_multi_macro(cfg):
    from src.envs.sim.multi_macro_env import Scenario, AMoD, GNNParser
    with open("src/envs/data/macro/calibrated_parameters.json", "r") as file:
        calibrated_params = json.load(file)
    cfg.simulator.cplexpath = cfg.model.cplexpath
    if not cfg.simulator.directory:
        cfg.simulator.directory = f"{cfg.model.name}/{cfg.simulator.city}"
    cfg = cfg.simulator
    city = cfg.city
    if cfg.constant_vehicle_count:
        supply_factor = cfg.firm_count
    else:
        supply_factor = 1
    scenario = Scenario(
        json_file=f"src/envs/data/macro/scenario_{city}.json",
        demand_ratio=calibrated_params[city]["demand_ratio"],
        json_hr=calibrated_params[city]["json_hr"],
        sd=cfg.seed,
        json_tstep=calibrated_params[city]["test_tstep"],
        tf=cfg.max_steps,
        supply_factor=supply_factor,
        firm_count=cfg.firm_count,
        demand_filter_type = cfg.demand_filter_type,
        initial_vehicle_distribution=cfg.initial_vehicle_distribution,
        pricing_model=cfg.pricing_model

    )

    filename = RUN_TIME + "_supply_factor_" + str(supply_factor) + "_firm_count_" + str(cfg.firm_count) + "_dm_" + str(cfg.demand_filter_type)
    save_sampled_demand(scenario.tripAttr, filename)
    env = AMoD(scenario, cfg=cfg, beta=calibrated_params[city]["beta"])
    parser = GNNParser(env, T=cfg.time_horizon, json_file=f"src/envs/data/macro/scenario_{city}.json")
    return env, parser

def setup_model(cfg, env, parser, device):
    model_name = cfg.model.name
    
    if model_name == "sac" or model_name =="cql":
        from src.algos.sac import SAC
        model= SAC(env=env, input_size=cfg.model.input_size, cfg=cfg.model, parser=parser, device=device).to(device)
        model.load_checkpoint(path=f"ckpt/{cfg.model.checkpoint_path}_best.pth")
        return model
    elif model_name == "a2c":
        from src.algos.a2c import A2C
        model = A2C(env=env, input_size=cfg.model.input_size,cfg=cfg.model, parser=parser, device=device).to(device)
        model.load_checkpoint(path=f"ckpt/{cfg.model.checkpoint_path}_best.pth")
        return model
    elif model_name == "iql":
        from src.algos.iql import IQL
        model = IQL(env=env, input_size=cfg.model.input_size,cfg=cfg.model, parser=parser, device=device).to(device)
        model.load_checkpoint(path=f"ckpt/{cfg.model.checkpoint_path}.pth")
        return model
    elif model_name == "bc":
        from src.algos.bc import BC
        model = BC(env=env, input_size=cfg.model.input_size,cfg=cfg.model, parser=parser, device=device).to(device)
        model.load_checkpoint(path=f"ckpt/{cfg.model.checkpoint_path}.pth")
        return model
    else:
        model_class = get_model(model_name)
        
        model_kwargs = {
            "cplexpath": cfg.simulator.cplexpath,
            "directory": cfg.simulator.directory,
            "T": cfg.simulator.time_horizon,
            "policy_name": cfg.model.name
        }
        for key, value in cfg.model.items():
            if key not in model_kwargs:
                model_kwargs[key] = value
        return model_class(**model_kwargs)

def multi_test(input_config):
    '''
    for Colab tutorial
    '''
    multi_config = [{**copy.deepcopy(input_config), "model.name" : model} for model in input_config["model.name"]]

    data = [{}, {}]
    for config in multi_config:

        with initialize(config_path="src/config"):
            cfg = compose(config_name="config", overrides= [f"{key}={value}" for key, value in config.items()])  # Load the configuration
            
        # Import simulator module based on the configuration
        simulator_name = cfg.simulator.name
        if simulator_name == "sumo":
            env, parser = setup_sumo(cfg)
        elif simulator_name == "macro":
            env, parser = setup_macro(cfg)
        elif simulator_name == "multi_macro":
            env, parser = setup_multi_macro(cfg)
        else:
            raise ValueError(f"Unknown simulator: {simulator_name}")
        
        use_cuda = not cfg.model.no_cuda and torch.cuda.is_available()
        device = torch.device("cuda" if use_cuda else "cpu")

        profit, inflows = test_approach(cfg, env, parser, device)
        data[0][config["model.name"]], data[1][config['model.name']] = profit, inflows

    control_data = get_no_control_performance(cfg, env, parser, device, use_saved_data=cfg.simulator.reuse_no_control)

    plot_comparison(cfg, env, control_data, data)

def save_vehicle_distribution(acc, file_str=None):
    """
    Save vehicle distribution across all timesteps and regions.
    
    Parameters:
        acc (defaultdict): Nested dict [region][timestep] = count.
        sim_time (str): Timestamp string for filename.
    """
    filename = f"saved_files/vehicle_distribution_{file_str}.json"
    if not os.path.exists("saved_files"):
        os.makedirs("saved_files")
    
    full_snapshot = defaultdict(dict)
    for region, timestep_dict in acc.items():
        for timestep, value in timestep_dict.items():
            full_snapshot[str(timestep)][str(region)] = float(value)
    
    with open(filename, "w") as f:
        json.dump(full_snapshot, f, indent=2)

    print(f"Saved full vehicle distribution to {filename}")

def save_sampled_demand(tripAttr, filename=None):
    # save the sampled demand to a file

    
    if not os.path.exists('saved_files'):
        os.makedirs('saved_files')
    with open(f'saved_files/sample_demand_{filename}.json', 'w') as f:
        print(f'Saving sampled demand to saved_files/sample_demand_{filename}.json')
        json.dump(tripAttr, f, indent=4)
        
def test_approach(cfg, env, parser, device):
    model = setup_model(cfg, env, parser, device)
    
    print(f'Testing model {cfg.model.name} on {cfg.simulator.name} environment')
  
    episode_reward, episode_served_demand, episode_rebalancing_cost, inflows = model.test(cfg.model.test_episodes, env)
    if cfg.simulator.constant_vehicle_count:
        supply_factor = cfg.simulator.firm_count
    else:
        supply_factor = 1
    file_str = RUN_TIME + "_" + str(cfg.model.name) + "_supply_factor_" + str(supply_factor) + "_firm_count_" + str(cfg.simulator.firm_count) + "_dm_" + str(cfg.simulator.demand_filter_type)
    save_vehicle_distribution(env.acc, file_str)


    print('Mean Episode Profit ($): ', np.mean(episode_reward))
    print('Mean Episode Served Demand- Proit($): ', np.mean(episode_served_demand))
    print('Mean Episode Rebalancing Cost($): ', np.mean(episode_rebalancing_cost))

    inflows = np.mean(inflows, axis=0)

    mean_reward = np.mean(episode_reward)
    mean_served_demand = np.mean(episode_served_demand)
    mean_rebalancing_cost = np.mean(episode_rebalancing_cost)

    mean_reward = round(mean_reward/1000,2)
    mean_served_demand = round(mean_served_demand/1000,2)
    mean_rebalancing_cost = round(mean_rebalancing_cost/1000,2)
    rl_means = (mean_reward, mean_served_demand, mean_rebalancing_cost)

    return rl_means, inflows
    

def get_no_control_performance(cfg, env, parser, device, use_saved_data=False):
    #check if no_control performance is saved
    path = f'./src/envs/data/{cfg.simulator.name}/{cfg.simulator.city}_no_control_performance.json'
    #check if path exists
    if os.path.exists(path) and use_saved_data:
        with open(path, 'r') as f:
            no_control_performance = json.load(f)
        no_reb_reward = no_control_performance['reward']
        no_reb_demand = no_control_performance['served_demand']
        no_reb_cost = no_control_performance['rebalancing_cost']
    else:
        print('No control performance not found. Calculating (this happens only the first time on a new environment)...')
        cfg_copy = cfg.copy()
        cfg_copy.model.name = 'no_rebalancing'
        model = setup_model(cfg_copy, env, parser, device)
        no_reb_reward, no_reb_demand, no_reb_cost, _ = model.test(10, env)
        no_reb_reward = round(np.mean(no_reb_reward)/1000,2)
        no_reb_demand = round(np.mean(no_reb_demand)/1000,2)
        no_reb_cost = round(np.mean(no_reb_cost)/1000,2)
        no_control_performance = {'reward': no_reb_reward, 'served_demand': no_reb_demand, 'rebalancing_cost': no_reb_cost}
        print(f'No control performance calculated. Saving in {path}...')
        print(os.getcwd())
        with open(path, 'w') as f:
            json.dump(no_control_performance, f)
    
    return (no_reb_reward, no_reb_demand, no_reb_cost)
    

def plot_comparison(cfg, env, control_data, comparison_data):

    mpl.rcParams.update({
        "font.size": 15,            # base
        "axes.titlesize": 15,
        "axes.labelsize": 15,
        "legend.fontsize": 14,
        "xtick.labelsize": 15,
        "ytick.labelsize": 15,
        "axes.linewidth": 1.0,
        "pdf.fonttype": 42,        
        "ps.fonttype": 42,
    })

    # Colorblind-safe Okabe–Ito palette 
    # okabe_ito = ["#0072B2", "#E69F00", "#009E73", "#D55E00",
    #              "#CC79A7", "#56B4E9", "#000000", "#F0E442"]
    

    # Distinct hatch patterns so bars remain distinguishable in grayscale
    hatches = ["////", "xxxx", "----", "\\\\\\\\","++++", "....", "oooo", "****"]

    # Function to add value labels on top of bars
    def add_value_labels( rects, fs=14):
        for rect in rects:
            height = rect.get_height()
            ax1.annotate(f'{height:.1f}',
                        xy=(rect.get_x() + rect.get_width()/ 2, height),
                        xytext=(0, 3),  # 3 points vertical offset
                        textcoords="offset points",
                        ha='center', va='bottom',fontsize=fs)

    profit_data, inflows = comparison_data
    labels = ['Overall Profit', 'Served Demand Profit', 'Rebalancing Cost']
    x = np.arange(len(labels))  # the label locations
    width = 0.18  # the width of the bars
    num_bars = len(profit_data) + 1

    fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(20, 8))
    
    #fig, ax = plt.subplots(figsize=(8, 5))

    # colors = okabe_ito[:(len(profit_data) + 1)]
    colors = plt.get_cmap("Set2").colors

    # make the first subplot square for IEEE look
    try:
        ax1.set_box_aspect(1)
    except Exception:
        pass

    start_x = x - (num_bars - 1) * width / 2
    for ind, (key, data) in enumerate(profit_data.items()):
        if key == "equal_distribution":
            key = "Equal Distribution"
        elif key == "sac":
            key = "SAC"
        elif key == "random":
            key = "Random"
        rects1 = ax1.bar(
            start_x + ind * width, data, width,
            label=key,
            color=colors[ind % len(colors)],
            edgecolor="black", linewidth=0.9,
            hatch=hatches[ind % len(hatches)]
        )
        add_value_labels(rects1)  # Adding value labels to each bar

    rects2 = ax1.bar(
        start_x + (ind + 1) * width, control_data, width,
        label='No Control',
        color=colors[-1],
        edgecolor="black", linewidth=0.9,
        hatch=hatches[(ind + 1) % len(hatches)]
    )
    add_value_labels(rects2)

    # Add some text for labels, title and custom x-axis tick labels, etc.
    ax1.set_xlabel('Metrics')
    ax1.set_ylabel(r'USD ($\times 10^3$)')
    if cfg.simulator.firm_count == 1:
        # ax1.set_title(f'Comparison on {cfg.simulator.city} Environment with 1 Firm')
        ax1.set_title(f'Comparison on NYC - Brooklyn Environment with 1 Firm')
    else:
        # ax1.set_title(f'Comparison on {cfg.simulator.city} Environment with {cfg.simulator.firm_count} Firms')
        ax1.set_title(f'Comparison on NYC - Brooklyn Environment with {cfg.simulator.firm_count} Firms')
    ax1.set_xticks(x)
    ax1.set_xticklabels(labels)
    ax1.legend(loc='best', frameon=True, edgecolor='black')

    ax1.grid(True, axis='y', linestyle='--', linewidth=0.7, alpha=0.8)

    if cfg.simulator.city != 'nyc_brooklyn': 
        plt.show()
    else: 
        #plots for tutorial
        open_reqest = {0: 0,
            1: 414.0,
            2: 0,
            3: 0,
            4: 0,
            5: 49756.49999999998,
            6: 9948.600000000006,
            7: 98.99999999999999,
            8: 198.00000000000003,
            9: 881.9999999999998,
            10: 1232.9999999999993,
            11: 6492.600000000001,
            12: 23293.80000000004,
            13: 170.99999999999997}
        
        #open_reqest = {k: v / max(open_reqest.values()) for k,v in open_reqest.items()}

        #inflows = inflows / max(inflows)

        labels = range(14)
        x = np.arange(len(labels))  # the label locations
        width = 0.25  # the width of the bars

        r1 = np.arange(14)
        r2 = [x + width for x in r1]

        #fig, ax = plt.subplots(figsize=(8, 5))
        for key, data in inflows.items():
            ax2.bar(r2, data, width, label=f'Rebalancing Flows for {key}', color="#0072BD")
        ax3 = ax2.twinx()  # Create a second y-axis
        ax3.bar(r1, open_reqest.values(), width, label='Profit', color="#A2142F")

        # Add labels and title to the second plot
        ax2.set_xlabel('Regions')
        ax2.set_ylabel('Flows', color="#0072BD")
        ax3.set_ylabel('Profit', color="#A2142F")
        ax2.set_title('Comparison of Incoming Rebalancing Flows vs Profit')
        ax2.set_xticks(r1)
        ax2.set_xticklabels(labels)
        ax2.tick_params(axis='y', labelcolor="#0072BD")
        ax3.tick_params(axis='y', labelcolor="#A2142F")
        #ax2.legend()
        #ax3.legend()

        plt.tight_layout()  
        plt.show()


def test(config):
    '''
    for Colab tutorial
    '''
    with initialize(config_path="src/config"):
        cfg = compose(config_name="config", overrides= [f"{key}={value}" for key, value in config.items()])  # Load the configuration
        
    # Import simulator module based on the configuration
    simulator_name = cfg.simulator.name
    if simulator_name == "sumo":
        env, parser = setup_sumo(cfg)
    elif simulator_name == "macro":
        env, parser = setup_macro(cfg)
    elif simulator_name == "multi_macro":
        env, parser = setup_multi_macro(cfg)
    else:
        raise ValueError(f"Unknown simulator: {simulator_name}")
    
    use_cuda = not cfg.model.no_cuda and torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    
    model = setup_model(cfg, env, parser, device)
    
    print(f'Testing model {cfg.model.name} on {cfg.simulator.name} environment')
    episode_reward, episode_served_demand, episode_rebalancing_cost, inflows = model.test(cfg.model.test_episodes, env)
    # save_vehicle_distribution(env.acc, timestep=0, sim_time=time.strftime("%Y%m%d-%H%M%S"))

    print('Mean Episode Profit ($): ', np.mean(episode_reward))
    print('Mean Episode Served Demand- Proit($): ', np.mean(episode_served_demand))
    print('Mean Episode Rebalancing Cost($): ', np.mean(episode_rebalancing_cost))

    inflows = np.mean(inflows, axis=0)
    
    #check if no_control performance is saved
    path = f'./src/envs/data/{cfg.simulator.name}/{cfg.simulator.city}_no_control_performance.json'
    #check if path exists
    if os.path.exists(path):
        with open(path, 'r') as f:
            no_control_performance = json.load(f)
        no_reb_reward = no_control_performance['reward']
        no_reb_demand = no_control_performance['served_demand']
        no_reb_cost = no_control_performance['rebalancing_cost']
    else:
        print('No control performance not found. Calculating (this happens only the first time on a new environment)...')
        cfg_copy = cfg.copy()
        cfg_copy.model.name = 'no_rebalancing'
        model = setup_model(cfg_copy, env, parser, device)
        no_reb_reward, no_reb_demand, no_reb_cost, _ = model.test(10, env)
        no_reb_reward = round(np.mean(no_reb_reward)/1000,2)
        no_reb_demand = round(np.mean(no_reb_demand)/1000,2)
        no_reb_cost = round(np.mean(no_reb_cost)/1000,2)
        no_control_performance = {'reward': no_reb_reward, 'served_demand': no_reb_demand, 'rebalancing_cost': no_reb_cost}
        print(f'No control performance calculated. Saving in {path}...')
        print(os.getcwd())
        with open(path, 'w') as f:
            json.dump(no_control_performance, f)

    mean_reward = np.mean(episode_reward)
    mean_served_demand = np.mean(episode_served_demand)
    mean_rebalancing_cost = np.mean(episode_rebalancing_cost)

    mean_reward = round(mean_reward/1000,2)
    mean_served_demand = round(mean_served_demand/1000,2)
    mean_rebalancing_cost = round(mean_rebalancing_cost/1000,2)
    labels = ['Overall Profit', 'Served Demand Profit', 'Rebalancing Cost']
    rl_means = [mean_reward, mean_served_demand, mean_rebalancing_cost]
    
    no_control = [no_reb_reward, no_reb_demand, no_reb_cost]
    
    import matplotlib.pyplot as plt
    x = np.arange(len(labels))  # the label locations
    width = 0.15  # the width of the bars

    fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(15, 5))
    
    #fig, ax = plt.subplots(figsize=(8, 5))
    rects1 = ax1.bar(x - width/2, rl_means, width, label=cfg.model.name, color="#0072BD")
    rects2 = ax1.bar(x + width/2, no_control, width, label='No Control', color="#A2142F")

    # Add some text for labels, title and custom x-axis tick labels, etc.
    ax1.set_xlabel('Metrics')
    ax1.set_ylabel('$, x10^3')
    ax1.set_title(f'Comparison of {cfg.model.name} vs No Control')
    ax1.set_xticks(x)
    ax1.set_xticklabels(labels)
    ax1.legend()
    
    # Function to add value labels on top of bars
    def add_value_labels(rects):
        for rect in rects:
            height = rect.get_height()
            ax1.annotate(f'{height:.1f}',
                        xy=(rect.get_x() + rect.get_width() / 2, height),
                        xytext=(0, 3),  # 3 points vertical offset
                        textcoords="offset points",
                        ha='center', va='bottom')

    # Adding value labels to each bar
    add_value_labels(rects1)
    add_value_labels(rects2)

    #plt.tight_layout()
    plt.grid(True, axis='y', linestyle='--', alpha=0.7)

    if cfg.simulator.city != 'nyc_brooklyn': 
        plt.show()
    else: 
        #plots for tutorial
        open_reqest = {0: 0,
            1: 414.0,
            2: 0,
            3: 0,
            4: 0,
            5: 49756.49999999998,
            6: 9948.600000000006,
            7: 98.99999999999999,
            8: 198.00000000000003,
            9: 881.9999999999998,
            10: 1232.9999999999993,
            11: 6492.600000000001,
            12: 23293.80000000004,
            13: 170.99999999999997}
        
        #open_reqest = {k: v / max(open_reqest.values()) for k,v in open_reqest.items()}

        #inflows = inflows / max(inflows)

        labels = range(14)
        x = np.arange(len(labels))  # the label locations
        width = 0.25  # the width of the bars

        r1 = np.arange(14)
        r2 = [x + width for x in r1]

        #fig, ax = plt.subplots(figsize=(8, 5))
        ax2.bar(r2, inflows, width, label='Rebalancing Flows', color="#0072BD")
        ax3 = ax2.twinx()  # Create a second y-axis
        ax3.bar(r1, open_reqest.values(), width, label='Profit', color="#A2142F")

        # Add labels and title to the second plot
        ax2.set_xlabel('Regions')
        ax2.set_ylabel('Flows', color="#0072BD")
        ax3.set_ylabel('Profit', color="#A2142F")
        ax2.set_title('Comparison of Incoming Rebalancing Flows vs Profit')
        ax2.set_xticks(r1)
        ax2.set_xticklabels(labels)
        ax2.tick_params(axis='y', labelcolor="#0072BD")
        ax3.tick_params(axis='y', labelcolor="#A2142F")
        #ax2.legend()
        #ax3.legend()

        plt.tight_layout()  
        plt.show()

@hydra.main(version_base=None, config_path="src/config/", config_name="config")
def main(cfg: DictConfig):
   
    # Import simulator module based on the configuration
    simulator_name = cfg.simulator.name
    if simulator_name == "sumo":
        env, parser = setup_sumo(cfg)

    elif simulator_name == "macro":
        env, parser = setup_macro(cfg)

    elif simulator_name == "multi_macro":
        env, parser = setup_multi_macro(cfg)
    else:
        raise ValueError(f"Unknown simulator: {simulator_name}")

    use_cuda = not cfg.model.no_cuda and torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")

    model = setup_model(cfg, env, parser, device)
    
    print('Testing...')
    episode_reward, episode_served_demand, episode_rebalancing_cost, episode_inflows = model.test(cfg.model.test_episodes, env)

    print('Mean Episode Profit ($): ', np.mean(episode_reward), 'Std Episode Reward: ', np.std(episode_reward))
    print('Mean Episode Served Demand($): ', np.mean(episode_served_demand), 'Std Episode Served Demand: ', np.std(episode_served_demand))
    print('Mean Episode Rebalancing Cost($): ', np.mean(episode_rebalancing_cost), 'Std Episode Rebalancing Cost: ', np.std(episode_rebalancing_cost))

    ##TODO: ADD VISUALIZATION

if __name__ == "__main__":
    main()