#!/usr/bin/env python3
"""Transparent hourly renewable-plus-storage screening model.

This is a screening calculator, not a grid dispatch optimizer or engineering design.
Inputs are one-hour rows with load_mw, solar_mw, wind_mw and optional timestamp.
Storage capacities are internal usable MWh; power ratings are MW.
"""
from __future__ import annotations
import argparse
import csv
import json
import math
import sys
from datetime import datetime
from pathlib import Path
from typing import Any

REQUIRED = ("load_mw", "solar_mw", "wind_mw")


def load_profile(path: str) -> list[dict[str, float]]:
    rows: list[dict[str, float]] = []
    stamps: list[datetime] = []
    with open(path, newline="", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        if reader.fieldnames is None or not set(REQUIRED).issubset(reader.fieldnames):
            raise ValueError("CSV must contain load_mw, solar_mw and wind_mw columns")
        for number, raw in enumerate(reader, start=2):
            item: dict[str, float] = {}
            for key in REQUIRED:
                try:
                    value = float(raw[key])
                except (TypeError, ValueError) as exc:
                    raise ValueError(f"Invalid {key} at CSV line {number}") from exc
                if not math.isfinite(value) or value < 0:
                    raise ValueError(f"{key} must be finite and non-negative at line {number}")
                item[key] = value
            rows.append(item)
            stamp = raw.get("timestamp", "")
            if stamp:
                try:
                    stamps.append(datetime.fromisoformat(stamp.replace("Z", "+00:00")))
                except ValueError as exc:
                    raise ValueError(f"Invalid ISO timestamp at CSV line {number}") from exc
    if not rows:
        raise ValueError("CSV contains no hourly rows")
    if stamps and len(stamps) != len(rows):
        raise ValueError("Either every row must have timestamp or none may have it")
    if stamps:
        for a, b in zip(stamps, stamps[1:]):
            if (b - a).total_seconds() != 3600:
                raise ValueError("Timestamps must be strictly consecutive one-hour intervals")
    return rows


def _validate_asset(power: float, energy: float, rte: float, name: str) -> None:
    if power < 0 or energy < 0:
        raise ValueError(f"{name} power and energy capacities must be non-negative")
    if not (0 < rte <= 1):
        raise ValueError(f"{name} round-trip efficiency must be in (0, 1]")


def _asset(power: float, energy: float, rte: float, initial_soc: float) -> dict[str, float]:
    eta = math.sqrt(rte)
    return {
        "power": power, "capacity": energy, "eta_charge": eta,
        "eta_discharge": eta, "soc": energy * initial_soc,
        "charge_input": 0.0, "discharge_output": 0.0,
        "charge_loss": 0.0, "discharge_loss": 0.0,
    }


def simulate(
    rows: list[dict[str, float]],
    battery_power_mw: float = 0.0,
    battery_energy_mwh: float = 0.0,
    battery_rte: float = 0.90,
    ldes_power_mw: float = 0.0,
    ldes_energy_mwh: float = 0.0,
    ldes_rte: float = 0.60,
    initial_soc_fraction: float = 0.0,
    charge_priority: str = "battery-first",
    discharge_priority: str = "battery-first",
) -> dict[str, Any]:
    """Simulate a 1-hour renewable-only island balance with two storage assets."""
    _validate_asset(battery_power_mw, battery_energy_mwh, battery_rte, "Battery")
    _validate_asset(ldes_power_mw, ldes_energy_mwh, ldes_rte, "LDES")
    if not 0 <= initial_soc_fraction <= 1:
        raise ValueError("initial_soc_fraction must be between 0 and 1")
    if charge_priority not in ("battery-first", "ldes-first"):
        raise ValueError("charge_priority must be battery-first or ldes-first")
    if discharge_priority not in ("battery-first", "ldes-first"):
        raise ValueError("discharge_priority must be battery-first or ldes-first")

    assets = {
        "battery": _asset(battery_power_mw, battery_energy_mwh, battery_rte, initial_soc_fraction),
        "ldes": _asset(ldes_power_mw, ldes_energy_mwh, ldes_rte, initial_soc_fraction),
    }
    charge_order = [charge_priority.split("-")[0], "ldes" if charge_priority == "battery-first" else "battery"]
    discharge_order = [discharge_priority.split("-")[0], "ldes" if discharge_priority == "battery-first" else "battery"]
    total_load = total_renewable = direct_renewable = curtailment = unserved = 0.0
    baseline_unserved = baseline_curtailment = 0.0
    storage_delivered = 0.0
    hours_with_unserved = hours_with_curtailment = 0

    for row in rows:
        load = row["load_mw"]
        renewable = row["solar_mw"] + row["wind_mw"]
        direct = min(load, renewable)
        deficit = load - direct
        surplus = renewable - direct
        total_load += load
        total_renewable += renewable
        direct_renewable += direct
        baseline_unserved += deficit
        baseline_curtailment += surplus

        # Charge only from renewable surplus. No grid import or fossil backup is assumed.
        for name in charge_order:
            a = assets[name]
            if surplus <= 1e-12 or a["power"] <= 0 or a["capacity"] <= 0:
                continue
            max_by_power = a["power"]
            max_by_room = (a["capacity"] - a["soc"]) / a["eta_charge"]
            input_mwh = min(surplus, max_by_power, max(0.0, max_by_room))
            if input_mwh > 0:
                a["soc"] += input_mwh * a["eta_charge"]
                a["charge_input"] += input_mwh
                a["charge_loss"] += input_mwh * (1.0 - a["eta_charge"])
                surplus -= input_mwh

        if surplus > 1e-12:
            curtailment += surplus
            hours_with_curtailment += 1

        # Discharge only into the remaining load deficit.
        for name in discharge_order:
            a = assets[name]
            if deficit <= 1e-12 or a["power"] <= 0 or a["capacity"] <= 0:
                continue
            output_mwh = min(deficit, a["power"], a["soc"] * a["eta_discharge"])
            if output_mwh > 0:
                withdrawn = output_mwh / a["eta_discharge"]
                a["soc"] -= withdrawn
                a["discharge_output"] += output_mwh
                a["discharge_loss"] += withdrawn - output_mwh
                deficit -= output_mwh
                storage_delivered += output_mwh

        if deficit > 1e-12:
            unserved += deficit
            hours_with_unserved += 1

    def asset_result(a: dict[str, float]) -> dict[str, float]:
        return {
            "power_mw": a["power"],
            "energy_capacity_mwh": a["capacity"],
            "charge_input_mwh": a["charge_input"],
            "discharge_output_mwh": a["discharge_output"],
            "charge_loss_mwh": a["charge_loss"],
            "discharge_loss_mwh": a["discharge_loss"],
            "total_conversion_loss_mwh": a["charge_loss"] + a["discharge_loss"],
            "end_soc_mwh": a["soc"],
            "end_soc_fraction": a["soc"] / a["capacity"] if a["capacity"] else 0.0,
            "equivalent_full_discharge_cycles": (
                a["discharge_output"] / a["capacity"] if a["capacity"] else 0.0
            ),
        }

    served = total_load - unserved
    return {
        "model": "Darkoza GridFlex Screen v0.1",
        "time_step_hours": 1,
        "hours_simulated": len(rows),
        "assumptions": {
            "renewable_sources": ["solar_mw", "wind_mw"],
            "grid_import": "not modelled",
            "fossil_backup": "not modelled",
            "storage_dispatch": "greedy; user-selected charge and discharge priority",
            "round_trip_efficiency_split": "charge and discharge efficiencies are both sqrt(round-trip efficiency)",
            "initial_state_of_charge_fraction": initial_soc_fraction,
        },
        "energy_balance_mwh": {
            "load": total_load,
            "renewable_generation": total_renewable,
            "direct_renewable_to_load": direct_renewable,
            "storage_discharge_to_load": storage_delivered,
            "unserved_load": unserved,
            "curtailed_renewable": curtailment,
            "baseline_unserved_without_storage": baseline_unserved,
            "baseline_curtailment_without_storage": baseline_curtailment,
            "storage_conversion_losses": sum(a["charge_loss"] + a["discharge_loss"] for a in assets.values()),
        },
        "service_metrics": {
            "load_served_fraction": served / total_load if total_load else 1.0,
            "unserved_load_reduction_vs_no_storage_fraction": (
                (baseline_unserved - unserved) / baseline_unserved if baseline_unserved else 0.0
            ),
            "curtailment_reduction_vs_no_storage_fraction": (
                (baseline_curtailment - curtailment) / baseline_curtailment if baseline_curtailment else 0.0
            ),
            "hours_with_unserved_load": hours_with_unserved,
            "hours_with_curtailment": hours_with_curtailment,
        },
        "storage": {
            "battery": asset_result(assets["battery"]),
            "long_duration": asset_result(assets["ldes"]),
        },
        "limitations": [
            "Screening model only; not a least-cost optimizer, grid study or engineering design.",
            "One-hour resolution; ignores sub-hour ramps, reserves, voltage, frequency, transmission and interconnection constraints.",
            "No grid imports, fossil generation, hydropower, nuclear, demand response, exports or sector coupling.",
            "No degradation, temperature, self-discharge, maintenance, start-up, fuel supply, emissions, water, land or project financing model.",
            "Storage capacities are treated as generic energy buffers; the model does not represent rSOC, hydrogen, methane, CAES, pumped hydro or thermal process physics.",
            "Results depend on the user input time series, initial state of charge and dispatch priority. End-of-period stored energy is reported and is not forced to zero.",
            "No real-world or scientific performance claim may be inferred from the synthetic example."
        ],
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("csv", help="Hourly CSV with load_mw, solar_mw, wind_mw; optional ISO timestamp")
    parser.add_argument("--battery-power-mw", type=float, default=0.0)
    parser.add_argument("--battery-energy-mwh", type=float, default=0.0)
    parser.add_argument("--battery-rte", type=float, default=0.90)
    parser.add_argument("--ldes-power-mw", type=float, default=0.0)
    parser.add_argument("--ldes-energy-mwh", type=float, default=0.0)
    parser.add_argument("--ldes-rte", type=float, default=0.60)
    parser.add_argument("--initial-soc", type=float, default=0.0)
    parser.add_argument("--charge-priority", choices=("battery-first", "ldes-first"), default="battery-first")
    parser.add_argument("--discharge-priority", choices=("battery-first", "ldes-first"), default="battery-first")
    parser.add_argument("--output", help="Optional JSON output path; defaults to stdout")
    args = parser.parse_args()
    try:
        result = simulate(
            load_profile(args.csv),
            args.battery_power_mw, args.battery_energy_mwh, args.battery_rte,
            args.ldes_power_mw, args.ldes_energy_mwh, args.ldes_rte,
            args.initial_soc, args.charge_priority, args.discharge_priority,
        )
    except (OSError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2
    rendered = json.dumps(result, indent=2, ensure_ascii=False) + "\n"
    if args.output:
        Path(args.output).write_text(rendered, encoding="utf-8")
    else:
        sys.stdout.write(rendered)
    return 0


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