#!/usr/bin/env python3
"""Reproducible SMARD Germany hourly-data study of counterfactual demand flexibility.

Downloads public SMARD hourly series, validates timestamp coverage and official
residual-load identity, then tests same-day demand shifting under explicit caps.
This is a sensitivity analysis, not a forecast, dispatch plan, or measured DR result.
Python standard library only.
"""
from __future__ import annotations
import argparse, csv, datetime as dt, hashlib, json, math, statistics, time, urllib.error, urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from zoneinfo import ZoneInfo

BASE_URL = "https://www.smard.de/app/chart_data"
REGION = "DE"
TZ_NAME = "Europe/Berlin"
TZ = ZoneInfo(TZ_NAME)
USER_AGENT = "Darkoza-OERSI-GridFlex-Research/1.0 (reproducible public-data study)"
FILTERS = {
    "load": 410,
    "residual_load": 4359,
    "wind_onshore": 4067,
    "wind_offshore": 1225,
    "solar": 4068,
}
START_LOCAL = "2025-11-02T21:00:00"
END_LOCAL = "2026-11-02T21:00:00"  # exclusive; 365 local days

def fetch_json(url: str, attempts: int = 5) -> dict:
    last = None
    for attempt in range(attempts):
        try:
            req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT, "Accept": "application/json"})
            with urllib.request.urlopen(req, timeout=30) as response:
                if response.status != 200:
                    raise RuntimeError(f"HTTP {response.status}: {url}")
                return json.loads(response.read().decode("utf-8"))
        except Exception as exc:
            last = exc
            time.sleep(min(2.0, 0.25 * (attempt + 1)))
    raise RuntimeError(f"SMARD request failed after {attempts} attempts: {url}: {last}")

def epoch_ms(local_iso: str) -> int:
    value = dt.datetime.fromisoformat(local_iso).replace(tzinfo=TZ)
    return int(value.timestamp() * 1000)

def download_series(filter_id: int, start_ms: int, end_ms: int) -> dict[int, float | None]:
    index_url = f"{BASE_URL}/{filter_id}/{REGION}/index_hour.json"
    index = fetch_json(index_url)
    stamps = index.get("timestamps")
    if not isinstance(stamps, list) or not stamps:
        raise RuntimeError(f"Missing SMARD timestamps for filter {filter_id}")
    prior = [int(t) for t in stamps if int(t) <= start_ms]
    if not prior:
        raise RuntimeError(f"No SMARD index block at/before start for filter {filter_id}")
    blocks = sorted(set([max(prior)] + [int(t) for t in stamps if start_ms < int(t) < end_ms]))

    def fetch_block(week_start: int) -> tuple[int, list]:
        url = f"{BASE_URL}/{filter_id}/{REGION}/{filter_id}_{REGION}_hour_{week_start}.json"
        payload = fetch_json(url)
        series = payload.get("series")
        if not isinstance(series, list):
            raise RuntimeError(f"Missing series for filter {filter_id}, block {week_start}")
        return week_start, series

    points: dict[int, float | None] = {}
    # Bounded concurrency keeps the one-year public-data retrieval practical.
    with ThreadPoolExecutor(max_workers=8) as pool:
        for _, series in pool.map(fetch_block, blocks):
            for row in series:
                if not isinstance(row, list) or len(row) < 2:
                    continue
                ts = int(row[0])
                if start_ms <= ts < end_ms:
                    value = row[1]
                    points[ts] = None if value is None else float(value)
    return points

def percentile(values: list[float], q: float) -> float:
    if not values:
        return float("nan")
    data = sorted(values)
    pos = (len(data) - 1) * q
    lo = int(math.floor(pos)); hi = int(math.ceil(pos))
    if lo == hi:
        return data[lo]
    return data[lo] + (data[hi] - data[lo]) * (pos - lo)

def metrics(values: list[float]) -> dict:
    return {
        "hours": len(values),
        "mean_MW": statistics.fmean(values) if values else None,
        "peak_MW": max(values) if values else None,
        "p95_MW": percentile(values, 0.95) if values else None,
        "stddev_MW": statistics.pstdev(values) if values else None,
        "minimum_MW": min(values) if values else None,
        "negative_residual_hours": sum(v < 0 for v in values),
    }

