data_add.py

"""
CLI commands for populating the database
"""

from __future__ import annotations

from datetime import datetime, timedelta
from typing import Type, List
import isodate
import json
import yaml
from pathlib import Path
from io import TextIOBase

from marshmallow import validate
import pandas as pd
import pytz
from flask import current_app as app
from flask.cli import with_appcontext
import click
import getpass
from sqlalchemy.exc import IntegrityError
from sqlalchemy import func
from timely_beliefs.sensors.func_store.knowledge_horizons import x_days_ago_at_y_oclock
import timely_beliefs as tb
import timely_beliefs.utils as tb_utils
from workalendar.registry import registry as workalendar_registry

from flexmeasures.cli.utils import DeprecatedDefaultGroup, MsgStyle
from flexmeasures.data import db
from flexmeasures.data.scripts.data_gen import (
    add_transmission_zone_asset,
    populate_initial_structure,
    add_default_asset_types,
)
from flexmeasures.data.services.data_sources import get_or_create_source
from flexmeasures.data.services.forecasting import create_forecasting_jobs
from flexmeasures.data.services.scheduling import make_schedule, create_scheduling_job
from flexmeasures.data.services.users import create_user
from flexmeasures.data.models.user import Account, AccountRole, RolesAccounts
from flexmeasures.data.models.time_series import (
    Sensor,
    TimedBelief,
)
from flexmeasures.data.models.validation_utils import (
    check_required_attributes,
    MissingAttributeException,
)
from flexmeasures.data.models.annotations import Annotation, get_or_create_annotation
from flexmeasures.data.schemas import (
    AwareDateTimeField,
    DurationField,
    LatitudeField,
    LongitudeField,
    SensorIdField,
    TimeIntervalField,
)
from flexmeasures.data.schemas.times import TimeIntervalSchema
from flexmeasures.data.schemas.scheduling.storage import EfficiencyField
from flexmeasures.data.schemas.sensors import SensorSchema
from flexmeasures.data.schemas.units import QuantityField
from flexmeasures.data.schemas.generic_assets import (
    GenericAssetSchema,
    GenericAssetTypeSchema,
)
from flexmeasures.data.models.generic_assets import GenericAsset, GenericAssetType
from flexmeasures.data.models.user import User
from flexmeasures.data.services.data_sources import (
    get_source_or_none,
)
from flexmeasures.data.services.utils import get_or_create_model
from flexmeasures.utils import flexmeasures_inflection
from flexmeasures.utils.time_utils import server_now, apply_offset_chain
from flexmeasures.utils.unit_utils import convert_units, ur
from flexmeasures.data.utils import save_to_db
from flexmeasures.data.models.reporting import Reporter
from timely_beliefs import BeliefsDataFrame


@click.group("add")
def fm_add_data():
    """FlexMeasures: Add data."""


@fm_add_data.command("account-role")
@with_appcontext
@click.option("--name", required=True)
@click.option("--description")
def new_account_role(name: str, description: str):
    """
    Create an account role.
    """
    role = AccountRole.query.filter_by(name=name).one_or_none()
    if role is not None:
        click.secho(f"Account role '{name}' already exists.", **MsgStyle.ERROR)
        raise click.Abort()
    role = AccountRole(name=name, description=description)
    db.session.add(role)
    db.session.commit()
    click.secho(
        f"Account role '{name}' (ID: {role.id}) successfully created.",
        **MsgStyle.SUCCESS,
    )


@fm_add_data.command("account")
@with_appcontext
@click.option("--name", required=True)
@click.option("--roles", help="e.g. anonymous,Prosumer,CPO")
def new_account(name: str, roles: str):
    """
    Create an account for a tenant in the FlexMeasures platform.
    """
    account = db.session.query(Account).filter_by(name=name).one_or_none()
    if account is not None:
        click.secho(f"Account '{name}' already exists.", **MsgStyle.ERROR)
        raise click.Abort()
    account = Account(name=name)
    db.session.add(account)
    if roles:
        for role_name in roles.split(","):
            role = AccountRole.query.filter_by(name=role_name).one_or_none()
            if role is None:
                click.secho(f"Adding account role {role_name} ...", **MsgStyle.ERROR)
                role = AccountRole(name=role_name)
                db.session.add(role)
            db.session.flush()
            db.session.add(RolesAccounts(role_id=role.id, account_id=account.id))
    db.session.commit()
    click.secho(
        f"Account '{name}' (ID: {account.id}) successfully created.",
        **MsgStyle.SUCCESS,
    )


@fm_add_data.command("user")
@with_appcontext
@click.option("--username", required=True)
@click.option("--email", required=True)
@click.option(
    "--account-id",
    type=int,
    required=True,
    help="Add user to this account. Follow up with the account's ID.",
)
@click.option("--roles", help="e.g. anonymous,Prosumer,CPO")
@click.option(
    "--timezone",
    "timezone_optional",
    help="Timezone as string, e.g. 'UTC' or 'Europe/Amsterdam' (defaults to FLEXMEASURES_TIMEZONE config setting)",
)
def new_user(
    username: str,
    email: str,
    account_id: int,
    roles: list[str],
    timezone_optional: str | None,
):
    """
    Create a FlexMeasures user.

    The `users create` task from Flask Security Too is too simple for us.
    Use this to add email, timezone and roles.
    """
    if timezone_optional is None:
        timezone = app.config.get("FLEXMEASURES_TIMEZONE", "UTC")
        click.secho(
            f"Setting user timezone to {timezone} (taken from FLEXMEASURES_TIMEZONE config setting)...",
            **MsgStyle.WARN,
        )
    else:
        timezone = timezone_optional
    try:
        pytz.timezone(timezone)
    except pytz.UnknownTimeZoneError:
        click.secho(f"Timezone {timezone} is unknown!", **MsgStyle.ERROR)
        raise click.Abort()
    account = db.session.query(Account).get(account_id)
    if account is None:
        click.secho(f"No account with ID {account_id} found!", **MsgStyle.ERROR)
        raise click.Abort()
    pwd1 = getpass.getpass(prompt="Please enter the password:")
    pwd2 = getpass.getpass(prompt="Please repeat the password:")
    if pwd1 != pwd2:
        click.secho("Passwords do not match!", **MsgStyle.ERROR)
        raise click.Abort()
    created_user = create_user(
        username=username,
        email=email,
        password=pwd1,
        account_name=account.name,
        timezone=timezone,
        user_roles=roles,
        check_email_deliverability=False,
    )
    db.session.commit()
    click.secho(f"Successfully created user {created_user}", **MsgStyle.SUCCESS)


