result_processor.py

import pandas as pd
from pyomo.core.util import quicksum
from pyomo.environ import value
from pathlib import Path
from config import Config  # NEW: Import Config class for constants

def process_model_results(model):
    """
    Process the results from a solved model.
    
    Parameters:
    model (ConcreteModel): A solved Pyomo model
    
    Returns:
    dict: Processed results including generation, revenue, and capacity data
    """
    try:
        gen = calculate_total_generation(model)
        # net_rev = calculate_net_revenue(model)
        # total_cap = calculate_total_capacity(model)
        retirement_schedule = calculate_retirement_schedule(model)
        plant_cap = calculate_plant_capacity(model)
        # net_rev = calculate_net_revenue(model)
        annual_summary = calculate_annual_summary(model) 
        plant_netrev = calculate_plant_netrev(model)
        # NEW: per-technology aggregations
        tech_gen = calculate_generation_by_technology(model)
        tech_cap = calculate_capacity_by_technology(model)
        tech_netrev = calculate_net_revenue_by_technology(model)
        tech_targets = get_technology_targets(model)
        min_cap_capped = getattr(model, "_min_cap_capped", [])
        gen_goal_capped = getattr(model, "_gen_goal_capped", [])
        out = {
            "PlantGen": gen,
            "retire_sched": retirement_schedule,
            "plant_cap": plant_cap,
            "AnnualSummary": annual_summary,
            "plant_netrev": plant_netrev,
            "TechGen": tech_gen,
            "TechCap": tech_cap,
            "TechNetRev": tech_netrev,
            "TechTargets": tech_targets,
        }
        if min_cap_capped:
            out["MinCapacityCapped"] = min_cap_capped
        if gen_goal_capped:
            out["GenGoalCapped"] = gen_goal_capped
        return out
    except Exception as e:
        raise RuntimeError(f"Error processing model results: {str(e)}")

def save_results_to_excel(results, output_file, output_dir=None):
    """
    Save the processed results to an Excel file.
    
    Parameters:
    results (dict): The processed results to save
    output_file (str): The path to the output Excel file
    """
    try:
        # Save each scenario's data to separate files
        for key, result in results.items():
            # NEW: Using Config constant for scenario output file template
            scenario_output_file = Config.SCENARIO_OUTPUT_TEMPLATE.format(key=key)
            if output_dir:
                scenario_output_file = str(Path(output_dir) / scenario_output_file)
            with pd.ExcelWriter(scenario_output_file) as scenario_writer:
           
                # Save each component to a separate sheet
                for sheet_name, data in result.items():
                    if sheet_name == "annual_summary":
                        summary_df = pd.DataFrame.from_dict(data, orient='index')
                        summary_df.to_excel(scenario_writer, sheet_name=f"Summary")
                    elif sheet_name == "plant_netrev":
                        # Save annual net revenue
                        annual_df = pd.DataFrame.from_dict(
                            data["annual"], 
                            orient='index'
                        )
                        # Save depreciated capex as an additional column
                        annual_df['Depreciated_Capex_$m'] = data["depreciated_capex"]
                        annual_df.to_excel(scenario_writer, sheet_name=f"PlantNetRev")
                    elif sheet_name in ("TechGen", "TechCap", "TechNetRev", "TechTargets"):
                        df = pd.DataFrame(data).transpose()
                        df.to_excel(scenario_writer, sheet_name=f"{sheet_name}")
                    elif sheet_name == "MinCapacityCapped":
                        df = pd.DataFrame(data)
                        df.to_excel(scenario_writer, sheet_name="MinCapacityUnmet")
                    elif sheet_name == "GenGoalCapped":
                        df = pd.DataFrame(data)
                        df.to_excel(scenario_writer, sheet_name="GenGoalCapped")
                    else:
                        df = pd.DataFrame.from_dict(data, orient='index')
                        df.to_excel(scenario_writer, sheet_name=f"{sheet_name}")
            
            # NEW: Also create one Excel per technology with its own sheets
            if "TechGen" in result:
                tech_keys = list(result["TechGen"].keys())
                for tech in tech_keys:
                    tech_file = f"{key}_{tech}_results.xlsx"
                    if output_dir:
                        tech_file = str(Path(output_dir) / tech_file)
                    with pd.ExcelWriter(tech_file) as tech_writer:
                        # Generation by year for this technology
                        gen_series = pd.Series(result["TechGen"][tech])
                        gen_series.to_excel(tech_writer, sheet_name="GenByYear")
                        # Capacity by year for this technology
                        if "TechCap" in result and tech in result["TechCap"]:
                            cap_series = pd.Series(result["TechCap"][tech])
                            cap_series.to_excel(tech_writer, sheet_name="CapByYear")
                        # Net revenue by year for this technology
                        if "TechNetRev" in result and tech in result["TechNetRev"]:
                            nr_series = pd.Series(result["TechNetRev"][tech])
                            nr_series.to_excel(tech_writer, sheet_name="NetRevByYear")
                        # Targets by year for this technology
                        if "TechTargets" in result and tech in result["TechTargets"]:
                            tgt_series = pd.Series(result["TechTargets"][tech])
                            tgt_series.to_excel(tech_writer, sheet_name="TargetsByYear")
    except Exception as e:
        raise IOError(f"Error saving results to Excel: {str(e)}")
    try:
        with pd.ExcelWriter(output_file) as writer:
            for key, result in results.items():               
                # Save each component to a separate sheet
                for sheet_name, data in result.items():
                    if sheet_name == "annual_summary":
                        summary_df = pd.DataFrame.from_dict(data, orient='index')
                        summary_df.to_excel(writer, sheet_name=f"{key}_Summary")
                    elif sheet_name == "plant_netrev":
                        # Save annual net revenue
                        annual_df = pd.DataFrame.from_dict(
                            data["annual"], 
                            orient='index'
                        )
                         # Save depreciated capex as an additional column
                        annual_df['Depreciated_Capex_$m'] = data["depreciated_capex"]
                        annual_df.to_excel(writer, sheet_name=f"{key}_PlantNetRev")
                        
                        # # Save depreciated capex
                        # capex_df = pd.DataFrame.from_dict(
                        #     data["depreciated_capex"], 
                        #     orient='index',
                        #     columns=['Depreciated_Capex_$m']
                        # )
                        # capex_df.to_excel(writer, sheet_name=f"{key}_Plant")
                    elif sheet_name in ("TechGen", "TechCap", "TechNetRev", "TechTargets"):
                        df = pd.DataFrame(data).transpose()
                        df.to_excel(writer, sheet_name=f"{key}_{sheet_name}")
                    elif sheet_name == "MinCapacityCapped":
                        df = pd.DataFrame(data)
                        df.to_excel(writer, sheet_name=f"{key}_MinCapacityUnmet")
                    elif sheet_name == "GenGoalCapped":
                        df = pd.DataFrame(data)
                        df.to_excel(writer, sheet_name=f"{key}_GenGoalCapped")
                    else:
                        df = pd.DataFrame.from_dict(data, orient='index')
                        df.to_excel(writer, sheet_name=f"{key}_{sheet_name}")
                    # Save each scenario's data to separate files
    except Exception as e:
        raise IOError(f"Error saving results to Excel: {str(e)}")

