ad-nix/apps/daylight_adjuster.py

import logging
from contextlib import suppress
from dataclasses import InitVar, dataclass, field
from datetime import datetime, timedelta, tzinfo
from typing import Dict, Iterable

import astral
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import pandas as pd
from astral import Observer, SunDirection
from astral.sun import elevation, sun, time_at_elevation
from IPython.display import display

HOME_TZ = datetime.now().astimezone().tzinfo


def format_x_axis(fig):
    ax: plt.Axes = fig.axes[0]
    ax.xaxis.set_major_locator(mdates.HourLocator(byhour=range(0, 24, 2)))
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%I%p'))
    ax.grid(True)
    fig.autofmt_xdate()


def normalize(s: pd.Series, min=None, max=None):
    min = min or s.min()
    max = max or s.max()
    rng = max - min
    return ((s - min) / rng) * 100


def get_next_time_at_elevation(*args, **kwargs):
    time = time_at_elevation(*args, **kwargs)
    if time < (now := datetime.now(HOME_TZ)):
        time = time_at_elevation(
            *args,
            date=(now + timedelta(days=1)).date(),
            **kwargs
        )
    return time


def get_next_sun_time(named_time: str, observer: Observer):
    sun_times_dict = sun(observer, datetime.today().date())
    try:
        time = sun_times_dict[named_time].astimezone()
    except KeyError:
        time = datetime.combine(
            datetime.today().date(),
            datetime.strptime(named_time, '%I:%M:%S%p').time()
        ).astimezone()

    if time < (now := datetime.now(HOME_TZ)):
        tomorrow = (now + timedelta(days=1)).date()
        sun_times_dict = sun(observer, tomorrow)
        try:
            time = sun_times_dict[named_time].astimezone()
        except KeyError:
            time = datetime.combine(
                tomorrow,
                datetime.strptime(named_time, '%I:%M:%S%p').time()
            ).astimezone()

    return time


def parse_periods(observer: Observer, periods: Dict):
    now = datetime.now(HOME_TZ)

    for period in periods:
        if 'elevation' in period:
            if period['direction'] == 'rising':
                dir = SunDirection.RISING
            elif period['direction'] == 'setting':
                dir = SunDirection.SETTING

            if isinstance(period['elevation'], int):
                time = get_next_time_at_elevation(
                    observer,
                    elevation=period['elevation'],
                    direction=dir,
                    tzinfo=HOME_TZ,
                )

            elif period['elevation'] == 'min':
                pass
            elif period['elevation'] == 'max':
                pass
        elif 'time' in period:
            time = get_next_sun_time(period['time'], observer)

        # res = {'time': time.replace(tzinfo=None)}
        # res = {'time': time.replace(tzinfo=HOME_TZ)}
        res = {'time': time}
        res.update({k: period[k] for k in ['brightness', 'color_temp'] if k in period})
        yield res


@dataclass
class DaylightAdjuster:
    latitude: float
    longitude: float
    periods: InitVar[Dict]
    datetime: datetime = field(default_factory=datetime.now)
    resolution: InitVar[int] = field(default=200)

    def __post_init__(self, periods: Dict, resolution: int):
        self.logger: logging.Logger = logging.getLogger(type(self).__name__)
        now = datetime.now().astimezone()
        times = pd.date_range(
            start=now, end=now + timedelta(days=1),
            periods=resolution,
            tz=HOME_TZ
        )
        self.logger.debug(times)

        pytimes = [dt.to_pydatetime() for dt in times]

        el = pd.Series(
            (elevation(self.observer, dt) for dt in pytimes),
            index=times, name='elevation'
        )
        self.logger.debug(el)

        # el.index = el.index.tz_convert(HOME_TZ)
        # el.index = el.index.tz_convert(None)

        self.periods = parse_periods(self.observer, periods)

        # self.df = pd.DataFrame(el)
        self.df = pd.concat([pd.DataFrame(self.periods).set_index('time'), el], axis=1)
        self.df = self.df.sort_index().interpolate().bfill().ffill()

        # self.df.index = self.df.index.to_series().dt.tz_localize(None)

    @property
    def observer(self) -> astral.Observer:
        return astral.Observer(self.latitude, self.longitude)

    @property
    def current_settings(self) -> pd.Series:
        return self.df[:datetime.now().astimezone()].iloc[-1].drop('elevation').astype(int).to_dict()

    def elevation_fig(self):
        fig, ax = plt.subplots(figsize=(10, 7))
        elevation = self.df['elevation']
        elevation.index = elevation.index.to_series().dt.tz_localize(None)
        handles = ax.plot(elevation)

        ax.set_ylabel('Elevation')
        ax.set_ylim(-100, 100)
        format_x_axis(fig)
        ax.set_xlim(elevation.index[0], elevation.index[-1])
        print(elevation)

        # ax.xaxis_date(HOME_TZ)

        ax2 = ax.twinx()
        handles.extend(ax2.plot(normalize(self.df['brightness'], 1, 255), 'tab:orange'))
        handles.extend(ax2.plot(normalize(self.df['color_temp'], 150, 650), 'tab:green'))
        ax2.set_ylabel('Brightness')
        ax2.set_ylim(0, 100)

        # handles.append(ax.axvline(datetime.now().astimezone(),
        #                           linestyle='--',
        #                           color='g'))

        # handles.append(ax2.axhline(self.get_brightness(),
        #                            linestyle='--',
        #                            color='r'))

        # handles.append(ax.axhline(self.get_elevation(),
        #                           linestyle='--',
        #                           color=handles[0].get_color()))

        ax.legend(handles=handles, loc='lower center', labels=[
            'Sun Elevation Angle',
            'Brightness Setting',
            'Color Temp Setting',
            'Current Time',
            # 'Current Brightness',
            # 'Current Elevation'
        ])

        fig.tight_layout()
        plt.close(fig)
        return fig