@fm_add_data.command("sensor")
@with_appcontext
@click.option("--name", required=True)
@click.option("--unit", required=True, help="e.g. °C, m/s, kW/m²")
@click.option(
    "--event-resolution",
    required=True,
    type=str,
    help="Expected resolution of the data in ISO8601 duration string",
)
@click.option(
    "--timezone",
    required=True,
    help="Timezone as string, e.g. 'UTC' or 'Europe/Amsterdam'",
)
@click.option(
    "--asset-id",
    "generic_asset_id",
    required=True,
    type=int,
    help="Generic asset to assign this sensor to",
)
@click.option(
    "--attributes",
    required=False,
    type=str,
    default="{}",
    help='Additional attributes. Passed as JSON string, should be a dict. Hint: Currently, for sensors that measure power, use {"capacity_in_mw": 10} to set a capacity of 10 MW',
)
def add_sensor(**args):
    """Add a sensor."""
    check_timezone(args["timezone"])
    try:
        attributes = json.loads(args["attributes"])
    except json.decoder.JSONDecodeError as jde:
        click.secho(
            f"Error decoding --attributes. Please check your JSON: {jde}",
            **MsgStyle.ERROR,
        )
        raise click.Abort()
    del args["attributes"]  # not part of schema
    if args["event_resolution"].isdigit():
        click.secho(
            "DeprecationWarning: Use ISO8601 duration string for event-resolution, minutes in int will be depricated from v0.16.0",
            **MsgStyle.WARN,
        )
        timedelta_event_resolution = timedelta(minutes=int(args["event_resolution"]))
        isodate_event_resolution = isodate.duration_isoformat(
            timedelta_event_resolution
        )
        args["event_resolution"] = isodate_event_resolution
    check_errors(SensorSchema().validate(args))

    sensor = Sensor(**args)
    if not isinstance(attributes, dict):
        click.secho("Attributes should be a dict.", **MsgStyle.ERROR)
        raise click.Abort()
    sensor.attributes = attributes
    if sensor.measures_power:
        if "capacity_in_mw" not in sensor.attributes:
            click.secho(
                "A sensor which measures power needs a capacity (see --attributes).",
                **MsgStyle.ERROR,
            )
            raise click.Abort()
    db.session.add(sensor)
    db.session.commit()
    click.secho(f"Successfully created sensor with ID {sensor.id}", **MsgStyle.SUCCESS)
    click.secho(
        f"You can access it at its entity address {sensor.entity_address}",
        **MsgStyle.SUCCESS,
    )


@fm_add_data.command("asset-type")
@with_appcontext
@click.option("--name", required=True)
@click.option(
    "--description",
    type=str,
    help="Description (useful to explain acronyms, for example).",
)
def add_asset_type(**args):
    """Add an asset type."""
    check_errors(GenericAssetTypeSchema().validate(args))
    generic_asset_type = GenericAssetType(**args)
    db.session.add(generic_asset_type)
    db.session.commit()
    click.secho(
        f"Successfully created asset type with ID {generic_asset_type.id}.",
        **MsgStyle.SUCCESS,
    )
    click.secho("You can now assign assets to it.", **MsgStyle.SUCCESS)


@fm_add_data.command("asset")
@with_appcontext
@click.option("--name", required=True)
@click.option(
    "--latitude",
    type=LatitudeField(),
    help="Latitude of the asset's location",
)
@click.option(
    "--longitude",
    type=LongitudeField(),
    help="Longitude of the asset's location",
)
@click.option(
    "--account-id",
    type=int,
    required=False,
    help="Add asset to this account. Follow up with the account's ID. If not set, the asset will become public (which makes it accessible to all users).",
)
@click.option(
    "--asset-type-id",
    "generic_asset_type_id",
    required=True,
    type=int,
    help="Asset type to assign to this asset",
)
def add_asset(**args):
    """Add an asset."""
    check_errors(GenericAssetSchema().validate(args))
    generic_asset = GenericAsset(**args)
    if generic_asset.account_id is None:
        click.secho(
            "Creating a PUBLIC asset, as no --account-id is given ...",
            **MsgStyle.WARN,
        )
    db.session.add(generic_asset)
    db.session.commit()
    click.secho(
        f"Successfully created asset with ID {generic_asset.id}.", **MsgStyle.SUCCESS
    )
    click.secho("You can now assign sensors to it.", **MsgStyle.SUCCESS)


@fm_add_data.command("initial-structure")
@with_appcontext
def add_initial_structure():
    """Initialize useful structural data."""
    populate_initial_structure(db)


@fm_add_data.command("source")
@with_appcontext
@click.option(
    "--name",
    required=True,
    type=str,
    help="Name of the source (usually an organization)",
)
@click.option(
    "--model",
    required=False,
    type=str,
    help="Optionally, specify a model (for example, a class name, function name or url).",
)
@click.option(
    "--version",
    required=False,
    type=str,
    help="Optionally, specify a version (for example, '1.0'.",
)
@click.option(
    "--type",
    "source_type",
    required=True,
    type=str,
    help="Type of source (for example, 'forecaster' or 'scheduler').",
)
def add_source(name: str, model: str, version: str, source_type: str):
    source = get_or_create_source(
        source=name,
        model=model,
        version=version,
        source_type=source_type,
    )
    db.session.commit()
    click.secho(f"Added source {source.__repr__()}", **MsgStyle.SUCCESS)