def shift_day_profile(values: list[float], load: list[float], timestamps_ms: list[int],
                      flex_fraction: float, window_hours: int) -> tuple[list[float], float, list[float], list[float], float]:
    """Greedy pairwise energy-neutral shift; each direct transfer obeys the time window."""
    n = len(values)
    if len(load) != n or len(timestamps_ms) != n:
        raise ValueError("profile_length_mismatch")
    if not 0 <= flex_fraction <= 0.5:
        raise ValueError("flex_fraction_out_of_range")
    if window_hours < 1:
        raise ValueError("window_hours_must_be_positive")
    y = [float(v) for v in values]
    # A single shared per-hour flexibility budget prevents an hour from
    # consuming the full cap once as a source and again as a destination.
    cap_remaining = [max(0.0, float(x)) * flex_fraction for x in load]
    moved_out = [0.0] * n
    moved_in = [0.0] * n
    total_moved = 0.0
    max_shift_used = 0.0
    max_gap_ms = window_hours * 3_600_000

    # Process shorter direct time separations first. For a larger window this
    # preserves the exact smaller-window sequence, then considers extra pairs.
    for index_gap in range(1, min(window_hours, n - 1) + 1):
        candidates = []
        for i in range(n - index_gap):
            j = i + index_gap
            gap_ms = abs(timestamps_ms[j] - timestamps_ms[i])
            if gap_ms > max_gap_ms:
                continue
            candidates.append((abs(y[i] - y[j]), i, j, gap_ms))
        candidates.sort(reverse=True)
        for _, i, j, gap_ms in candidates:
            if y[i] >= y[j]:
                source, destination = i, j
            else:
                source, destination = j, i
            difference = y[source] - y[destination]
            if difference <= 1e-10 or cap_remaining[source] <= 1e-10 or cap_remaining[destination] <= 1e-10:
                continue
            amount = min(cap_remaining[source], cap_remaining[destination], difference / 2.0)
            if amount <= 1e-10:
                continue
            y[source] -= amount
            y[destination] += amount
            cap_remaining[source] -= amount
            cap_remaining[destination] -= amount
            moved_out[source] += amount
            moved_in[destination] += amount
            total_moved += amount
            max_shift_used = max(max_shift_used, gap_ms / 3_600_000)
    return y, total_moved, moved_out, moved_in, max_shift_used