# def calculate_net_revenue(model):
def calculate_annual_summary(model):
    """
    Calculate annual summary statistics including total generation and net revenue.
    Based on original GAMS code Summary calculations.
    
    Parameters:
    model (ConcreteModel): The solved Pyomo model
    
    Returns:
    dict: Dictionary containing summary statistics by year
    """
    try:
        summary = {}
        plant_netrev = calculate_net_revenue(model)

        for y in model.y:
            summary[y] = {
                # Total Coal Generation (TWh)
                # NEW: Using Config constants for time conversion (was hardcoded 8.76/1000)
                "Total Coal Gen TWh": sum(
                    model.Gen[g, y, t].value * model.Price_dur[t] * Config.HOURS_PER_YEAR / Config.USD_TO_MILLIONS
                    for g in model.g 
                    for t in model.t
                ),
                
                # Total Capacity (GW)
                # NEW: Using Config constant for MW to GW conversion (was hardcoded /1000)
                "Total Capacity GW": sum(
                    model.Cap[g, y].value for g in model.g
                ) / Config.MW_TO_GW,
                
                # Total Undiscounted Net Revenue ($b)
                # NEW: Using Config constant for conversion to billions (was hardcoded /1000)
                "Total Undiscounted Net Revenue $b":
                   sum(
                    plant_netrev[g][y] for g in model.g
                ) / Config.USD_TO_THOUSANDS,  # Convert to billions
                "Discounted Net Revenue $b": sum(
                    plant_netrev[g][y] * model.DR[y] for g in model.g
                ) / Config.USD_TO_THOUSANDS  # Convert to billions
                }
        
        return summary
    except Exception as e:
        raise RuntimeError(f"Error calculating annual summary: {str(e)}")