@fm_add_data.command("beliefs")
@with_appcontext
@click.argument("file", type=click.Path(exists=True))
@click.option(
    "--sensor-id",
    "sensor",
    required=True,
    type=SensorIdField(),
    help="Record the beliefs under this sensor. Follow up with the sensor's ID. ",
)
@click.option(
    "--source",
    required=True,
    type=str,
    help="Source of the beliefs (an existing source id or name, or a new name).",
)
@click.option(
    "--unit",
    required=False,
    type=str,
    help="Unit of the data, for conversion to the sensor unit, if possible (a string unit such as 'kW' or 'm³/h').\n"
    "Measurements of time itself that are formatted as a 'datetime' or 'timedelta' can be converted to a sensor unit representing time (such as 's' or 'h'),\n"
    "where datetimes are represented as a duration with respect to the UNIX epoch."
    "Hint: to switch the sign of the data, prepend a minus sign.\n"
    "For example, when assigning kW consumption data to a kW production sensor, use '-kW'.",
)
@click.option(
    "--horizon",
    required=False,
    type=int,
    help="Belief horizon in minutes (use positive horizon for ex-ante beliefs or negative horizon for ex-post beliefs).",
)
@click.option(
    "--cp",
    required=False,
    type=click.FloatRange(0, 1),
    help="Cumulative probability in the range [0, 1].",
)
@click.option(
    "--resample/--do-not-resample",
    default=True,
    help="Resample the data to fit the sensor's event resolution. "
    " Only downsampling is currently supported (for example, from hourly to daily data).",
)
@click.option(
    "--allow-overwrite/--do-not-allow-overwrite",
    default=False,
    help="Allow overwriting possibly already existing data.\n"
    "Not allowing overwriting can be much more efficient",
)
@click.option(
    "--skiprows",
    required=False,
    default=1,
    type=int,
    help="Number of rows to skip from the top. Set to >1 to skip additional headers.",
)
@click.option(
    "--na-values",
    required=False,
    multiple=True,
    help="Additional strings to recognize as NaN values. This argument can be given multiple times.",
)
@click.option(
    "--keep-default-na",
    default=False,
    type=bool,
    help="Whether or not to keep NaN values in the data.",
)
@click.option(
    "--nrows",
    required=False,
    type=int,
    help="Number of rows to read (from the top, after possibly skipping rows). Leave out to read all rows.",
)
@click.option(
    "--datecol",
    required=False,
    default=0,
    type=int,
    help="Column number with datetimes (0 is 1st column, the default)",
)
@click.option(
    "--valuecol",
    required=False,
    default=1,
    type=int,
    help="Column number with values (1 is 2nd column, the default)",
)
@click.option(
    "--beliefcol",
    required=False,
    type=int,
    help="Column number with datetimes",
)
@click.option(
    "--timezone",
    required=False,
    default=None,
    help="Timezone as string, e.g. 'UTC' or 'Europe/Amsterdam'",
)
@click.option(
    "--filter-column",
    "filter_columns",
    multiple=True,
    help="Set a column number to filter data. Use together with --filter-value.",
)
@click.option(
    "--filter-value",
    "filter_values",
    multiple=True,
    help="Set a column value to filter data. Only rows with this value will be added. Use together with --filter-column.",
)
@click.option(
    "--delimiter",
    required=True,
    type=str,
    default=",",
    help="[For CSV files] Character to delimit columns per row, defaults to comma",
)
@click.option(
    "--decimal",
    required=False,
    default=".",
    type=str,
    help="[For CSV files] decimal character, e.g. '.' for 10.5",
)
@click.option(
    "--thousands",
    required=False,
    default=None,
    type=str,
    help="[For CSV files] thousands separator, e.g. '.' for 10.035,2",
)
@click.option(
    "--sheet_number",
    required=False,
    type=int,
    help="[For xls or xlsx files] Sheet number with the data (0 is 1st sheet)",
)
def add_beliefs(
    file: str,
    sensor: Sensor,
    source: str,
    filter_columns: list[int],
    filter_values: list[int],
    unit: str | None = None,
    horizon: int | None = None,
    cp: float | None = None,
    resample: bool = True,
    allow_overwrite: bool = False,
    skiprows: int = 1,
    na_values: list[str] | None = None,
    keep_default_na: bool = False,
    nrows: int | None = None,
    datecol: int = 0,
    valuecol: int = 1,
    beliefcol: int | None = None,
    timezone: str | None = None,
    delimiter: str = ",",
    decimal: str = ".",
    thousands: str | None = None,
    sheet_number: int | None = None,
    **kwargs,  # in-code calls to this CLI command can set additional kwargs for use in pandas.read_csv or pandas.read_excel
):
    """Add sensor data from a CSV or Excel file.

    To use default settings, structure your CSV file as follows:

        - One header line (will be ignored!)
        - UTC datetimes in 1st column
        - values in 2nd column

    For example:

        Date,Inflow (cubic meter)
        2020-12-03 14:00,212
        2020-12-03 14:10,215.6
        2020-12-03 14:20,203.8

    In case no --horizon is specified and no beliefcol is specified,
    the moment of executing this CLI command is taken as the time at which the beliefs were recorded.
    """
    _source = parse_source(source)

    # Set up optional parameters for read_csv
    if file.split(".")[-1].lower() == "csv":
        kwargs["infer_datetime_format"] = True
        kwargs["delimiter"] = delimiter
        kwargs["decimal"] = decimal
        kwargs["thousands"] = thousands
    if sheet_number is not None:
        kwargs["sheet_name"] = sheet_number
    if horizon is not None:
        kwargs["belief_horizon"] = timedelta(minutes=horizon)
    elif beliefcol is None:
        kwargs["belief_time"] = server_now().astimezone(pytz.timezone(sensor.timezone))

    # Set up optional filters:
    if len(filter_columns) != len(filter_values):
        raise ValueError(
            "The number of filter columns and filter values should be the same."
        )
    filter_by_column = (
        dict(zip(filter_columns, filter_values)) if filter_columns else None
    )
    bdf = tb.read_csv(
        file,
        sensor,
        source=_source,
        cumulative_probability=cp,
        resample=resample,
        header=None,
        skiprows=skiprows,
        nrows=nrows,
        usecols=[datecol, valuecol]
        if beliefcol is None
        else [datecol, beliefcol, valuecol],
        parse_dates=True,
        na_values=na_values,
        keep_default_na=keep_default_na,
        timezone=timezone,
        filter_by_column=filter_by_column,
        **kwargs,
    )
    duplicate_rows = bdf.index.duplicated(keep="first")
    if any(duplicate_rows) > 0:
        click.secho(
            "Duplicates found. Dropping duplicates for the following records:",
            **MsgStyle.WARN,
        )
        click.secho(bdf[duplicate_rows], **MsgStyle.WARN)
        bdf = bdf[~duplicate_rows]
    if unit is not None:
        bdf["event_value"] = convert_units(
            bdf["event_value"],
            from_unit=unit,
            to_unit=sensor.unit,
            event_resolution=sensor.event_resolution,
        )
    try:
        TimedBelief.add(
            bdf,
            expunge_session=True,
            allow_overwrite=allow_overwrite,
            bulk_save_objects=True,
            commit_transaction=True,
        )
        click.secho(f"Successfully created beliefs\n{bdf}", **MsgStyle.SUCCESS)
    except IntegrityError as e:
        db.session.rollback()
        click.secho(
            f"Failed to create beliefs due to the following error: {e.orig}",
            **MsgStyle.ERROR,
        )
        if not allow_overwrite:
            click.secho(
                "As a possible workaround, use the --allow-overwrite flag.",
                **MsgStyle.ERROR,
            )


@fm_add_data.command("annotation")
@with_appcontext
@click.option(
    "--content",
    required=True,
    prompt="Enter annotation",
)
@click.option(
    "--at",
    "start_str",
    required=True,
    help="Annotation is set (or starts) at this datetime. Follow up with a timezone-aware datetime in ISO 6801 format.",
)
@click.option(
    "--until",
    "end_str",
    required=False,
    help="Annotation ends at this datetime. Follow up with a timezone-aware datetime in ISO 6801 format. Defaults to one (nominal) day after the start of the annotation.",
)
@click.option(
    "--account-id",
    "account_ids",
    type=click.INT,
    multiple=True,
    help="Add annotation to this organisation account. Follow up with the account's ID. This argument can be given multiple times.",
)
@click.option(
    "--asset-id",
    "generic_asset_ids",
    type=int,
    multiple=True,
    help="Add annotation to this asset. Follow up with the asset's ID. This argument can be given multiple times.",
)
@click.option(
    "--sensor-id",
    "sensor_ids",
    type=int,
    multiple=True,
    help="Add annotation to this sensor. Follow up with the sensor's ID. This argument can be given multiple times.",
)
@click.option(
    "--user-id",
    type=int,
    required=True,
    help="Attribute annotation to this user. Follow up with the user's ID.",
)
def add_annotation(
    content: str,
    start_str: str,
    end_str: str | None,
    account_ids: list[int],
    generic_asset_ids: list[int],
    sensor_ids: list[int],
    user_id: int,
):
    """Add annotation to accounts, assets and/or sensors."""

    # Parse input
    start = pd.Timestamp(start_str)
    end = (
        pd.Timestamp(end_str)
        if end_str is not None
        else start + pd.offsets.DateOffset(days=1)
    )
    accounts = (
        db.session.query(Account).filter(Account.id.in_(account_ids)).all()
        if account_ids
        else []
    )
    assets = (
        db.session.query(GenericAsset)
        .filter(GenericAsset.id.in_(generic_asset_ids))
        .all()
        if generic_asset_ids
        else []
    )
    sensors = (
        db.session.query(Sensor).filter(Sensor.id.in_(sensor_ids)).all()
        if sensor_ids
        else []
    )
    user = db.session.query(User).get(user_id)
    _source = get_or_create_source(user)

    # Create annotation
    annotation = get_or_create_annotation(
        Annotation(
            content=content,
            start=start,
            end=end,
            source=_source,
            type="label",
        )
    )
    for account in accounts:
        account.annotations.append(annotation)
    for asset in assets:
        asset.annotations.append(annotation)
    for sensor in sensors:
        sensor.annotations.append(annotation)
    db.session.commit()
    click.secho("Successfully added annotation.", **MsgStyle.SUCCESS)