def make_days(points: list[dict]) -> tuple[list[list[dict]], list[str]]:
    groups: dict[str, list[dict]] = {}
    for point in points:
        local = dt.datetime.fromtimestamp(point["timestamp_ms"] / 1000, dt.timezone.utc).astimezone(TZ)
        groups.setdefault(local.date().isoformat(), []).append(point)
    days, excluded = [], []
    for date_text in sorted(groups):
        group = sorted(groups[date_text], key=lambda p: p["timestamp_ms"])
        day = dt.date.fromisoformat(date_text)
        start = dt.datetime.combine(day, dt.time(0), TZ)
        end = dt.datetime.combine(day + dt.timedelta(days=1), dt.time(0), TZ)
        expected = int((end.timestamp() - start.timestamp()) // 3600)
        if len(group) != expected or any(p["load_MW"] is None or p["residual_MW"] is None for p in group):
            excluded.append(date_text)
            continue
        expected_ts = [int((start.timestamp() + h * 3600) * 1000) for h in range(expected)]
        if [p["timestamp_ms"] for p in group] != expected_ts:
            excluded.append(date_text)
            continue
        days.append(group)
    return days, excluded

def run(output_dir: Path) -> dict:
    start_ms, end_ms = epoch_ms(START_LOCAL), epoch_ms(END_LOCAL)
    expected_hours = int((end_ms - start_ms) // 3_600_000)
    if expected_hours != 8760:
        raise RuntimeError(f"Expected a 365-day period / 8760 hours, got {expected_hours}")
    all_series = {}
    for name, filter_id in FILTERS.items():
        all_series[name] = download_series(filter_id, start_ms, end_ms)

    points = []
    for ts in range(start_ms, end_ms, 3_600_000):
        row = {"timestamp_ms": ts}
        for name in FILTERS:
            row[name] = all_series[name].get(ts)
        points.append({
            "timestamp_ms": ts,
            "load_MW": row["load"],
            "residual_MW": row["residual_load"],
            "wind_onshore_MW": row["wind_onshore"],
            "wind_offshore_MW": row["wind_offshore"],
            "solar_MW": row["solar"],
        })

    days, excluded_days = make_days(points)
    valid_points = [p for day in days for p in day]
    if not valid_points:
        raise RuntimeError("No complete local days available")
    coverage = len(valid_points) / expected_hours
    if coverage < 0.995 or len(excluded_days) > 2:
        raise RuntimeError(f"Data coverage gate failed: complete_hours={len(valid_points)}/{expected_hours}; excluded_days={len(excluded_days)}")

    load = [p["load_MW"] for p in valid_points]
    residual = [p["residual_MW"] for p in valid_points]
    component_errors = []
    component_count = 0
    for p in valid_points:
        comps = (p["wind_onshore_MW"], p["wind_offshore_MW"], p["solar_MW"])
        if all(v is not None for v in comps):
            reconstructed = p["load_MW"] - sum(comps)
            component_errors.append(abs(reconstructed - p["residual_MW"]))
            component_count += 1
    crosscheck = {
        "eligible_hours": component_count,
        "mean_absolute_difference_MW": statistics.fmean(component_errors) if component_errors else None,
        "p95_absolute_difference_MW": percentile(component_errors, .95) if component_errors else None,
        "max_absolute_difference_MW": max(component_errors) if component_errors else None,
        "interpretation": "Independent arithmetic cross-check against SMARD's published residual-load series; small differences may arise from rounding/aggregation."
    }

    scenario_results = []
    for flex in (0.025, 0.05, 0.10):
        for window in (2, 4, 6):
            shifted_all, total_shift = [], 0.0
            max_shift_used = 0.0
            daily_balance_errors = []
            for day in days:
                values = [p["residual_MW"] for p in day]
                day_load = [p["load_MW"] for p in day]
                stamps = [p["timestamp_ms"] for p in day]
                shifted, moved, _, _, day_max_shift = shift_day_profile(values, day_load, stamps, flex, window)
                balance_error = abs(sum(values) - sum(shifted))
                daily_balance_errors.append(balance_error)
                if balance_error > 1e-6:
                    raise RuntimeError(f"Daily energy conservation failed: {balance_error} MWh")
                shifted_all.extend(shifted)
                total_shift += moved
                max_shift_used = max(max_shift_used, day_max_shift)
            max_daily_balance_error = max(daily_balance_errors, default=0.0)
            base_metrics = metrics(residual)
            shifted_metrics = metrics(shifted_all)
            scenario_results.append({
                "flexible_load_cap_fraction": flex,
                "max_shift_window_hours": window,
                "baseline": base_metrics,
                "counterfactual": shifted_metrics,
                "peak_change_percent": 100 * (shifted_metrics["peak_MW"] - base_metrics["peak_MW"]) / abs(base_metrics["peak_MW"]),
                "p95_change_percent": 100 * (shifted_metrics["p95_MW"] - base_metrics["p95_MW"]) / abs(base_metrics["p95_MW"]),
                "one_way_shifted_energy_MWh": total_shift,
                "max_observed_direct_shift_hours": max_shift_used,
                "max_daily_energy_balance_error_MWh": max_daily_balance_error,
                "share_of_observed_annual_load_percent": 100 * total_shift / sum(load),
                "daily_energy_conserved": max_daily_balance_error <= 1e-6
            })

    output_dir.mkdir(parents=True, exist_ok=True)
    csv_path = output_dir / "smard_hourly_2025-11-02_to_2026-11-02.csv"
    with csv_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["timestamp_utc", "timestamp_local", "load_hourly_value", "residual_load_hourly_value",
                         "wind_onshore_hourly_value", "wind_offshore_hourly_value", "solar_hourly_value"])
        for p in points:
            utc = dt.datetime.fromtimestamp(p["timestamp_ms"] / 1000, dt.timezone.utc)
            local = utc.astimezone(TZ)
            writer.writerow([utc.isoformat(), local.isoformat(), p["load_MW"], p["residual_MW"],
                             p["wind_onshore_MW"], p["wind_offshore_MW"], p["solar_MW"]])

    result = {
        "study_id": "gridflex-germany-flexibility-2026-11-03",
        "title": "Germany's hourly load and residual-load flexibility: a one-year SMARD counterfactual",
        "created_date": "2026-11-03",
        "period": {"start_local": START_LOCAL + " Europe/Berlin", "end_exclusive_local": END_LOCAL + " Europe/Berlin",
                   "expected_hours": expected_hours, "complete_hours_used": len(valid_points),
                   "complete_local_days_used": len(days), "excluded_local_days": excluded_days},
        "sources": {
            "publisher": "Bundesnetzagentur / SMARD",
            "portal": "https://www.smard.de/",
            "data_download_guidance": "https://www.smard.de/en/all-about-our-data-download-section-210130",
            "hourly_chart_api": "https://www.smard.de/app/chart_data/{filter}/DE/index_hour.json and weekly JSON blocks",
            "license": "Creative Commons Attribution 4.0 International (CC BY 4.0), per SMARD data-use notice",
            "filters": FILTERS,
            "resolution": "hourly",
            "region": "Germany (DE)",
            "timezone_for_daily_grouping": TZ_NAME,
            "attribution": "Bundesnetzagentur | SMARD.de. Hourly electricity-market data for Germany; hourly series joined and used in a counterfactual analysis. License: CC BY 4.0."
        },
        "data_quality": {
            "coverage_fraction": coverage,
            "expected_hourly_rows": expected_hours,
            "valid_complete_day_hours": len(valid_points),
            "unique_timestamps": len({p["timestamp_ms"] for p in points}),
            "residual_identity_crosscheck": crosscheck,
            "missing_values_by_series": {name: sum(all_series[name].get(ts) is None for ts in range(start_ms, end_ms, 3_600_000)) for name in FILTERS}
        },
        "observed_baseline": {
            "network_load": metrics(load),
            "published_residual_load": metrics(residual),
            "description": "Residual load is SMARD net load minus wind onshore, wind offshore and photovoltaic feed-in; it is a system-level proxy for the remaining demand not met by those variable renewables."
        },
        "counterfactual_method": {
            "type": "deterministic pairwise, same-local-day demand-shifting sensitivity; shorter direct time gaps processed first",
            "assumed_flexibility_scenarios": [0.025, 0.05, 0.10],
            "max_shift_windows_hours": [2, 4, 6],
            "interpretation": "At each hour, the combined absolute energy shifted out plus shifted in is capped at the stated fraction of observed load. Within each local day, hourly pairs are processed by increasing direct time gap (then current residual-load spread); each pair is shifted at most once, with transfers bounded by the requested window and the shared per-hour flexibility budget. This is a transparent heuristic, not a globally optimal dispatch.",
            "not_modelled": ["actual flexible-device inventory or participation", "user comfort/process constraints", "market prices and dispatch payments", "network congestion, voltage, frequency and reserves", "cross-border flows or unit commitment response", "storage, generation dispatch, ramp rates and forecast error", "emissions or fuel costs", "industrial/household rebound beyond the same-day constraint"],
            "not_a_claim": "The results are mathematical counterfactual sensitivity bounds under assumed flexibility, not measured demand-response availability, realized savings, avoided emissions, an operational dispatch recommendation, or proof that all shifted load is technically/economically feasible."
        },
        "scenarios": scenario_results,
        "reproducibility": {
            "python": "standard library only",
            "script": "run_study.py",
            "tests": "test_run_study.py",
            "raw_data_file": csv_path.name
        }
    }
    json_path = output_dir / "results.json"
    result["reproducibility"]["raw_data_sha256"] = hashlib.sha256(csv_path.read_bytes()).hexdigest()
    result["reproducibility"]["results_sha256_file"] = "results.json.sha256"
    json_path.write_text(json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False) + "\n", encoding="utf-8")
    checksum = hashlib.sha256(json_path.read_bytes()).hexdigest()
    (output_dir / "results.json.sha256").write_text(f"{checksum}  results.json\n", encoding="utf-8")
    return result

def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", default=".")
    args = parser.parse_args()
    result = run(Path(args.output_dir))
    print(json.dumps({
        "study_id": result["study_id"],
        "coverage_fraction": result["data_quality"]["coverage_fraction"],
        "complete_hours": result["period"]["complete_hours_used"],
        "excluded_days": len(result["period"]["excluded_local_days"]),
        "crosscheck": result["data_quality"]["residual_identity_crosscheck"],
        "baseline": result["observed_baseline"]["published_residual_load"],
        "scenarios": result["scenarios"]
    }, indent=2, ensure_ascii=False))
    return 0

if __name__ == "__main__":
    raise SystemExit(main())