def calculate_plant_netrev(model):
    """
    Calculate net revenue for each plant and year.
    Based on original GAMS code NetRev calculations.
    
    Parameters:
    model (ConcreteModel): The solved Pyomo model
    
    Returns:
    dict: Dictionary containing plant net revenue and depreciated capex
    """
    try:
        plant_netrev = {"annual": {}, "depreciated_capex": {}}
        
        # Calculate annual net revenue for each plant
        for g in model.g:
            plant_netrev["annual"][g] = {}
            for y in model.y:
                # Calculate net revenue according to GAMS formula
                '''
                if y == min(model.y):  # Use dynamic base year
                    print("--------------------------------")
                    print('BAU',model.GenData[g]["CAPACITY"])
                    print('AD',model.Cap[g, y].value)
                    print(model.GenData[g]["CAPACITY"] * model.SetScenario[model.s[1]] +
                        model.Cap[g, y].value * (1 - model.SetScenario[model.s[1]]))
                    print('Gen',sum(model.Gen[g, y, t].value * model.Price_dur[t] * 8.76 for t in model.t))
                    print('PPA',model.FC_PPA[g, y])
                    print('SetScenario', [model.SetScenario[s] for s in model.s])
                    print('SetPriceScenario', [model.SetPriceScenario[p] for p in model.p])
                    print("--------------------------------") 
                    print
                '''
                netrev = -(
                            # Use the corresponding capacity based on scenario
                            model.Cap[g, y].value if model.s.at(1) == "AD" else model.GenData[g]["CAPACITY"]
                        ) *(
                            # NEW: Using Config constants for price scenario handling (was hardcoded /1e3 and 100)
                            model.FC_PPA[g, y]/Config.USD_TO_THOUSANDS if model.p.at(1) == "AvgPPAPrice" else Config.DEFAULT_COST_PER_MW_MarketPrice
                        ) + quicksum(
                        (
                            model.rev_unit[g, y, model.p.at(1)] * 
                            model.Price_Dist1[y, model.p.at(1), t] - 
                            model.cost[g, y]
                        ) * model.Gen[g, y, t].value * model.Price_dur[t] * Config.HOURS_PER_YEAR
                        for t in model.t
                    )
               
                # NEW: Using Config constant for conversion to millions (was hardcoded /1e6)
                plant_netrev["annual"][g][y] = netrev/Config.USD_TO_MILLIONS
            
            # Calculate depreciated capex
            # NEW: Using Config constant for conversion to thousands (was hardcoded /1000)
            # Get technology type for this plant
            tech_type = value(model.GenData[g]["TECHNOLOGY"])
            plant_netrev["depreciated_capex"][g] = max(
                model.GenData[g]["CAPACITY"] * 
                model.TechParams[tech_type, "CoalCapex $/kW"] *
                (1 - model.TechParams[tech_type, "Straight-line depreciation"] * model.life[g]), 
                0
            ) / Config.USD_TO_THOUSANDS
        return plant_netrev
    except Exception as e:
        raise RuntimeError(f"Error calculating plant net revenue: {str(e)}")

def calculate_annual_total_netrev(model):
    """
    Calculate total net revenue for each year (both nominal and discounted).
    
    Parameters:
    model (ConcreteModel): The solved Pyomo model
    
    Returns:
    dict: Dictionary containing nominal and discounted net revenue by year
    """
    try:
        annual_netrev = {"nominal": {}, "discounted": {}}
        
        # Use the existing calculate_net_revenue function to get the net revenue for each plant
        plant_netrev = calculate_net_revenue(model)
        
        # Get all years
        years = list(plant_netrev[list(plant_netrev.keys())[0]].keys())
        
        # Calculate total net revenue for each year
        for year in years:
            # Calculate nominal value: sum of net revenue for all plants in the year
            annual_netrev["nominal"][year] = sum(
                plant_netrev[plant][year] 
                for plant in plant_netrev.keys()
            )
            
            # NEW: Using Config constants for discount rate and base year (was hardcoded 0.06 and 2021)
            annual_netrev["discounted"][year] = (
                annual_netrev["nominal"][year] / 
                (1 + Config.DISCOUNT_RATE) ** (int(year) - Config.BASE_YEAR)
            )

        return annual_netrev
    except Exception as e:
        raise RuntimeError(f"Error calculating annual total net revenue: {str(e)}")

def calculate_total_generation(model):
    """
    Calculate total generation for each plant and year.
    
    Parameters:
    model (ConcreteModel): The solved Pyomo model
    
    Returns:
    dict: Total generation by plant and year in TWh
    """
    try:
        gen = {}
        '''
        for g in model.g:
            total_gen[g] = {}
            for y in model.y:
                # Sum generation across all time blocks and convert to TWh
                total_gen[g][y] = sum(
                    model.Gen[g, y, t].value * model.Price_dur[t] * 8.76 / 1000 
                    for t in model.t
                )
        '''
        for g in model.g:
            gen[g] = {}
            for y in model.y:
            #         for t in model.t:
            #             print(model.Gen[g, y, t])
                # NEW: Using Config constant for hours per year (was hardcoded 8.76)
                gen[g][y] = round(sum(
                    model.Gen[g, y, t].value * model.Price_dur[t] * Config.HOURS_PER_YEAR/Config.USD_TO_THOUSANDS for t in model.t),5)
        
        return gen
    except Exception as e:
        raise RuntimeError(f"Error calculating total generation: {str(e)}")