@fm_add_data.command("holidays")
@with_appcontext
@click.option(
    "--year",
    type=click.INT,
    help="The year for which to look up holidays",
)
@click.option(
    "--country",
    "countries",
    type=click.STRING,
    multiple=True,
    help="The ISO 3166-1 country/region or ISO 3166-2 sub-region for which to look up holidays (such as US, BR and DE). This argument can be given multiple times.",
)
@click.option(
    "--asset-id",
    "generic_asset_ids",
    type=click.INT,
    multiple=True,
    help="Add annotations to this asset. Follow up with the asset's ID. This argument can be given multiple times.",
)
@click.option(
    "--account-id",
    "account_ids",
    type=click.INT,
    multiple=True,
    help="Add annotations to this account. Follow up with the account's ID. This argument can be given multiple times.",
)
def add_holidays(
    year: int,
    countries: list[str],
    generic_asset_ids: list[int],
    account_ids: list[int],
):
    """Add holiday annotations to accounts and/or assets."""
    calendars = workalendar_registry.get_calendars(countries)
    num_holidays = {}

    accounts = (
        db.session.query(Account).filter(Account.id.in_(account_ids)).all()
        if account_ids
        else []
    )
    assets = (
        db.session.query(GenericAsset)
        .filter(GenericAsset.id.in_(generic_asset_ids))
        .all()
        if generic_asset_ids
        else []
    )
    annotations = []
    for country, calendar in calendars.items():
        _source = get_or_create_source(
            "workalendar", model=country, source_type="CLI script"
        )
        holidays = calendar().holidays(year)
        for holiday in holidays:
            start = pd.Timestamp(holiday[0])
            end = start + pd.offsets.DateOffset(days=1)
            annotations.append(
                get_or_create_annotation(
                    Annotation(
                        content=holiday[1],
                        start=start,
                        end=end,
                        source=_source,
                        type="holiday",
                    )
                )
            )
        num_holidays[country] = len(holidays)
    db.session.add_all(annotations)
    for account in accounts:
        account.annotations += annotations
    for asset in assets:
        asset.annotations += annotations
    db.session.commit()
    click.secho(
        f"Successfully added holidays to {len(accounts)} {flexmeasures_inflection.pluralize('account', len(accounts))} and {len(assets)} {flexmeasures_inflection.pluralize('asset', len(assets))}:\n{num_holidays}",
        **MsgStyle.SUCCESS,
    )


@fm_add_data.command("forecasts")
@with_appcontext
@click.option(
    "--sensor-id",
    "sensor_ids",
    multiple=True,
    required=True,
    help="Create forecasts for this sensor. Follow up with the sensor's ID. This argument can be given multiple times.",
)
@click.option(
    "--from-date",
    "from_date_str",
    default="2015-02-08",
    help="Forecast from date (inclusive). Follow up with a date in the form yyyy-mm-dd.",
)
@click.option(
    "--to-date",
    "to_date_str",
    default="2015-12-31",
    help="Forecast to date (inclusive). Follow up with a date in the form yyyy-mm-dd.",
)
@click.option(
    "--resolution",
    type=int,
    help="Resolution of forecast in minutes. If not set, resolution is determined from the sensor to be forecasted",
)
@click.option(
    "--horizon",
    "horizons_as_hours",
    multiple=True,
    type=click.Choice(["1", "6", "24", "48"]),
    default=["1", "6", "24", "48"],
    help="Forecasting horizon in hours. This argument can be given multiple times. Defaults to all possible horizons.",
)
@click.option(
    "--as-job",
    is_flag=True,
    help="Whether to queue a forecasting job instead of computing directly. "
    "To process the job, run a worker (on any computer, but configured to the same databases) to process the 'forecasting' queue. Defaults to False.",
)
def create_forecasts(
    sensor_ids: list[int],
    from_date_str: str = "2015-02-08",
    to_date_str: str = "2015-12-31",
    horizons_as_hours: list[str] = ["1"],
    resolution: int | None = None,
    as_job: bool = False,
):
    """
    Create forecasts.

    For example:

        --from-date 2015-02-02 --to-date 2015-02-04 --horizon 6 --sensor-id 12 --sensor-id 14

        This will create forecast values from 0am on May 2nd to 0am on May 5th,
        based on a 6-hour horizon, for sensors 12 and 14.

    """
    # make horizons
    horizons = [timedelta(hours=int(h)) for h in horizons_as_hours]

    # apply timezone and set forecast_end to be an inclusive version of to_date
    timezone = app.config.get("FLEXMEASURES_TIMEZONE")
    forecast_start = pd.Timestamp(from_date_str).tz_localize(timezone)
    forecast_end = (pd.Timestamp(to_date_str) + pd.Timedelta("1D")).tz_localize(
        timezone
    )

    event_resolution: timedelta | None
    if resolution is not None:
        event_resolution = timedelta(minutes=resolution)
    else:
        event_resolution = None

    if as_job:
        num_jobs = 0
        for sensor_id in sensor_ids:
            for horizon in horizons:
                # Note that this time period refers to the period of events we are forecasting, while in create_forecasting_jobs
                # the time period refers to the period of belief_times, therefore we are subtracting the horizon.
                jobs = create_forecasting_jobs(
                    sensor_id=sensor_id,
                    horizons=[horizon],
                    start_of_roll=forecast_start - horizon,
                    end_of_roll=forecast_end - horizon,
                )
                num_jobs += len(jobs)
        click.secho(
            f"{num_jobs} new forecasting job(s) added to the queue.",
            **MsgStyle.SUCCESS,
        )
    else:
        from flexmeasures.data.scripts.data_gen import populate_time_series_forecasts

        populate_time_series_forecasts(  # this function reports its own output
            db=app.db,
            sensor_ids=sensor_ids,
            horizons=horizons,
            forecast_start=forecast_start,
            forecast_end=forecast_end,
            event_resolution=event_resolution,
        )


