"""
Per-(device, origin) Wiener-process parameters for first-hitting-time RUL.

Models temp-corrected daily voltage as Brownian motion with drift; EOL is the
first passage of the validated 2.40 V threshold. For each (device, origin) we
estimate, from data <= origin only:

    wnr_m   - drift toward the threshold (V/day, positive = declining)
    wnr_s2  - diffusion variance (V^2/day) from drift-adjusted increments
    wnr_n   - number of increments behind the estimate (for shrinkage)

Estimation window: last FIT_DAYS days of in-band daily medians. Cached to
dataset/train/fht_features.parquet, keyed like features.parquet.
"""

from pathlib import Path
import os
import sys

import numpy as np
import pandas as pd

MIN_TEMP, MAX_TEMP = 10, 30
BETA = 0.005          # V/degC temperature correction (validated in features.py)
FIT_DAYS = 180

FHT_FEATURES = ['wnr_m', 'wnr_s2', 'wnr_n']


def daily_corrected(timeseries: pd.DataFrame) -> pd.DataFrame:
    df = timeseries.reset_index()
    df['device_id'] = df['device_id'].astype(str)
    df = df[(df['temperature'] > MIN_TEMP) & (df['temperature'] < MAX_TEMP)]
    df['day'] = df['end_time'].dt.normalize()
    g = df.groupby(['device_id', 'day'])
    out = pd.DataFrame({'v': g['voltage'].median(), 't': g['temperature'].median()})
    out['vc'] = out['v'] - BETA * (out['t'] - 20.0)
    return out.reset_index().sort_values('day')


def wiener_at_origin(daily: pd.DataFrame, origin: pd.Timestamp) -> pd.DataFrame:
    h = daily[(daily['day'] <= origin) &
              (daily['day'] > origin - pd.Timedelta(days=FIT_DAYS))]

    def fit(g):
        if len(g) < 4:
            return pd.Series({'wnr_m': np.nan, 'wnr_s2': np.nan,
                              'wnr_n': float(len(g)) - 1})
        tt = (g['day'] - g['day'].iloc[0]).dt.days.values.astype(float)
        vv = g['vc'].values
        span = tt[-1] - tt[0]
        if span < 7:
            return pd.Series({'wnr_m': np.nan, 'wnr_s2': np.nan,
                              'wnr_n': float(len(g)) - 1})
        # drift: endpoints estimator (MLE for BM); positive = declining voltage
        m = (vv[0] - vv[-1]) / span
        dt = np.diff(tt)
        dv = -np.diff(vv)          # positive = decline
        s2 = np.mean((dv - m * dt) ** 2 / dt)
        return pd.Series({'wnr_m': m, 'wnr_s2': s2, 'wnr_n': float(len(dt))})

    return h.groupby('device_id').apply(fit, include_groups=False)


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'))
    daily = daily_corrected(timeseries)
    rows = []
    for _, r in origins.iterrows():
        f = wiener_at_origin(daily, r['origin'])
        f['origin_idx'] = r['origin_idx']
        rows.append(f.reset_index().rename(columns={'index': 'device_id'}))
        print(f"fht 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' / 'fht_features.parquet')
    return out


def load(dataset_path: Path = Path('dataset'), force: bool = False) -> pd.DataFrame:
    p = dataset_path / 'train' / 'fht_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[FHT_FEATURES].describe().round(5).to_string())