#!/usr/bin/env python
"""
Example: Curtailment as Percentage of Total Generation

This example demonstrates how to calculate and visualize curtailment 
as a percentage of total energy generated by region per day.

This helps understand the relative impact of curtailment on the grid.
"""

import os
from datetime import datetime, timedelta

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from dotenv import load_dotenv

from openelectricity import OEClient
from openelectricity.types import DataMetric, MarketMetric

# Load environment variables
load_dotenv()

# Set up the seaborn style
sns.set_theme(style="whitegrid")
plt.rcParams["figure.figsize"] = (16, 10)


def fetch_generation_and_curtailment(client: OEClient, days_back: int = 14) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Fetch total generation and curtailment energy data for all NEM regions.
    
    Args:
        client: OpenElectricity API client
        days_back: Number of days to fetch data for
        
    Returns:
        Tuple of (generation_df, curtailment_df)
    """
    # Calculate date range (omit end_date to get latest)
    start_date = datetime.now() - timedelta(days=days_back)
    
    print(f"Fetching data for the last {days_back} days")
    
    # Fetch total energy generation by region
    print("  Fetching total energy generation...")
    generation_response = client.get_network_data(
        network_code="NEM",
        metrics=[DataMetric.ENERGY],
        interval="1d",
        date_start=pd.to_datetime(start_date),
        # date_end omitted to get latest data
        primary_grouping="network_region",
    )
    
    # Fetch curtailment energy by region
    print("  Fetching curtailment energy...")
    curtailment_response = client.get_market(
        network_code="NEM",
        metrics=[MarketMetric.CURTAILMENT_ENERGY],
        interval="1d",
        date_start=pd.to_datetime(start_date),
        # date_end omitted to get latest data
        primary_grouping="network_region",
    )
    
    # Process generation data
    gen_data = []
    for timeseries in generation_response.data:
        for result in timeseries.results:
            # Extract region from result name or group_by
            name_parts = result.name.split("_")
            region = name_parts[-1] if len(name_parts) > 1 else "Unknown"
            
            for data_point in result.data:
                timestamp = data_point.timestamp if hasattr(data_point, "timestamp") else data_point.root[0]
                value = data_point.value if hasattr(data_point, "value") else data_point.root[1]
                if value is not None:
                    gen_data.append({
                        "date": pd.to_datetime(timestamp).date(),
                        "region": region,
                        "generation_mwh": value
                    })
    
    # Process curtailment data
    curt_data = []
    for timeseries in curtailment_response.data:
        for result in timeseries.results:
            # Extract region from result name
            name_parts = result.name.split("_")
            region = name_parts[-1] if len(name_parts) > 1 else "Unknown"
            
            for data_point in result.data:
                timestamp = data_point.timestamp if hasattr(data_point, "timestamp") else data_point.root[0]
                value = data_point.value if hasattr(data_point, "value") else data_point.root[1]
                if value is not None:
                    curt_data.append({
                        "date": pd.to_datetime(timestamp).date(),
                        "region": region,
                        "curtailment_mwh": value
                    })
    
    generation_df = pd.DataFrame(gen_data)
    curtailment_df = pd.DataFrame(curt_data)
    
    return generation_df, curtailment_df


def calculate_curtailment_percentage(generation_df: pd.DataFrame, curtailment_df: pd.DataFrame) -> pd.DataFrame:
    """
    Calculate curtailment as a percentage of total generation.
    
    Args:
        generation_df: DataFrame with generation data
        curtailment_df: DataFrame with curtailment data
        
    Returns:
        DataFrame with curtailment percentages
    """
    # Merge generation and curtailment data
    merged_df = pd.merge(
        generation_df,
        curtailment_df,
        on=["date", "region"],
        how="outer"
    )
    
    # Fill NaN values with 0
    merged_df = merged_df.fillna(0)
    
    # Calculate curtailment percentage
    # Note: Curtailment is already included in generation, so the actual generation without curtailment
    # would be generation - curtailment. The percentage is curtailment / generation * 100
    merged_df["curtailment_pct"] = (merged_df["curtailment_mwh"] / merged_df["generation_mwh"]) * 100
    
    # Handle any inf or NaN values
    merged_df["curtailment_pct"] = merged_df["curtailment_pct"].replace([float('inf'), -float('inf')], 0)
    merged_df["curtailment_pct"] = merged_df["curtailment_pct"].fillna(0)
    
    # Sort by date and region
    merged_df = merged_df.sort_values(["date", "region"])
    
    return merged_df


def create_percentage_line_chart(df: pd.DataFrame):
    """
    Create a line chart showing curtailment percentage over time by region.
    
    Args:
        df: DataFrame with curtailment percentage data
    """
    if df.empty:
        print("No data to plot")
        return
    
    # Pivot data for line chart
    pivot_df = df.pivot(index="date", columns="region", values="curtailment_pct")
    
    # Define colors for each region
    region_colors = {
        "NSW1": "#FF6B00",  # Orange
        "QLD1": "#0066CC",  # Blue
        "VIC1": "#228B22",  # Green
        "SA1": "#DC143C",   # Red
        "TAS1": "#8B008B"   # Purple
    }
    
    # Create the plot
    fig, ax = plt.subplots(figsize=(14, 8))
    
    # Plot line for each region
    for region in pivot_df.columns:
        color = region_colors.get(region, "#888888")
        pivot_df[region].plot(
            ax=ax,
            label=region,
            color=color,
            linewidth=2,
            marker='o',
            markersize=4,
            alpha=0.8
        )
    
    # Format the chart
    ax.set_title("Curtailment as Percentage of Total Generation by Region", fontsize=16, fontweight="bold")
    ax.set_xlabel("Date", fontsize=12)
    ax.set_ylabel("Curtailment (%)", fontsize=12)
    ax.legend(title="Region", bbox_to_anchor=(1.05, 1), loc="upper left")
    ax.grid(True, alpha=0.3)
    
    # Format x-axis labels
    ax.tick_params(axis='x', rotation=45)
    
    # Add horizontal line at zero
    ax.axhline(y=0, color="black", linestyle="-", linewidth=0.5, alpha=0.5)
    
    # Add percentage formatting to y-axis
    ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda y, _: f'{y:.1f}%'))
    
    plt.tight_layout()
    
    # Save the figure
    output_file = "curtailment_percentage.png"
    plt.savefig(output_file, dpi=150, bbox_inches="tight")
    print(f"\nChart saved to {output_file}")
    
    # plt.show()  # Uncomment to display the chart interactively


def create_heatmap(df: pd.DataFrame):
    """
    Create a heatmap showing curtailment percentage by date and region.
    
    Args:
        df: DataFrame with curtailment percentage data
    """
    if df.empty:
        print("No data to plot")
        return
    
    # Pivot data for heatmap
    pivot_df = df.pivot(index="region", columns="date", values="curtailment_pct")
    
    # Create the plot
    fig, ax = plt.subplots(figsize=(16, 6))
    
    # Create heatmap
    sns.heatmap(
        pivot_df,
        annot=True,
        fmt=".1f",
        cmap="YlOrRd",
        cbar_kws={'label': 'Curtailment (%)'},
        ax=ax,
        vmin=0,
        vmax=pivot_df.max().max()
    )
    
    # Format the chart
    ax.set_title("Curtailment Percentage Heatmap by Region and Date", fontsize=16, fontweight="bold")
    ax.set_xlabel("Date", fontsize=12)
    ax.set_ylabel("Region", fontsize=12)
    
    # Format date labels
    date_labels = [d.strftime('%m/%d') for d in pivot_df.columns]
    ax.set_xticklabels(date_labels, rotation=45, ha='right')
    
    plt.tight_layout()
    
    # Save the figure
    output_file = "curtailment_percentage_heatmap.png"
    plt.savefig(output_file, dpi=150, bbox_inches="tight")
    print(f"\nHeatmap saved to {output_file}")
    
    # plt.show()  # Uncomment to display the chart interactively


def print_summary_statistics(df: pd.DataFrame):
    """
    Print summary statistics for curtailment percentage data.
    
    Args:
        df: DataFrame with curtailment percentage data
    """
    if df.empty:
        print("No data to summarize")
        return
    
    print("\n" + "=" * 60)
    print("CURTAILMENT PERCENTAGE SUMMARY")
    print("=" * 60)
    
    # Overall statistics by region
    print("\nAverage Curtailment Percentage by Region:")
    print("-" * 40)
    
    regions = sorted(df["region"].unique())
    
    for region in regions:
        region_df = df[df["region"] == region]
        avg_pct = region_df["curtailment_pct"].mean()
        max_pct = region_df["curtailment_pct"].max()
        max_date = region_df.loc[region_df["curtailment_pct"].idxmax(), "date"]
        
        total_gen = region_df["generation_mwh"].sum()
        total_curt = region_df["curtailment_mwh"].sum()
        
        print(f"\n{region}:")
        print(f"  Average: {avg_pct:.2f}%")
        print(f"  Maximum: {max_pct:.2f}% (on {max_date})")
        print(f"  Total Generation: {total_gen:,.0f} MWh")
        print(f"  Total Curtailed: {total_curt:,.0f} MWh")
    
    # NEM-wide statistics
    print("\n" + "-" * 40)
    print("NEM-WIDE STATISTICS:")
    print("-" * 40)
    
    # Calculate NEM-wide daily percentages
    daily_totals = df.groupby("date").agg({
        "generation_mwh": "sum",
        "curtailment_mwh": "sum"
    })
    daily_totals["curtailment_pct"] = (daily_totals["curtailment_mwh"] / daily_totals["generation_mwh"]) * 100
    
    avg_nem_pct = daily_totals["curtailment_pct"].mean()
    max_nem_pct = daily_totals["curtailment_pct"].max()
    max_nem_date = daily_totals["curtailment_pct"].idxmax()
    
    total_nem_gen = df["generation_mwh"].sum()
    total_nem_curt = df["curtailment_mwh"].sum()
    overall_pct = (total_nem_curt / total_nem_gen) * 100
    
    print(f"  Average Daily: {avg_nem_pct:.2f}%")
    print(f"  Maximum Daily: {max_nem_pct:.2f}% (on {max_nem_date})")
    print(f"  Overall Period: {overall_pct:.2f}%")
    print(f"  Total Generation: {total_nem_gen:,.0f} MWh")
    print(f"  Total Curtailed: {total_nem_curt:,.0f} MWh")
    
    # Find days with highest curtailment
    print("\n" + "-" * 40)
    print("TOP 5 CURTAILMENT DAYS (NEM-wide):")
    print("-" * 40)
    
    top_days = daily_totals.nlargest(5, "curtailment_pct")
    for date, row in top_days.iterrows():
        print(f"  {date}: {row['curtailment_pct']:.2f}% ({row['curtailment_mwh']:,.0f} MWh)")


def main():
    """Main function to run the curtailment percentage analysis."""
    
    # Initialize the client
    api_key = os.getenv("OPENELECTRICITY_API_KEY")
    api_url = os.getenv("OPENELECTRICITY_API_URL", "https://api.openelectricity.org.au/v4")
    
    if not api_key:
        print("Error: OPENELECTRICITY_API_KEY environment variable not set")
        print("Please set your API key in the .env file or environment")
        return
    
    client = OEClient(api_key=api_key, base_url=api_url)
    
    print("OpenElectricity Curtailment Percentage Analysis")
    print("=" * 60)
    print()
    
    try:
        # Fetch generation and curtailment data for the last 14 days
        generation_df, curtailment_df = fetch_generation_and_curtailment(client, days_back=14)
        
        if generation_df.empty or curtailment_df.empty:
            print("No data retrieved")
            return
        
        print(f"\nRetrieved {len(generation_df)} generation and {len(curtailment_df)} curtailment data points")
        
        # Calculate curtailment percentages
        percentage_df = calculate_curtailment_percentage(generation_df, curtailment_df)
        
        # Print summary statistics
        print_summary_statistics(percentage_df)
        
        # Create visualizations
        print("\nGenerating charts...")
        create_percentage_line_chart(percentage_df)
        create_heatmap(percentage_df)
        
    except Exception as e:
        print(f"Error during analysis: {e}")
        import traceback
        traceback.print_exc()
    
    print("\n" + "=" * 60)
    print("Analysis complete!")


if __name__ == "__main__":
    main()