# todo: repurpose `flexmeasures add schedule` (deprecated since v0.12),
#       - see https://github.com/FlexMeasures/flexmeasures/pull/537#discussion_r1048680231
#       - hint for repurposing to invoke custom logic instead of a default subcommand:
#             @fm_add_data.group("schedule", invoke_without_command=True)
#             def create_schedule():
#                 if ctx.invoked_subcommand:
#                     ...
@fm_add_data.group(
    "schedule",
    cls=DeprecatedDefaultGroup,
    default="storage",
    deprecation_message="The command 'flexmeasures add schedule' is deprecated. Please use `flexmeasures add schedule for-storage` instead.",
)
@click.pass_context
@with_appcontext
def create_schedule(ctx):
    """(Deprecated) Create a new schedule for a given power sensor.

    THIS COMMAND HAS BEEN RENAMED TO `flexmeasures add schedule for-storage`
    """
    pass


@create_schedule.command("for-storage")
@with_appcontext
@click.option(
    "--sensor-id",
    "power_sensor",
    type=SensorIdField(),
    required=True,
    help="Create schedule for this sensor. Should be a power sensor. Follow up with the sensor's ID.",
)
@click.option(
    "--consumption-price-sensor",
    "consumption_price_sensor",
    type=SensorIdField(),
    required=False,
    help="Optimize consumption against this sensor. The sensor typically records an electricity price (e.g. in EUR/kWh), but this field can also be used to optimize against some emission intensity factor (e.g. in kg CO₂ eq./kWh). Follow up with the sensor's ID.",
)
@click.option(
    "--production-price-sensor",
    "production_price_sensor",
    type=SensorIdField(),
    required=False,
    help="Optimize production against this sensor. Defaults to the consumption price sensor. The sensor typically records an electricity price (e.g. in EUR/kWh), but this field can also be used to optimize against some emission intensity factor (e.g. in kg CO₂ eq./kWh). Follow up with the sensor's ID.",
)
@click.option(
    "--optimization-context-id",
    "optimization_context_sensor",
    type=SensorIdField(),
    required=False,
    help="To be deprecated. Use consumption-price-sensor instead.",
)
@click.option(
    "--inflexible-device-sensor",
    "inflexible_device_sensors",
    type=SensorIdField(),
    multiple=True,
    help="Take into account the power flow of inflexible devices. Follow up with the sensor's ID."
    " This argument can be given multiple times.",
)
@click.option(
    "--start",
    "start",
    type=AwareDateTimeField(format="iso"),
    required=True,
    help="Schedule starts at this datetime. Follow up with a timezone-aware datetime in ISO 6801 format.",
)
@click.option(
    "--duration",
    "duration",
    type=DurationField(),
    required=True,
    help="Duration of schedule, after --start. Follow up with a duration in ISO 6801 format, e.g. PT1H (1 hour) or PT45M (45 minutes).",
)
@click.option(
    "--soc-at-start",
    "soc_at_start",
    type=QuantityField("%", validate=validate.Range(min=0, max=1)),
    required=True,
    help="State of charge (e.g 32.8%, or 0.328) at the start of the schedule.",
)
@click.option(
    "--soc-target",
    "soc_target_strings",
    type=click.Tuple(
        types=[QuantityField("%", validate=validate.Range(min=0, max=1)), str]
    ),
    multiple=True,
    required=False,
    help="Target state of charge (e.g 100%, or 1) at some datetime. Follow up with a float value and a timezone-aware datetime in ISO 6081 format."
    " This argument can be given multiple times."
    " For example: --soc-target 100% 2022-02-23T13:40:52+00:00",
)
@click.option(
    "--soc-min",
    "soc_min",
    type=QuantityField("%", validate=validate.Range(min=0, max=1)),
    required=False,
    help="Minimum state of charge (e.g 20%, or 0.2) for the schedule.",
)
@click.option(
    "--soc-max",
    "soc_max",
    type=QuantityField("%", validate=validate.Range(min=0, max=1)),
    required=False,
    help="Maximum state of charge (e.g 80%, or 0.8) for the schedule.",
)
@click.option(
    "--roundtrip-efficiency",
    "roundtrip_efficiency",
    type=EfficiencyField(),
    required=False,
    default=1,
    help="Round-trip efficiency (e.g. 85% or 0.85) to use for the schedule. Defaults to 100% (no losses).",
)
@click.option(
    "--storage-efficiency",
    "storage_efficiency",
    type=EfficiencyField(),
    required=False,
    default=1,
    help="Storage efficiency (e.g. 95% or 0.95) to use for the schedule,"
    " applied over each time step equal to the sensor resolution."
    " For example, a storage efficiency of 99 percent per (absolute) day, for scheduling a 1-hour resolution sensor, should be passed as a storage efficiency of 0.99**(1/24)."
    " Defaults to 100% (no losses).",
)
@click.option(
    "--as-job",
    is_flag=True,
    help="Whether to queue a scheduling job instead of computing directly. "
    "To process the job, run a worker (on any computer, but configured to the same databases) to process the 'scheduling' queue. Defaults to False.",
)
def add_schedule_for_storage(
    power_sensor: Sensor,
    consumption_price_sensor: Sensor,
    production_price_sensor: Sensor,
    optimization_context_sensor: Sensor,
    inflexible_device_sensors: list[Sensor],
    start: datetime,
    duration: timedelta,
    soc_at_start: ur.Quantity,
    soc_target_strings: list[tuple[ur.Quantity, str]],
    soc_min: ur.Quantity | None = None,
    soc_max: ur.Quantity | None = None,
    roundtrip_efficiency: ur.Quantity | None = None,
    storage_efficiency: ur.Quantity | None = None,
    as_job: bool = False,
):
    """Create a new schedule for a storage asset.

    Current limitations:

    - Limited to power sensors (probably possible to generalize to non-electric assets)
    - Only supports datetimes on the hour or a multiple of the sensor resolution thereafter
    """
    # todo: deprecate the 'optimization-context-id' argument in favor of 'consumption-price-sensor' (announced v0.11.0)
    tb_utils.replace_deprecated_argument(
        "optimization-context-id",
        optimization_context_sensor,
        "consumption-price-sensor",
        consumption_price_sensor,
    )

    # Parse input and required sensor attributes
    if not power_sensor.measures_power:
        click.secho(
            f"Sensor with ID {power_sensor.id} is not a power sensor.",
            **MsgStyle.ERROR,
        )
        raise click.Abort()
    if production_price_sensor is None:
        production_price_sensor = consumption_price_sensor
    end = start + duration

    # Convert SoC units (we ask for % in this CLI) to MWh, given the storage capacity
    try:
        check_required_attributes(power_sensor, [("max_soc_in_mwh", float)])
    except MissingAttributeException:
        click.secho(
            f"Sensor {power_sensor} has no max_soc_in_mwh attribute.", **MsgStyle.ERROR
        )
        raise click.Abort()
    capacity_str = f"{power_sensor.get_attribute('max_soc_in_mwh')} MWh"
    soc_at_start = convert_units(soc_at_start.magnitude, soc_at_start.units, "MWh", capacity=capacity_str)  # type: ignore
    soc_targets = []
    for soc_target_tuple in soc_target_strings:
        soc_target_value_str, soc_target_dt_str = soc_target_tuple
        soc_target_value = convert_units(
            soc_target_value_str.magnitude,
            str(soc_target_value_str.units),
            "MWh",
            capacity=capacity_str,
        )
        soc_target_datetime = pd.Timestamp(soc_target_dt_str)
        soc_targets.append(dict(value=soc_target_value, datetime=soc_target_datetime))
    if soc_min is not None:
        soc_min = convert_units(soc_min.magnitude, str(soc_min.units), "MWh", capacity=capacity_str)  # type: ignore
    if soc_max is not None:
        soc_max = convert_units(soc_max.magnitude, str(soc_max.units), "MWh", capacity=capacity_str)  # type: ignore
    if roundtrip_efficiency is not None:
        roundtrip_efficiency = roundtrip_efficiency.magnitude / 100.0
    if storage_efficiency is not None:
        storage_efficiency = storage_efficiency.magnitude / 100.0

    scheduling_kwargs = dict(
        start=start,
        end=end,
        belief_time=server_now(),
        resolution=power_sensor.event_resolution,
        flex_model={
            "soc-at-start": soc_at_start,
            "soc-targets": soc_targets,
            "soc-min": soc_min,
            "soc-max": soc_max,
            "soc-unit": "MWh",
            "roundtrip-efficiency": roundtrip_efficiency,
            "storage-efficiency": storage_efficiency,
        },
        flex_context={
            "consumption-price-sensor": consumption_price_sensor.id,
            "production-price-sensor": production_price_sensor.id,
            "inflexible-device-sensors": [s.id for s in inflexible_device_sensors],
        },
    )
    if as_job:
        job = create_scheduling_job(sensor=power_sensor, **scheduling_kwargs)
        if job:
            click.secho(
                f"New scheduling job {job.id} has been added to the queue.",
                **MsgStyle.SUCCESS,
            )
    else:
        success = make_schedule(sensor_id=power_sensor.id, **scheduling_kwargs)
        if success:
            click.secho("New schedule is stored.", **MsgStyle.SUCCESS)