def calculate_net_revenue(model):
    """
    Calculate net revenue for each plant and year.
    
    Parameters:
    model (ConcreteModel): The solved Pyomo model
    
    Returns:
    dict: Net revenue by plant and year in millions of currency units
    """
    try:
        net_revenue = calculate_plant_netrev(model)["annual"]
        return net_revenue
    except Exception as e:
        raise RuntimeError(f"Error calculating net revenue: {str(e)}")

def calculate_total_capacity(model):
    """
    Calculate total capacity for each plant and year.
    
    Parameters:
    model (ConcreteModel): The solved Pyomo model
    
    Returns:
    dict: Total capacity by plant and year in GW
    """
    try:
        total_capacity = {}
        for y in model.y:    
            total_capacity_mw = sum(
                model.Cap[g, y].value for g in model.g
            )
            # print(f"Total Capacity: {total_capacity_mw}")
            # NEW: Using Config constant for MW to GW conversion (was hardcoded /1000)
            total_capacity[y] = total_capacity_mw/Config.MW_TO_GW
        # print(total_capacity)
        return total_capacity
    except Exception as e:
        raise RuntimeError(f"Error calculating total capacity: {str(e)}")
    
def calculate_plant_capacity(model):
    """
    Calculate capacity for each plant and year.
    
    Parameters:
    model (ConcreteModel): The solved Pyomo model
    
    Returns:
    dict: Capacity by plant and year in MW
    """
    try:
        plant_capacity = {}
        for g in model.g:
            plant_capacity[g] = {}
            for y in model.y:
                plant_capacity[g][y] = model.Cap[g, y].value
        return plant_capacity
    except Exception as e:
        raise RuntimeError(f"Error calculating plant capacity: {str(e)}")


def calculate_retirement_schedule(model):
    """
    Calculate retirement schedule for each plant.
    
    Parameters:
    model (ConcreteModel): The solved Pyomo model
    
    Returns:
    dict: Retirement schedule by plant and year (1 if retired, 0 if operating)
    """
    try:
        retirement = {}
        for g in model.g:
            retirement[g] = {}
            for y in model.y:
                retirement[g][y] = model.Retire[g, y].value
        return retirement
    except Exception as e:
        raise RuntimeError(f"Error calculating retirement schedule: {str(e)}")

# NEW: Per-technology aggregations
def calculate_generation_by_technology(model):
    try:
        gen_by_tech = {}
        for tech in model.tech:
            gen_by_tech[tech] = {}
            for y in model.y:
                gen_by_tech[tech][y] = round(sum(
                    sum(model.Gen[g, y, t].value * model.Price_dur[t] * Config.HOURS_PER_YEAR/Config.USD_TO_THOUSANDS for t in model.t)
                    for g in model.plants_by_tech[tech]
                ), 5)
        return gen_by_tech
    except Exception as e:
        raise RuntimeError(f"Error calculating generation by technology: {str(e)}")

def calculate_capacity_by_technology(model):
    try:
        cap_by_tech = {}
        for tech in model.tech:
            cap_by_tech[tech] = {}
            for y in model.y:
                cap_by_tech[tech][y] = sum(model.Cap[g, y].value for g in model.plants_by_tech[tech])
        return cap_by_tech
    except Exception as e:
        raise RuntimeError(f"Error calculating capacity by technology: {str(e)}")

def calculate_net_revenue_by_technology(model):
    try:
        plant_nr = calculate_plant_netrev(model)["annual"]
        nr_by_tech = {}
        for tech in model.tech:
            nr_by_tech[tech] = {}
            for y in model.y:
                nr_by_tech[tech][y] = sum(
                    plant_nr[g][y] for g in model.plants_by_tech[tech]
                )
        return nr_by_tech
    except Exception as e:
        raise RuntimeError(f"Error calculating net revenue by technology: {str(e)}")

def get_technology_targets(model):
    try:
        targets = {}
        for tech in model.tech:
            targets[tech] = {}
            for y in model.y:
                targets[tech][y] = float(model.PriceGenTech[y, tech]) if (y, tech) in model.PriceGenTech else 0.0
        return targets
    except Exception as e:
        raise RuntimeError(f"Error collecting technology targets: {str(e)}")