Clean viz man!

Description

The below description is supplied in free-text by the user
#!/usr/bin/env python3

import argparse
import inspect
import json
from pathlib import Path
from datetime import datetime, timedelta
import sys

from novem import Plot
import pandas as pd
import requests


def fetch_electricity_prices(date, region="NO1"):
    """
    Fetch electricity prices for a specific date in a Norwegian region.

    Args:
        date (datetime): Date to fetch prices for
        region (str): Norwegian electricity region code (default: 'NO1')

    Returns:
        list: Formatted list of tuples with hour and price including MVA
    """
    # Format date in required format (YYYY/MM-DD)
    formatted_date = date.strftime("%Y/%m-%d")

    # Construct the API URL
    url = (
        f"https://www.hvakosterstrommen.no/api/v1/prices/{formatted_date}_{region}.json"
    )

    try:
        # Fetch data from API
        response = requests.get(url, verify=False)
        response.raise_for_status()  # Raise exception for HTTP errors
        return response.json()

    except requests.exceptions.RequestException as e:
        print(f"Error fetching data for {formatted_date}: {e}", file=sys.stderr)
        return []


def process_data(df: pd.DataFrame):
    # Compute price incl. MVA tax rate in Norway (25%)
    mva_rate = 0.25
    df["NOK_per_kWh"] = (df["NOK_per_kWh"] * (1 + mva_rate)).round(4)

    # Extract date and hour from time_start
    df["date"] = pd.to_datetime(df["time_start"]).dt.date
    df["hour"] = pd.to_datetime(df["time_start"]).dt.hour

    return df.pivot_table(
        index="date", columns="hour", values="NOK_per_kWh", aggfunc="first"
    )


def parse_date(date_str):
    if date_str.lower() == "today":
        date = datetime.now()
    elif date_str.lower() == "tomorrow":
        date = datetime.now() + timedelta(days=1)
    elif date_str.lower() == "yesterday":
        date = datetime.now() - timedelta(days=1)
    else:
        date = datetime.strptime(date_str, "%Y-%m-%d")

    return date.replace(hour=0, minute=0, second=0, microsecond=0)


def main(argv: list[str]):
    parser = argparse.ArgumentParser(
        description="Fetch and display Norwegian electricity prices for specified dates."
    )
    parser.add_argument(
        "--fetch",
        action="store_true",
        help="Fetch and store API JSON data to files only.",
    )
    parser.add_argument(
        "--load",
        action="store_true",
        help="Load data from local JSON files instead of fetching from API.",
    )
    parser.add_argument(
        "--region",
        type=str,
        default="NO1",
        help="Norwegian electricity region code (default: NO1).",
    )
    parser.add_argument(
        "--plot",
        type=str,
        help="Name of novem plot to update.",
    )
    parser.add_argument(
        "dates",
        nargs="*",
        help="Dates to fetch prices for (format: YYYY-MM-DD). Use 'today', 'tomorrow', or 'yesterday'. Default is today if no dates are provided.",
    )
    args = parser.parse_args(argv)

    # Parse input dates
    dates = [parse_date(date) for date in args.dates]
    if not dates:
        dates = [parse_date("today")]
 
    if args.load:
        data = []
        for date in dates:
            filename = Path(f"strompriser_{date.strftime('%Y-%m-%d')}_{args.region}.json")
            if filename.exists():
                json_data = json.loads(filename.read_text())
                data.append(json_data)
            else:
                print(f"File {filename} not found. Skipping.", file=sys.stderr)
    else:
        data = [fetch_electricity_prices(date, args.region) for date in dates]

    if args.fetch:
        if args.load:
            print("Cannot use --fetch and --load together.", file=sys.stderr)
            return
        for date, json_data in zip(dates, data):
            filename = Path(f"strompriser_{date.strftime('%Y-%m-%d')}_{args.region}.json")
            filename.write_text(json.dumps(json_data, indent=2))
            print(f"Saved {filename}")
        return

    # Fetch prices for each date
    prices = pd.concat(process_data(pd.DataFrame(d)) for d in data).transpose()

    # Print csv results
    if args.plot:
        plt = Plot(args.plot)
        plt.name = "Strømpriser i Norge"
        plt.data = prices
        plt.description = '\n'.join([
            '```python',
            inspect.getsource(sys.modules[__name__]),
            '```',
        ])
    else:
        print(prices.to_csv(index_label="date"))


if __name__ == "__main__":
    main(sys.argv[1:])