@create_schedule.command("for-process")
@with_appcontext
@click.option(
    "--sensor-id",
    "power_sensor",
    type=SensorIdField(),
    required=True,
    help="Create schedule for this sensor. Should be a power sensor. Follow up with the sensor's ID.",
)
@click.option(
    "--consumption-price-sensor",
    "consumption_price_sensor",
    type=SensorIdField(),
    required=False,
    help="Optimize consumption against this sensor. The sensor typically records an electricity price (e.g. in EUR/kWh), but this field can also be used to optimize against some emission intensity factor (e.g. in kg CO₂ eq./kWh). Follow up with the sensor's ID.",
)
@click.option(
    "--start",
    "start",
    type=AwareDateTimeField(format="iso"),
    required=True,
    help="Schedule starts at this datetime. Follow up with a timezone-aware datetime in ISO 6801 format.",
)
@click.option(
    "--duration",
    "duration",
    type=DurationField(),
    required=True,
    help="Duration of schedule, after --start. Follow up with a duration in ISO 6801 format, e.g. PT1H (1 hour) or PT45M (45 minutes).",
)
@click.option(
    "--process-duration",
    "process_duration",
    type=DurationField(),
    required=True,
    help="Duration of the process. Follow up with a duration in ISO 6801 format, e.g. PT1H (1 hour) or PT45M (45 minutes).",
)
@click.option(
    "--process-type",
    "process_type",
    type=click.Choice(["INFLEXIBLE", "BREAKABLE", "SHIFTABLE"], case_sensitive=False),
    required=False,
    default="SHIFTABLE",
    help="Process schedule policy: INFLEXIBLE, BREAKABLE or SHIFTABLE.",
)
@click.option(
    "--process-power",
    "process_power",
    type=ur.Quantity,
    required=True,
    help="Constant power of the process during the activation period, e.g. 4kW.",
)
@click.option(
    "--forbid",
    type=TimeIntervalField(),
    multiple=True,
    required=False,
    help="Add time restrictions to the optimization, where the load will not be scheduled into."
    'Use the following format to define the restrictions: `{"start":<timezone-aware datetime in ISO 6801>, "duration":<ISO 6801 duration>}`'
    "This options allows to define multiple time restrictions by using the --forbid for different periods.",
)
@click.option(
    "--as-job",
    is_flag=True,
    help="Whether to queue a scheduling job instead of computing directly. "
    "To process the job, run a worker (on any computer, but configured to the same databases) to process the 'scheduling' queue. Defaults to False.",
)
def add_schedule_process(
    power_sensor: Sensor,
    consumption_price_sensor: Sensor,
    start: datetime,
    duration: timedelta,
    process_duration: timedelta,
    process_type: str,
    process_power: ur.Quantity,
    forbid: List | None = None,
    as_job: bool = False,
):
    """Create a new schedule for a process asset.

    Current limitations:
    - Only supports consumption blocks.
    - Not taking into account grid constraints or other processes.
    """

    if forbid is None:
        forbid = []

    # Parse input and required sensor attributes
    if not power_sensor.measures_power:
        click.secho(
            f"Sensor with ID {power_sensor.id} is not a power sensor.",
            **MsgStyle.ERROR,
        )
        raise click.Abort()

    end = start + duration

    process_power = convert_units(process_power.magnitude, process_power.units, "MW")  # type: ignore

    scheduling_kwargs = dict(
        start=start,
        end=end,
        belief_time=server_now(),
        resolution=power_sensor.event_resolution,
        flex_model={
            "duration": pd.Timedelta(process_duration).isoformat(),
            "process-type": process_type,
            "power": process_power,
            "time-restrictions": [TimeIntervalSchema().dump(f) for f in forbid],
        },
    )

    if consumption_price_sensor is not None:
        scheduling_kwargs["flex_context"] = {
            "consumption-price-sensor": consumption_price_sensor.id,
        }

    if as_job:
        job = create_scheduling_job(sensor=power_sensor, **scheduling_kwargs)
        if job:
            click.secho(
                f"New scheduling job {job.id} has been added to the queue.",
                **MsgStyle.SUCCESS,
            )
    else:
        success = make_schedule(sensor_id=power_sensor.id, **scheduling_kwargs)
        if success:
            click.secho("New schedule is stored.", **MsgStyle.SUCCESS)


