"""Synthetic fleet generator for the SCALING evaluation.

Fleet at scale k = k time-shifted clones of each real building. Each clone
building gets ONE shift (preserves within-building death correlation,
decorrelates clone sets). Ground-truth EOL is DERIVED from the shifted
telemetry with the official mechanism (7d-median < 2.40, validated 81/81
by the team repo) - labels stay mechanically consistent with the data.
Travel: distance between clones = distance between base buildings;
same-base clone pairs get the median off-diagonal hop.

Writes /tmp/c5/fleet_k{K}.pkl with: timeseries, locations, eol, travel fn.
"""
import sys, pickle, time
from pathlib import Path
import numpy as np
import pandas as pd

TEAM = Path('/tmp/teamrepo')
OURS = Path('/Users/THIF/repos/BatterySwapAI2026-Example')
sys.path.insert(0, str(TEAM))
sys.path.insert(0, str(OURS))
from batteryswap_public.utils import load_dataset
from features import daily_panel, eol_from_panel

K = int(sys.argv[1]) if len(sys.argv) > 1 else 2
SEED = 7

locations, ts, eol_times, scenarios = load_dataset(OURS / 'dataset/train')
df = ts.reset_index()
df['device_id'] = df['device_id'].astype(str)
dev = pd.read_csv(OURS / 'dataset/train/devices.csv').drop_duplicates('device_id')
span_lo = df['end_time'].min()
span_hi = df['end_time'].max()

rng = np.random.default_rng(SEED)
frames, loc_rows = [], []
for c in range(K):
    for b, sub in dev.groupby('building_id'):
        # clone 0 of each building = the untouched original (scale anchor)
        shift = 0 if c == 0 else int(rng.integers(-120, 121))
        sb = b if c == 0 else f'{b}_c{c}'
        for _, r in sub.iterrows():
            sd = r['device_id'] if c == 0 else f"{r['device_id']}_c{c}"
            d = df[df['device_id'] == r['device_id']].copy()
            d['device_id'] = sd
            d['end_time'] = d['end_time'] + pd.Timedelta(days=shift)
            d = d[(d['end_time'] >= span_lo) & (d['end_time'] <= span_hi)]
            if len(d) < 50:
                continue
            frames.append(d)
            loc_rows.append(dict(battery=sd, building=sb,
                                 room=(r['room_id'] if c == 0
                                       else f"{r['room_id']}_c{c}"),
                                 base_building=b,
                                 start_time=pd.Timestamp(r['start_time'])
                                 + pd.Timedelta(days=shift),
                                 end_time=pd.Timestamp(r['end_time'])
                                 + pd.Timedelta(days=shift)))

sts = pd.concat(frames, ignore_index=True)
locs = pd.DataFrame(loc_rows)
for col in ('start_time', 'end_time'):
    locs[col] = pd.to_datetime(locs[col], utc=True)

t0 = time.time()
panel = daily_panel(sts.set_index(['device_id', 'end_time']))
eol = eol_from_panel(panel)
print(f'k={K}: {locs.battery.nunique()} devices, '
      f'{locs.building.nunique()} buildings, {len(eol)} EOLs '
      f'(panel {time.time()-t0:.0f}s)')

# sanity: clone-0 EOLs must match the official file exactly
official = pd.read_csv(OURS / 'dataset/train/eol_times.csv')
col = [c for c in official.columns if c != 'device_id'][0]
truth = pd.to_datetime(official.set_index('device_id')[col]).dt.tz_localize(None)
c0 = eol[~eol.index.str.contains('_c')]
both = truth.dropna().index.intersection(c0.index)
mism = int((pd.to_datetime(c0.reindex(both)).dt.normalize()
            != truth.reindex(both).dt.normalize()).sum())
print(f'clone-0 EOL check: {len(both)} deaths, {mism} mismatches')

# travel-cost template: base 24x24 matrix from scenario 0
tc0 = scenarios[0]['travel_costs']
base_h = tc0.set_index(['from', 'to'])['hours']
med_hop = tc0.loc[tc0['from'] != tc0['to'], 'hours'].median()
blds = sorted(locs['building'].unique())
base_of = locs.drop_duplicates('building').set_index('building')['base_building']
rows = []
for i in blds:
    for j in blds:
        if i == j:
            h = 0.0
        else:
            bi, bj = base_of[i], base_of[j]
            h = med_hop if bi == bj else float(base_h.get((bi, bj), med_hop))
        rows.append(dict(**{'from': i, 'to': j}, hours=h))
travel = pd.DataFrame(rows)

with open(f'/tmp/c5/fleet_k{K}.pkl', 'wb') as f:
    pickle.dump(dict(ts=sts, locations=locs, eol=pd.to_datetime(eol),
                     travel=travel), f)
print(f'wrote /tmp/c5/fleet_k{K}.pkl')