"""
Baseline planners and a forward-split evaluation harness.

The harness fits nothing (the heuristic baselines have no parameters and
`current` is pre-fitted), but it establishes the reference points any real
model has to beat:

    current    - whatever planners/best.pickle holds (the static fallback
                 planner; currently an ECRoutePlanner from planner.train_full)
    defer      - park every battery past the horizon (costs nothing but emergencies)
    voltage-k  - swap the k lowest-voltage batteries
    oracle     - knows the true EOL, swaps on the exact day

Scored on origins s_24..s_41, which sit after the training block (origins whose
42-day window closes before s_24) and before the origins whose windows run past
the end of the data.
"""

import os
from pathlib import Path
from typing import Optional

import numpy
import pandas
import structlog
from pydantic_settings import BaseSettings, SettingsConfigDict

from batteryswap_public.interfaces import Planner
from batteryswap_public.utils import load_dataset, iterate_scenarios
from batteryswap_public.evaluate import evaluate_plan, check_plan_valid

log = structlog.get_logger()

# Origins whose 42-day window closes before s_24 are the training block.
TEST_ORIGINS = range(24, 42)

# Temperature band for trusting a voltage reading, as in utils.smooth_series
MIN_TEMP = 10
MAX_TEMP = 30


def infer_start_time(battery_data: pandas.DataFrame) -> pandas.Timestamp:
    """
    Recover the planning start from the truncated timeseries.

    The Planner interface does not receive start_time - iterate_scenarios cuts the
    series at it, so the last observation is the best available estimate.
    """
    return battery_data.reset_index()['end_time'].max().normalize()


def voltage_feature(battery_data: pandas.DataFrame) -> pandas.Series:
    """
    Last daily-median voltage per device, using only in-band readings.

    Daily median rejects within-day noise; the temperature band rejects the part of
    the voltage swing that reflects ambient rather than state-of-charge. No recency
    limit - a stale low reading still predicts EOL far better than no reading, and
    the 95th percentile of staleness is ~148 days.

    Deliberately uses no fleet-wide statistic, so slicing history at an origin can
    never leak information from after it.
    """
    df = battery_data.reset_index()
    df = df[(df['temperature'] > MIN_TEMP) & (df['temperature'] < MAX_TEMP)]
    if df.empty:
        return pandas.Series(dtype=float, name='voltage')

    df['day'] = df['end_time'].dt.normalize()
    last_day = df.groupby('device_id')['day'].transform('max')
    latest = df[df['day'] == last_day]
    return latest.groupby('device_id')['voltage'].median().rename('voltage')


def _assemble(selected_days: pandas.Series, locations, start_time, window_days):
    """
    Build a valid plan: scheduled batteries on their days, everything else parked
    one day past the horizon where evaluate_plan discards it for free.
    """
    parked_day = start_time + pandas.Timedelta(days=window_days + 1)
    scheduled = set(selected_days.index)
    parked = [b for b in locations['battery'] if b not in scheduled]

    plan = pandas.concat([
        pandas.DataFrame({'day': selected_days.values, 'battery': selected_days.index}),
        pandas.DataFrame({'day': parked_day, 'battery': parked}),
    ], ignore_index=True)
    plan = plan.sort_values('day', kind='stable').reset_index(drop=True)

    check_plan_valid(plan, locations, start_time=start_time)
    return plan


class DeferAllPlanner(Planner):
    """Park everything. Pays only forced emergency visits."""

    def plan(self, battery_data, locations, travel_costs, settings):
        start_time = infer_start_time(battery_data)
        empty = pandas.Series(dtype='datetime64[ns]')
        return _assemble(empty, locations, start_time, settings.planning_window_days)


class LowestVoltagePlanner(Planner):
    """
    Swap the k batteries with the lowest last in-band voltage, worst first, one per
    day from the start of the window.

    k is deliberately allowed to exceed the expected number of deaths: missing a
    battery that dies costs an emergency visit, while including a healthy one costs
    only early_replacement_penalty_daily per day.
    """

    def __init__(self, k: int = 10, feature_cache: Optional[dict] = None):
        self.k = k
        self.feature_cache = feature_cache if feature_cache is not None else {}

    def _voltage(self, battery_data, start_time):
        key = pandas.Timestamp(start_time)
        if key not in self.feature_cache:
            self.feature_cache[key] = voltage_feature(battery_data)
        return self.feature_cache[key]

    def plan(self, battery_data, locations, travel_costs, settings):
        start_time = infer_start_time(battery_data)
        voltage = self._voltage(battery_data, start_time)

        # Restrict to batteries still in the fleet; unmeasured ones sort last.
        v = voltage.reindex(locations['battery']).fillna(numpy.inf)
        worst = v.nsmallest(self.k, keep='first')

        days = start_time + pandas.to_timedelta(numpy.arange(len(worst)), unit='D')
        selected = pandas.Series(days, index=worst.index)
        return _assemble(selected, locations, start_time, settings.planning_window_days)