@fm_add_data.command("report")
@with_appcontext
@click.option(
    "--sensor-id",
    "sensor",
    type=SensorIdField(),
    required=False,
    help="Sensor used to save the report. Follow up with the sensor's ID. Can be defined in the parameters file, as well"
    " If needed, use `flexmeasures add sensor` to create a new sensor first.",
)
@click.option(
    "--config",
    "config_file",
    required=False,
    type=click.File("r"),
    help="Path to the JSON or YAML file with the configuration of the reporter.",
)
@click.option(
    "--parameters",
    "parameters_file",
    required=False,
    type=click.File("r"),
    help="Path to the JSON or YAML file with the report parameters (passed to the compute step).",
)
@click.option(
    "--reporter",
    "reporter_class",
    default="PandasReporter",
    type=click.STRING,
    help="Reporter class registered in flexmeasures.data.models.reporting or in an available flexmeasures plugin."
    " Use the command `flexmeasures show reporters` to list all the available reporters.",
)
@click.option(
    "--start",
    "start",
    type=AwareDateTimeField(format="iso"),
    required=False,
    help="Report start time. `--start-offset` can be used instead. Follow up with a timezone-aware datetime in ISO 6801 format.",
)
@click.option(
    "--start-offset",
    "start_offset",
    type=str,
    required=False,
    help="Report start offset time from now. Use multiple Pandas offset strings separated by commas, e.g: -3D,DB,1W. Use DB or HB to offset to the begin of the day or hour, respectively.",
)
@click.option(
    "--end-offset",
    "end_offset",
    type=str,
    required=False,
    help="Report end offset time from now. Use multiple Pandas offset strings separated by commas, e.g: -3D,DB,1W. Use DB or HB to offset to the begin of the day or hour, respectively.",
)
@click.option(
    "--end",
    "end",
    type=AwareDateTimeField(format="iso"),
    required=False,
    help="Report end time. `--end-offset` can be used instead. Follow up with a timezone-aware datetime in ISO 6801 format.",
)
@click.option(
    "--resolution",
    "resolution",
    type=DurationField(format="iso"),
    required=False,
    help="Time resolution of the input time series to employ for the calculations. Follow up with a ISO 8601 duration string",
)
@click.option(
    "--output-file",
    "output_file",
    required=False,
    type=click.Path(),
    help="Path to save the report to file. Will override any previous file contents."
    " Use the `.csv` suffix to save the results as Comma Separated Values and `.xlsx` to export them as Excel sheets.",
)
@click.option(
    "--timezone",
    "timezone",
    required=False,
    help="Timezone as string, e.g. 'UTC' or 'Europe/Amsterdam' (defaults to the timezone of the sensor used to save the report).",
)
@click.option(
    "--dry-run",
    "dry_run",
    is_flag=True,
    help="Add this flag to avoid saving the results to the database.",
)
@click.option(
    "--edit-config",
    "edit_config",
    is_flag=True,
    help="Add this flag to edit the configuration of the Reporter in your default text editor (e.g. nano).",
)
@click.option(
    "--edit-parameters",
    "edit_parameters",
    is_flag=True,
    help="Add this flag to edit the parameters passed to the Reporter in your default text editor (e.g. nano).",
)
@click.option(
    "--save-config",
    "save_config",
    is_flag=True,
    help="Add this flag to save the `config` in the attributes of the DataSource for future reference.",
)
def add_report(  # noqa: C901
    reporter_class: str,
    sensor: Sensor | None = None,
    config_file: TextIOBase | None = None,
    parameters_file: TextIOBase | None = None,
    start: datetime | None = None,
    end: datetime | None = None,
    start_offset: str | None = None,
    end_offset: str | None = None,
    resolution: timedelta | None = None,
    output_file: Path | None = None,
    dry_run: bool = False,
    edit_config: bool = False,
    edit_parameters: bool = False,
    save_config: bool = False,
    timezone: str | None = None,
):
    """
    Create a new report using the Reporter class and save the results
    to the database or export them as CSV or Excel file.
    """

    config = dict()

    if config_file:
        config = yaml.safe_load(config_file)

    if edit_config:
        config = launch_editor("/tmp/config.yml")

    parameters = dict()

    if parameters_file:
        parameters = yaml.safe_load(parameters_file)

    if edit_parameters:
        parameters = launch_editor("/tmp/parameters.yml")

    if sensor is not None:
        parameters["sensor"] = sensor.id

    # check if sensor is not provided either in the parameters or the CLI
    # click parameter
    if parameters.get("sensor") is None:
        click.secho(
            "Report sensor needs to be defined, either on the `parameters` file or trough the --sensor CLI parameter...",
            **MsgStyle.ERROR,
        )
        raise click.Abort()

    sensor = Sensor.query.get(parameters.get("sensor"))

    # compute now in the timezone local to the output sensor
    if timezone is not None:
        check_timezone(timezone)

    now = pytz.timezone(
        zone=timezone if timezone is not None else sensor.timezone
    ).localize(datetime.now())

    # apply offsets, if provided
    if start_offset is not None:
        if start is None:
            start = now
        start = apply_offset_chain(start, start_offset)

    if end_offset is not None:
        if end is None:
            end = now
        end = apply_offset_chain(end, end_offset)

    # the case of not getting --start or --start-offset
    if start is None:
        click.secho(
            "Either --start or --start-offset should be provided."
            " Trying to use the latest datapoint of the report sensor as the start time...",
            **MsgStyle.WARN,
        )
        last_value_datetime = (
            db.session.query(func.max(TimedBelief.event_start))
            .filter(TimedBelief.sensor_id == sensor.id)
            .one_or_none()
        )

        # If there's data saved to the reporter sensors
        if last_value_datetime[0] is not None:
            start = last_value_datetime[0]
        else:
            click.secho(
                f"Could not find any data for the report sensor {sensor}.",
                **MsgStyle.ERROR,
            )
            raise click.Abort()

    # the case of not getting --end or --end-offset
    if end is None:
        click.secho(
            "Either --end or --end-offset should be provided."
            " Trying to use the current time as the end...",
            **MsgStyle.WARN,
        )
        end = now

    click.echo(f"Report scope:\n\tstart: {start}\n\tend:   {end}")

    click.echo(
        f"Looking for the Reporter {reporter_class} among all the registered reporters...",
    )

    # get reporter class
    ReporterClass: Type[Reporter] = app.data_generators.get("reporter").get(
        reporter_class
    )

    # check if it exists
    if ReporterClass is None:
        click.secho(
            f"Reporter class `{reporter_class}` not available.",
            **MsgStyle.ERROR,
        )
        raise click.Abort()

    click.secho(f"Reporter {reporter_class} found.", **MsgStyle.SUCCESS)

    # initialize reporter class with the reporter sensor and reporter config
    reporter: Reporter = ReporterClass(config=config, save_config=save_config)

    click.echo("Report computation is running...")

    if ("start" not in parameters) and (start is not None):
        parameters["start"] = start.isoformat()
    if ("end" not in parameters) and (end is not None):
        parameters["end"] = end.isoformat()
    if ("resolution" not in parameters) and (resolution is not None):
        parameters["resolution"] = pd.Timedelta(resolution).isoformat()

    # compute the report
    result: BeliefsDataFrame = reporter.compute(parameters=parameters)

    if not result.empty:
        click.secho("Report computation done.", **MsgStyle.SUCCESS)
    else:
        click.secho(
            "Report computation done, but the report is empty.", **MsgStyle.WARN
        )

    # save the report if it's not running in dry mode
    if not dry_run:
        click.echo("Saving report to the database...")
        save_to_db(result.dropna())
        db.session.commit()
        click.secho(
            "Success. The report has been saved to the database.",
            **MsgStyle.SUCCESS,
        )
    else:
        click.echo(
            f"Not saving report to the database (because of --dry-run), but this is what I computed:\n{result}"
        )

    # if an output file path is provided, save the results
    if output_file:
        suffix = str(output_file).split(".")[-1] if "." in str(output_file) else ""

        if suffix == "xlsx":  # save to EXCEL
            result.to_excel(output_file)
            click.secho(
                f"Success. The report has been exported as EXCEL to the file `{output_file}`",
                **MsgStyle.SUCCESS,
            )

        elif suffix == "csv":  # save to CSV
            result.to_csv(output_file)
            click.secho(
                f"Success. The report has been exported as CSV to the file `{output_file}`",
                **MsgStyle.SUCCESS,
            )

        else:  # default output format: CSV.
            click.secho(
                f"File suffix not provided. Exporting results as CSV to file {output_file}",
                **MsgStyle.WARN,
            )
            result.to_csv(output_file)
    else:
        click.secho(
            "Success.",
            **MsgStyle.SUCCESS,
        )


