import numpy as np
import pyomo.environ as pyo
from typing import Dict
from build_consumer_agent import build_consumer_agent


def solve_consumer_agent(model: pyo.AbstractModel, data: Dict, instance_data: Dict = None) -> pyo.ConcreteModel:
    """Solve consumer agent's optimization problem."""
    number_of_timesteps = data["General"]["nTimesteps"]
    # First delete existing model components if they were previously instantiated
    if hasattr(model, 'var_g'):
        model.del_component(model.var_g)
    if hasattr(model, 'objective_function'):
        model.del_component(model.objective_function)
    if hasattr(model, 'energy_balance'):
        model.del_component(model.energy_balance)
    # Default instance data
    default_instance_data = {
        None: {
            'λ_EOM': dict(enumerate(np.zeros(number_of_timesteps))),
            'g_bar': dict(enumerate(np.zeros(number_of_timesteps))),
            'ρ_EOM': {None: 1},
        }
    }
    # Use default instance data if none is provided
    if instance_data is None:
        instance_data = default_instance_data
    # Reconstruct and instantiate model for next round of solving
    model = build_consumer_agent(model)
    model_instance = model.create_instance(data=instance_data)
    solver = pyo.SolverFactory('gurobi')
    solver.solve(model_instance)
    
    return model_instance


# # Example usage
# import yaml
# import pandas as pd
# from define_common_parameters import define_common_parameters
# from define_consumer_parameters import define_consumer_parameters


# def read_config(config_file):
#     with open(config_file, 'r') as file:
#         config = yaml.safe_load(file)
#     return config

# data = read_config('config.yaml')
# consumer_agents = [id for id in data['Consumers'].keys()]
# ts = pd.read_csv('timeseries.csv', delimiter=';')
# mdict = {m: define_common_parameters(m, data) for m in consumer_agents}

# instance_data_1 = {
#     None: {
#         'λ_EOM': dict(enumerate(np.ones(24)*10)),
#         'g_bar': dict(enumerate(np.ones(24)*150)),
#         'ρ_EOM': {None: 1},
#     }
# }

# instance_data_2 = {
#     None: {
#         'λ_EOM': dict(enumerate(np.ones(24)*50)),
#         'g_bar': dict(enumerate(np.ones(24)*300)),
#         'ρ_EOM': {None: 1},
#     }
# }

# for agent, model in mdict.items():
#     model = define_consumer_parameters(agent, model, data, ts)
#     zero_model_instance = solve_consumer_agent(model, data)
#     print(zero_model_instance.name)
#     print()
#     zero_model_instance.objective_function.display()
#     first_model_instance = solve_consumer_agent(model, data, instance_data_1)
#     print(first_model_instance.name)
#     print()
#     first_model_instance.objective_function.display()
#     print()
#     # second_model_instance = solve_consumer_agent(model, data, instance_data_2)
#     # print(second_model_instance.name)
#     # print()
#     # second_model_instance.display()