#!/usr/bin/env python3
# Eclipse SUMO, Simulation of Urban MObility; see https://eclipse.dev/sumo
# Copyright (C) 2008-2026 German Aerospace Center (DLR) and others.
# This program and the accompanying materials are made available under the
# terms of the Eclipse Public License 2.0 which is available at
# https://www.eclipse.org/legal/epl-2.0/
# This Source Code may also be made available under the following Secondary
# Licenses when the conditions for such availability set forth in the Eclipse
# Public License 2.0 are satisfied: GNU General Public License, version 2
# or later which is available at
# https://www.gnu.org/licenses/old-licenses/gpl-2.0-standalone.html
# SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-or-later

# @file    runner.py
# @author  Michael Behrisch
# @date    2026-03-04

"""
Test that SUMO parquet and CSV output match XML output for the main output types.

For each output type the script:
  1. Runs SUMO once writing that output as XML.
  2. Runs SUMO again writing it as parquet.
  3. Runs SUMO another time writing it as CSV.
  4. Converts the XML file to CSV with xml2csv.py.
  5. Compare everything numerically by loading it into pandas data frames.

Parquet may expose columns that the XML format omits (e.g. vehicle_edge in
FCD output).  Such extra parquet columns are reported but not counted as
failures.  Missing columns (XML has data that parquet drops) are failures.

Detector outputs (e1, e1instant) use tools/output/generate*Detectors.py to
produce the additional file pointing to the output path; SUMO is then run
three times with each additional file (XML, parquet and CSV respectively).
"""

import os
import subprocess
import sys

import pandas as pd

SUMO_HOME = os.environ.get(
    "SUMO_HOME",
    os.path.abspath(os.path.join(os.path.dirname(__file__), "..")),
)
SUMO_BIN = os.path.join(SUMO_HOME, "bin", "sumo")
XML2CSV = os.path.join(SUMO_HOME, "tools", "xml", "xml2csv.py")
TOOLS_OUTPUT = os.path.join(SUMO_HOME, "tools", "output")

# (sumo-option, file-stem, end-time, skip-columns-or-None)
# end-time of 1000 s ensures vehicles complete their trips for
# tripinfo / vehroute output.
# skip-columns lists column names that legitimately differ between runs
# (e.g. wall-clock timestamps in statistic-output) and are excluded from
# the comparison.
OUTPUTS = [
    ("fcd-output",        "fcd",        120,  {}),
    ("emission-output",   "emissions",  120,  {}),
    ("tripinfo-output",   "tripinfo",   1000, {}),
    ("summary-output",    "summary",    1000, {}),
    ("vehroute-output",   "vehroutes",  1000, {}),
    ("queue-output",      "queue",      200,  {}),
    ("lanechange-output", "lanechanges", 1000, {}),
    ("lanechange-output", "lanechanges", 1000, {"extra_args": ["--routes", "input_routes2.rou.xml"]}),
    ("edgedata-output",   "edgedata",   1000, {}),
    ("lanedata-output",   "lanedata",   1000, {}),
    # performance_clock* and rate columns are wall-clock measurements that
    # differ between the two independent SUMO runs.
    ("statistic-output",  "statistic",  1000, {"skip_cols": {
        "performance_clockBegin", "performance_clockEnd", "performance_clockDuration",
        "performance_realTimeFactor", "performance_vehicleUpdatesPerSecond",
        "performance_personUpdatesPerSecond", "performance_traciDuration",
    }}),
    ("stop-output",       "stopout",    1000, {}),
    ("stop-output",       "stopout",    1000, {"extra_args": ["--routes", "input_routes2.rou.xml"]}),
    # amitran-output writes heterogeneous sibling elements (actorConfig/vehicle/motionState)
    # which are fundamentally incompatible with the column-oriented parquet schema.
    # full-output also has mixed structure (vehicles/persons/edges/lanes/tls at different depths).
    ("person-summary-output", "personsummary", 1000, {}),
    ("link-output",       "link",       120,  {}),
    # battery-output requires the battery device to be active on all vehicles
    ("battery-output",    "battery",    120,  {"extra_args": ["--device.battery.probability", "1"]}),
    # bt-output requires btreceiver and btsender devices on all vehicles
    ("bt-output",         "bt",         120,  {"extra_args": ["--device.btreceiver.probability", "1",
                                                              "--device.btsender.probability", "1"]}),
]