def launch_editor(filename: str) -> dict:
    """Launch editor to create/edit a json object"""
    click.edit("{\n}", filename=filename)

    with open(filename, "r") as f:
        content = yaml.safe_load(f)
        if content is None:
            return dict()

        return content


@fm_add_data.command("toy-account")
@with_appcontext
@click.option(
    "--kind",
    default="battery",
    type=click.Choice(["battery", "process"]),
    help="What kind of toy account. Defaults to a battery.",
)
@click.option("--name", type=str, default="Toy Account", help="Name of the account")
def add_toy_account(kind: str, name: str):
    """
    Create a toy account, for tutorials and trying things.
    """
    asset_types = add_default_asset_types(db=db)
    location = (52.374, 4.88969)  # Amsterdam

    # make an account (if not exist)
    account = Account.query.filter(Account.name == name).one_or_none()
    if account:
        click.secho(
            f"Account '{account}' already exists. Skipping account creation. Use `flexmeasures delete account --id {account.id}` if you need to remove it.",
            **MsgStyle.WARN,
        )

    # make an account user (account-admin?)
    email = "toy-user@flexmeasures.io"
    user = User.query.filter_by(email=email).one_or_none()
    if user is not None:
        click.secho(
            f"User with email {email} already exists in account {user.account.name}.",
            **MsgStyle.WARN,
        )
    else:
        user = create_user(
            email=email,
            check_email_deliverability=False,
            password="toy-password",
            user_roles=["account-admin"],
            account_name=name,
        )
        click.secho(
            f"Toy account {name} with user {user.email} created successfully. You might want to run `flexmeasures show account --id {user.account.id}`",
            **MsgStyle.SUCCESS,
        )

    db.session.commit()

    # add public day-ahead market (as sensor of transmission zone asset)
    nl_zone = add_transmission_zone_asset("NL", db=db)
    day_ahead_sensor = get_or_create_model(
        Sensor,
        name="day-ahead prices",
        generic_asset=nl_zone,
        unit="EUR/MWh",
        timezone="Europe/Amsterdam",
        event_resolution=timedelta(minutes=60),
        knowledge_horizon=(
            x_days_ago_at_y_oclock,
            {"x": 1, "y": 12, "z": "Europe/Paris"},
        ),
    )
    db.session.commit()
    click.secho(
        f"The sensor recording day-ahead prices is {day_ahead_sensor} (ID: {day_ahead_sensor.id}).",
        **MsgStyle.SUCCESS,
    )

    def create_power_asset(
        asset_name: str, asset_type: str, sensor_name: str, **attributes
    ):
        asset = get_or_create_model(
            GenericAsset,
            name=asset_name,
            generic_asset_type=asset_types[asset_type],
            owner=user.account,
            latitude=location[0],
            longitude=location[1],
        )
        asset.attributes = attributes
        power_sensor_specs = dict(
            generic_asset=asset,
            unit="MW",
            timezone="Europe/Amsterdam",
            event_resolution=timedelta(minutes=15),
        )
        power_sensor = get_or_create_model(
            Sensor,
            name=sensor_name,
            **power_sensor_specs,
        )
        return power_sensor

    if kind == "battery":
        # create battery
        discharging_sensor = create_power_asset(
            "toy-battery",
            "battery",
            "discharging",
            capacity_in_mw=0.5,
            min_soc_in_mwh=0.05,
            max_soc_in_mwh=0.45,
        )

        # create solar
        production_sensor = create_power_asset(
            "toy-solar",
            "solar",
            "production",
        )

        # add day-ahead price sensor and PV production sensor to show on the battery's asset page
        db.session.flush()
        battery = discharging_sensor.generic_asset
        battery.attributes["sensors_to_show"] = [
            day_ahead_sensor.id,
            [
                production_sensor.id,
                discharging_sensor.id,
            ],
        ]

        db.session.commit()

        click.secho(
            f"The sensor recording battery discharging is {discharging_sensor} (ID: {discharging_sensor.id}).",
            **MsgStyle.SUCCESS,
        )
        click.secho(
            f"The sensor recording solar forecasts is {production_sensor} (ID: {production_sensor.id}).",
            **MsgStyle.SUCCESS,
        )
    elif kind == "process":
        inflexible_power = create_power_asset(
            "toy-process",
            "process",
            "Power (Inflexible)",
        )

        breakable_power = create_power_asset(
            "toy-process",
            "process",
            "Power (Breakable)",
        )

        shiftable_power = create_power_asset(
            "toy-process",
            "process",
            "Power (Shiftable)",
        )

        db.session.flush()

        process = shiftable_power.generic_asset
        process.attributes["sensors_to_show"] = [
            day_ahead_sensor.id,
            inflexible_power.id,
            breakable_power.id,
            shiftable_power.id,
        ]

        db.session.commit()

        click.secho(
            f"The sensor recording the power of the inflexible load is {inflexible_power} (ID: {inflexible_power.id}).",
            **MsgStyle.SUCCESS,
        )
        click.secho(
            f"The sensor recording the power of the breakable load is {breakable_power} (ID: {breakable_power.id}).",
            **MsgStyle.SUCCESS,
        )
        click.secho(
            f"The sensor recording the power of the shiftable load is {shiftable_power} (ID: {shiftable_power.id}).",
            **MsgStyle.SUCCESS,
        )


app.cli.add_command(fm_add_data)


def check_timezone(timezone):
    try:
        pytz.timezone(timezone)
    except pytz.UnknownTimeZoneError:
        click.secho("Timezone %s is unknown!" % timezone, **MsgStyle.ERROR)
        raise click.Abort()


def check_errors(errors: dict[str, list[str]]):
    if errors:
        click.secho(
            f"Please correct the following errors:\n{errors}.\n Use the --help flag to learn more.",
            **MsgStyle.ERROR,
        )
        raise click.Abort()


def parse_source(source):
    if source.isdigit():
        _source = get_source_or_none(int(source))
        if not _source:
            click.secho(f"Failed to find source {source}.", **MsgStyle.ERROR)
            raise click.Abort()
    else:
        _source = get_or_create_source(source, source_type="CLI script")
    return _source