"""
Auxiliary (device, origin) features from raw hourly readings.

Everything in features.py is built on daily *medians*; this module mines the
within-day and meta-signal structure instead - discharge depth (daily voltage
range/std), telemetry intensity, thermal exposure, and reporting-gap dynamics.
Motivation: 27 of 60 remaining test-block misses had p42 < 0.10 because their
median voltage had not moved by the origin; usage stress may lead the droop.

Cached to dataset/train/aux_features.parquet, keyed (device_id, origin_idx),
aligned with features.parquet's origins. Only data <= origin is used.
"""

from pathlib import Path
import os
import sys

import numpy as np
import pandas as pd

MIN_TEMP, MAX_TEMP = 10, 30

AUX_FEATURES = ['rng_30', 'rng_90', 'rng_ratio', 'vstd_30', 'hot_90',
                't_med_90', 'reads_30', 'reads_90', 'oob_90',
                'gap_max_60', 'v_min_life']


def day_stats(timeseries: pd.DataFrame) -> pd.DataFrame:
    df = timeseries.reset_index()
    df['device_id'] = df['device_id'].astype(str)
    df['day'] = df['end_time'].dt.normalize()
    inband = df[(df['temperature'] > MIN_TEMP) & (df['temperature'] < MAX_TEMP)]

    g = inband.groupby(['device_id', 'day'])['voltage']
    stats = pd.DataFrame({
        'v_rng': g.max() - g.min(),
        'v_std': g.std(),
        'v_med': g.median(),
    })
    ga = df.groupby(['device_id', 'day'])
    stats['n_reads'] = ga.size()
    stats['t_med'] = ga['temperature'].median()
    stats['hot'] = ga['temperature'].apply(lambda t: (t > 28).mean())
    stats['oob'] = ga['temperature'].apply(
        lambda t: ((t <= MIN_TEMP) | (t >= MAX_TEMP)).mean())
    return stats.reset_index().sort_values('day')


def aux_at_origin(stats: pd.DataFrame, origin: pd.Timestamp) -> pd.DataFrame:
    h = stats[stats['day'] <= origin]
    out = {}

    for w in (30, 90):
        hw = h[h['day'] > origin - pd.Timedelta(days=w)]
        gw = hw.groupby('device_id')
        if w == 30:
            out['rng_30'] = gw['v_rng'].mean()
            out['vstd_30'] = gw['v_std'].mean()
            out['reads_30'] = gw['n_reads'].mean()
        else:
            out['rng_90'] = gw['v_rng'].mean()
            out['reads_90'] = gw['n_reads'].mean()
            out['hot_90'] = gw['hot'].mean()
            out['t_med_90'] = gw['t_med'].mean()
            out['oob_90'] = gw['oob'].mean()

    h60 = h[h['day'] > origin - pd.Timedelta(days=60)]

    def max_gap(days):
        d = np.sort(days.values)
        edges = np.concatenate([d, [np.datetime64(origin)]])
        return np.diff(edges).max() / np.timedelta64(1, 'D') if len(d) else np.nan

    out['gap_max_60'] = h60.groupby('device_id')['day'].apply(max_gap)
    out['v_min_life'] = h.groupby('device_id')['v_med'].min()

    f = pd.DataFrame(out)
    f['rng_ratio'] = f['rng_30'] / f['rng_90'].replace(0, np.nan)
    return f[AUX_FEATURES]


def build(dataset_path: Path) -> pd.DataFrame:
    from batteryswap_public.utils import load_dataset
    locations, timeseries, eol_times, scenarios = load_dataset(dataset_path / 'train')
    table = pd.read_parquet(dataset_path / 'train' / 'features.parquet')
    origins = (table[['origin_idx', 'origin']].drop_duplicates()
               .sort_values('origin_idx'))

    stats = day_stats(timeseries)
    rows = []
    for _, r in origins.iterrows():
        f = aux_at_origin(stats, r['origin'])
        f['origin_idx'] = r['origin_idx']
        rows.append(f.reset_index().rename(columns={'index': 'device_id'}))
        print(f"aux origin {r['origin_idx']} {r['origin'].date()}: {len(f)}",
              file=sys.stderr)
    out = pd.concat(rows, ignore_index=True)
    out.to_parquet(dataset_path / 'train' / 'aux_features.parquet')
    return out


def load(dataset_path: Path = Path('dataset'), force: bool = False) -> pd.DataFrame:
    p = dataset_path / 'train' / 'aux_features.parquet'
    if p.exists() and not force:
        return pd.read_parquet(p)
    return build(dataset_path)


if __name__ == '__main__':
    dataset_path = Path(os.environ.get('BATTERYSWAP_DATASET_PATH', 'dataset'))
    t = load(dataset_path, force='--force' in sys.argv)
    print(t.describe().round(3).to_string())