# Detector outputs that need an additional file generated by a helper script.
# Each entry: (stem, generator_script, extra_generator_args, end_time)
# The generator is called twice: once with a .xml result path and once with a
# .parquet result path; SUMO is then run with each additional file in turn.
DETECTOR_OUTPUTS = [
    ("e1detector",
     os.path.join(TOOLS_OUTPUT, "generateTLSE1Detectors.py"),
     ["-f", "100"],
     500),
    ("e1instant",
     os.path.join(TOOLS_OUTPUT, "generateDetectors.py"),
     ["-t", "instantInductionLoop"],
     500),
    ("e2detector",
     os.path.join(TOOLS_OUTPUT, "generateTLSE2Detectors.py"),
     ["-f", "100", "-l", "100"],
     500),
    ("e3detector",
     os.path.join(TOOLS_OUTPUT, "generateTLSE3Detectors.py"),
     ["-f", "100", "-l", "100"],
     500),
]

# Route probe outputs — the output file is embedded in the additional file,
# so we generate the additional file directly without a helper script.
# Each entry: (stem, edge_id, period, end_time)
ROUTEPROBE_OUTPUTS = [
    ("routeprobe", "1fi", 100, 1000),
]

# Absolute tolerance for floating-point comparison.
# XML output rounds values to 2 decimal places.
FLOAT_ATOL = 0.0101


def run_sumo(outdir, opt, stem, ext, end_time, args):
    """Run SUMO for a single output type with the given file extension.

    Returns (success, output_path).
    """
    outfile = os.path.join(outdir, f"{stem}.{ext}")
    cmd = [
        SUMO_BIN,
        "-c", "sumo.sumocfg",
        "-b", "0", "-e", str(end_time),
        f"--{opt}", outfile,
    ]
    if "extra_args" in args:
        cmd += args["extra_args"]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        last_line = result.stderr.strip().splitlines()[-1] if result.stderr.strip() else "(no stderr)"
        print(f"  SUMO ({ext}) failed (rc={result.returncode}): {last_line}")
        return None
    return outfile


def generate_detector_additional(outdir, stem, ext, generator_script, extra_args):
    """Generate a detector additional file via a generator script.

    The result (detector output) file is set to ``<outdir>/<stem>.<ext>``.
    Returns ``(add_file_path, output_file_path)`` on success, ``None`` on failure.
    """
    outfile = os.path.abspath(os.path.join(outdir, f"{stem}.{ext}"))
    add_file = os.path.join(outdir, f"{stem}_{ext}.add.xml")
    cmd = [
        sys.executable, generator_script,
        "-n", "input_net.net.xml",
        "-r", outfile,
        "-o", add_file,
    ] + extra_args
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"  Generator failed: {result.stderr.strip()[:200]}")
        return None, None
    return add_file, outfile


def generate_routeprobe_additional(outdir, stem, ext, edge_id, period):
    """Generate a routeProbe additional file pointing the output to outdir/stem.ext.

    Returns (add_file_path, output_file_path).
    """
    outfile = os.path.abspath(os.path.join(outdir, f"{stem}.{ext}"))
    add_file = os.path.join(outdir, f"{stem}_{ext}.add.xml")
    with open(add_file, "w") as f:
        f.write(f'<additional>\n'
                f'    <routeProbe id="rp_{edge_id}" edge="{edge_id}"'
                f' period="{period}" file="{outfile}"/>\n'
                f'</additional>\n')
    return add_file, outfile


