# SPDX-License-Identifier: MIT
"""YS-EN-0006 · synthetic office-CO2 series, rebuilt from the 2023 report's description.

Hourly values. Workdays: a truncated sine from ~08:00 to ~17:00 with a lunch dip
(the report's "double peak"), peaking near 15 above a baseline of 5. Non-workdays
(weekends and stat holidays): one low hump near 7. Gaussian noise on every hour.
An optional slow level offset stands in for the report's "trend" experiment.

This is a reconstruction from the report's text and figures, not its code.
"""
from __future__ import annotations

import datetime as dt

import numpy as np

BASE = 5.0


def day_profile(workday: bool) -> np.ndarray:
    h = np.arange(24) + 0.5
    if workday:
        occ = np.clip(1.6 * np.sin(np.pi * (h - 7.5) / 10.0), 0, 1)   # truncated sine
        occ *= 1 - 0.13 * np.exp(-((h - 12.5) / 0.9) ** 2)            # lunch dip → double peak
        return BASE + 10.0 * occ
    occ = np.clip(1.3 * np.sin(np.pi * (h - 9.5) / 8.0), 0, 1) * ((h > 9.5) & (h < 17.5))
    return BASE + 2.0 * occ


def make(start: dt.date, n_days: int, holidays: set[dt.date] = frozenset(),
         noise: float = 0.2, level: np.ndarray | None = None, seed: int = 0):
    """Return (y, workday) as hourly arrays of length 24*n_days.

    level: optional per-hour additive offset (same length) for trend experiments.
    """
    rng = np.random.default_rng(seed)
    days = [start + dt.timedelta(d) for d in range(n_days)]
    wd = np.array([d.weekday() < 5 and d not in holidays for d in days])
    y = np.concatenate([day_profile(w) for w in wd])
    y = y + rng.normal(0, noise, y.size)
    if level is not None:
        y = y + level
    return y.astype(np.float32), np.repeat(wd, 24).astype(np.float32), days


def random_level(n_hours: int, lo: float, hi: float, seg_days: int = 14, seed: int = 0):
    """Piecewise-linear level wandering within [lo, hi], new knot every seg_days."""
    rng = np.random.default_rng(seed + 1000)
    knots = rng.uniform(lo, hi, n_hours // (24 * seg_days) + 2)
    x = np.arange(n_hours) / (24 * seg_days)
    return np.interp(x, np.arange(knots.size), knots).astype(np.float32)