class OraclePlanner(Planner):
    """Upper bound: knows the true EOL and swaps on exactly that day."""

    def __init__(self, eol_times: pandas.Series):
        self.eol_times = eol_times

    def plan(self, battery_data, locations, travel_costs, settings):
        start_time = infer_start_time(battery_data)
        window = pandas.Timedelta(days=settings.planning_window_days)
        eol = self.eol_times.reindex(locations['battery']).dropna()
        doomed = eol[(eol > start_time) & (eol <= start_time + window)].sort_values()
        return _assemble(doomed, locations, start_time, settings.planning_window_days)


def check_feature_parity(battery_data, cache_key, feature_cache):
    """
    Guard against train/serve skew: the harness prefills the same cache the planner
    reads, so assert the cached value matches a fresh online computation.
    """
    cached = feature_cache[cache_key]
    online = voltage_feature(battery_data)
    pandas.testing.assert_series_equal(cached.sort_index(), online.sort_index())


def load_shipped_planner(path='batteryswap_example/planners/best.pickle'):
    """
    Load planners/best.pickle. The current pickle (written by planner.py's
    main()) records its classes under importable module paths, but a pickle
    written by running train.py as a script records OrderedPlanner /
    DummyRULModel under __main__ - keep registering them there so such
    pickles load too.
    """
    import pickle
    import __main__
    from batteryswap_example import train as train_module
    for name in ('OrderedPlanner', 'DummyRULModel'):
        setattr(__main__, name, getattr(train_module, name))
    with open(path, 'rb') as f:
        return pickle.load(f)


def run(dataset_path: Path, origins=TEST_ORIGINS, k_values=(5, 10, 15, 20, 30)):
    locations, timeseries, eol_times, scenarios = load_dataset(dataset_path / 'train')
    log.info('harness-loaded', devices=len(locations), origins=list(origins))

    shipped = load_shipped_planner()

    feature_cache: dict = {}
    policies = {'current': shipped, 'defer': DeferAllPlanner(),
                'oracle': OraclePlanner(eol_times)}
    for k in k_values:
        policies[f'voltage-{k}'] = LowestVoltagePlanner(k, feature_cache=feature_cache)

    records = []
    parity_checked = False
    gen = iterate_scenarios(locations, timeseries, eol_times, scenarios)
    for index, (scenario, locs, cut, eol) in enumerate(gen):
        if index not in origins:
            continue
        start_time = pandas.Timestamp(scenario['start_time'])
        inferred = infer_start_time(cut)
        assert inferred == start_time, (scenario['name'], inferred, start_time)

        for name, planner in policies.items():
            plan = planner.plan(cut, locs, scenario['travel_costs'], scenario['settings'])
            _, _, overall = evaluate_plan(
                plan, locs, scenario['travel_costs'], scenario['settings'],
                eol_times=eol, start_time=start_time, verbose=0)
            records.append(dict(overall, policy=name, origin=scenario['name']))

        if not parity_checked:
            check_feature_parity(cut, start_time, feature_cache)
            parity_checked = True
            log.info('feature-parity-ok', origin=scenario['name'])

        log.info('origin-done', origin=scenario['name'])

    return pandas.DataFrame(records)


def report(df: pandas.DataFrame):
    order = ['total_cost', 'late_swap', 'early_swap', 'daily_limit', 'weekly_limit',
             'overtime', 'travel', 'battery_swap', 'building_change', 'room_change']
    mean = df.groupby('policy')[order].mean().sort_values('total_cost')

    defer = mean.loc['defer', 'total_cost']
    per_origin = df.pivot(index='origin', columns='policy', values='total_cost')
    beats = (per_origin.lt(per_origin['defer'], axis=0)).sum()
    spread = per_origin.std()

    summary = mean.copy()
    summary['vs_defer_%'] = (100 * (mean['total_cost'] / defer - 1)).round(1)
    summary['origins_beating_defer'] = beats.reindex(summary.index).astype(str) + f'/{len(per_origin)}'
    summary['std_across_origins'] = spread.reindex(summary.index).round(1)

    print(f'\nForward split: scored on {len(per_origin)} origins '
          f'({per_origin.index[0]}..{per_origin.index[-1]})\n')
    print(summary.round(2).to_string())
    return summary


class Config(BaseSettings):
    model_config = SettingsConfigDict(env_prefix='', cli_parse_args=True,
                                      cli_ignore_unknown_args=True)
    dataset_path: Optional[Path] = None


def main():
    cfg = Config()
    dataset_path = cfg.dataset_path or Path(
        os.environ.get('BATTERYSWAP_DATASET_PATH', 'dataset'))
    df = run(Path(dataset_path))
    report(df)


if __name__ == '__main__':
    main()