def run_sumo_with_additional(add_file, end_time):
    """Run SUMO using sumo.sumocfg plus an extra additional file.

    The base additional-files value from the config (input_routes.rou.xml) is
    re-stated on the command line together with the detector additional file so
    that routes are still loaded when the command-line option overrides the
    config-file value.

    Returns True on success.
    """
    cmd = [
        SUMO_BIN,
        "-c", "sumo.sumocfg",
        "-b", "0", "-e", str(end_time),
        "--additional-files", add_file,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        last_line = result.stderr.strip().splitlines()[-1] if result.stderr.strip() else "(no stderr)"
        print(f"  SUMO failed (rc={result.returncode}): {last_line}")
        return False
    return True


def xml_to_df(xml_path):
    """Convert an XML output file to a DataFrame via xml2csv.py.

    xml2csv writes <stem>.csv next to the input file.
    Returns None if the CSV is empty or missing.
    """
    csv_path = os.path.splitext(xml_path)[0] + "_xml.csv"
    result = subprocess.run(
        [sys.executable, XML2CSV, xml_path, "-o", csv_path],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        print(f"  xml2csv error: {result.stderr.strip()[:200]}")
        return None
    if not os.path.exists(csv_path) or os.path.getsize(csv_path) == 0:
        return None
    try:
        df = pd.read_csv(csv_path, sep=";", low_memory=False)
    except Exception as e:
        print(f"  Failed to read xml CSV: {e}")
        return None
    return df if not df.empty else None


def compare(df_xml, df_parquet, skip_cols=None):
    """Compare two DataFrames column by column.

    Columns listed in *skip_cols* are excluded before comparison (they are
    neither checked for presence nor compared by value).  This is used for
    columns whose values legitimately differ between two independent SUMO
    runs (e.g. wall-clock timestamps).

    Returns (ok, notes) where ok is True if all XML columns match and
    notes lists any informational messages (e.g. extra columns in parquet).
    """
    notes = []
    xml_cols = set(df_xml.columns)
    pq_cols = set(df_parquet.columns)

    if skip_cols:
        xml_cols -= skip_cols
        pq_cols -= skip_cols

    missing = xml_cols - pq_cols   # data present in XML but absent in parquet
    extra = pq_cols - xml_cols     # extra columns that parquet adds

    if missing:
        return False, [f"parquet is missing XML columns: {sorted(missing)}"]
    if extra:
        notes.append(f"parquet has extra columns not in XML: {sorted(extra)}")

    # Compare only the columns that appear in the XML output.
    cols = sorted(xml_cols)
    df_xml = df_xml[cols].copy()
    df_parquet = df_parquet[cols].copy()

    if len(df_xml) != len(df_parquet):
        return False, notes + [f"row count differs: xml={len(df_xml)}, parquet={len(df_parquet)}"]

    ok = True
    for col in cols:
        s_xml = df_xml[col]
        s_pq = df_parquet[col]
        diff_pairs = None
        try:
            pd.testing.assert_series_equal(
                s_xml.astype(float),
                s_pq.astype(float),
                check_names=False,
                rtol=0,
                atol=FLOAT_ATOL,
            )
        except (ValueError, TypeError):
            try:
                pd.testing.assert_series_equal(
                    s_xml.astype(str).replace((None, 'None', 'nan'), ''),
                    s_pq.astype(str).replace((None, 'None', 'nan'), ''),
                    check_names=False,
                )
            except AssertionError as e:
                notes.append(f"column '{col}' differs: {str(e)[:200]}")
                diff_pairs = zip(list(s_xml), list(s_pq))
        except AssertionError as e:
            notes.append(f"column '{col}' differs: {str(e)[:200]}")
            diff_pairs = zip(list(s_xml), list(s_pq))
        if diff_pairs:
            err = 0
            for i, p in enumerate(diff_pairs):
                if p[0] != p[1]:
                    notes.append(str((i, p)))
                    err += 1
                    if err == 10:
                        break
            ok = False

    return ok, notes


def check_and_compare(stem, xml_path, other_path, label, skip_cols):
    """Load the XML reference and another format output, compare them.

    *label* is used in printed output (e.g. "parquet" or "csv").
    Returns (ok, skipped).
    """
    xml_exists = os.path.exists(xml_path) and os.path.getsize(xml_path) > 0
    other_exists = os.path.exists(other_path) and os.path.getsize(other_path) > 0

    if not xml_exists and not other_exists:
        print(f"[SKIP]  {stem}: no output generated by either format")
        return True, True
    if not xml_exists:
        print(f"[FAIL]  {stem}: XML output file is empty")
        return False, False

    df_xml = xml_to_df(xml_path)

    if not other_exists:
        if df_xml is None:
            print(f"[SKIP]  {stem}: both are empty")
            return True, True
        print(f"[FAIL]  {stem}: {label} output file is empty")
        return False, False

    try:
        if label == "parquet":
            df_other = pd.read_parquet(other_path)
        else:
            df_other = pd.read_csv(other_path, sep=";", low_memory=False)
    except pd.errors.EmptyDataError:
        df_other = None
    except Exception as e:
        print(f"  Failed to read {other_path}: {e}")
        return False, False

    if df_xml is None and (df_other is None or len(df_other) == 0):
        print(f"[SKIP]  {stem}: both are empty")
        return True, True
    if df_xml is None:
        print(f"[FAIL]  {stem}: XML output is empty but {label} is not")
        return False, False
    if df_other is None or len(df_other) == 0:
        print(f"[FAIL]  {stem}: {label} output is empty but XML is not")
        return False, False

    ok, notes = compare(df_xml, df_other, skip_cols=skip_cols)
    for note in notes:
        print(f"  note: {note}")
    if ok:
        print(f"[OK]    {stem} ({label}): {len(df_xml)} rows match")
    else:
        print(f"[FAIL]  {stem} ({label}): outputs differ")
    return ok, False


def check(stem, xml_path, parquet_path, csv_path, skip_cols=None):
    if not xml_path:
        print(f"[ERROR] {stem}: XML SUMO run failed")
        return False, False
    if not parquet_path:
        print(f"[ERROR]  {stem}: parquet SUMO run failed")
        return False, False
    if not csv_path:
        print(f"[ERROR]  {stem}: CSV SUMO run failed")
        return False, False

    ok, was_skipped = check_and_compare(stem, xml_path, parquet_path, "parquet", skip_cols)
    ok2, was_skipped2 = check_and_compare(stem, xml_path, csv_path, "csv", skip_cols)
    print()
    return ok and ok2, was_skipped and was_skipped2


def main():
    tmpdir = "sumo_parquet_test"
    os.makedirs(tmpdir, exist_ok=True)
    errors = 0
    skipped = 0

    for opt, stem, end_time, args in OUTPUTS:
        print(f"Testing {stem} ({opt}, end={end_time}s)...")
        ok, was_skipped = check(stem, run_sumo(tmpdir, opt, stem, "xml", end_time, args),
                                run_sumo(tmpdir, opt, stem, "parquet", end_time, args),
                                run_sumo(tmpdir, opt, stem, "csv", end_time, args), args.get("skip_cols"))
        if not ok:
            errors += 1
        elif was_skipped:
            skipped += 1

    for stem, generator_script, extra_args, end_time in DETECTOR_OUTPUTS:
        print(f"Testing {stem} (detector, end={end_time}s)...")
        add_xml, xml_path = generate_detector_additional(tmpdir, stem, "xml", generator_script, extra_args)
        add_pq, parquet_path = generate_detector_additional(tmpdir, stem, "parquet", generator_script, extra_args)
        add_csv, csv_path = generate_detector_additional(tmpdir, stem, "csv", generator_script, extra_args)
        if add_xml is None:
            print(f"[ERROR] {stem}: XML additional file generation failed")
            errors += 1
            continue
        if add_pq is None:
            print(f"[ERROR] {stem}: parquet additional file generation failed")
            errors += 1
            continue
        if add_csv is None:
            print(f"[ERROR] {stem}: CSV additional file generation failed")
            errors += 1
            continue
        ok, was_skipped = check(stem, xml_path if run_sumo_with_additional(add_xml, end_time) else None,
                                parquet_path if run_sumo_with_additional(add_pq, end_time) else None,
                                csv_path if run_sumo_with_additional(add_csv, end_time) else None)
        if not ok:
            errors += 1
        elif was_skipped:
            skipped += 1

    for stem, edge_id, period, end_time in ROUTEPROBE_OUTPUTS:
        print(f"Testing {stem} (routeprobe, end={end_time}s)...")
        add_xml, xml_path = generate_routeprobe_additional(tmpdir, stem, "xml", edge_id, period)
        add_pq, parquet_path = generate_routeprobe_additional(tmpdir, stem, "parquet", edge_id, period)
        add_csv, csv_path = generate_routeprobe_additional(tmpdir, stem, "csv", edge_id, period)
        ok, was_skipped = check(stem, xml_path if run_sumo_with_additional(add_xml, end_time) else None,
                                parquet_path if run_sumo_with_additional(add_pq, end_time) else None,
                                csv_path if run_sumo_with_additional(add_csv, end_time) else None)
        if not ok:
            errors += 1
        elif was_skipped:
            skipped += 1

    if errors:
        print(f"Result: FAILED ({errors} output type(s) differ, {skipped} skipped)")
        sys.exit(1)
    else:
        print(f"Result: PASSED (all outputs match, {skipped} skipped)")


if __name__ == "__main__":
    main()
