#!/usr/bin/env python3
import json
import sys
import os
import math
# step5reportResults.py
def load_json(filename):
    """Load a JSON file and return its data."""
    with open(filename, "r") as f:
        return json.load(f)

def extract_rows(data, label):
    """
    Convert yellowBalls and blueBalls into a list of:
    (id, color, x, y, label)
    where x,y correspond to either x0,y0 or x1,y1 depending on label.
    """
    rows = []

    for b in data.get("yellowBalls", []):
        rows.append((b["id"], "yellow", b["x"], b["y"], label))

    for b in data.get("blueBalls", []):
        rows.append((b["id"], "blue", b["x"], b["y"], label))

    return rows

def merge_rows(base_rows, measured_rows):
    """
    Merge base (x0,y0) and measured (x1,y1) rows by (id,color).
    Returns a list of:
    (id, color, x0, y0, x1, y1, d0, d1, error)
    """
    merged = {}

    # Insert base rows
    for r in base_rows:
        ball_id, color, x0, y0, _ = r
        merged[(ball_id, color)] = {
            "x0": x0, "y0": y0,
            "x1": None, "y1": None
        }

    # Insert measured rows
    for r in measured_rows:
        ball_id, color, x1, y1, _ = r
        if (ball_id, color) not in merged:
            merged[(ball_id, color)] = {
                "x0": None, "y0": None,
                "x1": x1, "y1": y1
            }
        else:
            merged[(ball_id, color)]["x1"] = x1
            merged[(ball_id, color)]["y1"] = y1

    # Compute distances and error
    out = []
    for (ball_id, color), vals in merged.items():
        x0, y0 = vals["x0"], vals["y0"]
        x1, y1 = vals["x1"], vals["y1"]

        d0 = math.sqrt(x0*x0 + y0*y0) if x0 is not None else None
        d1 = math.sqrt(x1*x1 + y1*y1) if x1 is not None else None

        error = (d1 - d0) if (d0 is not None and d1 is not None) else None

        out.append((ball_id, color, x0, y0, x1, y1, d0, d1, error))

    out.sort(key=lambda r: r[0])
    return out

def fmt(v):
    """Format numbers nicely, leave None blank."""
    if v is None:
        return " "
    return f"{v:8.2f}"

def print_table(rows):
    """Print the merged table with distances and error."""
    print(
        f"{'ID':>3}  {'Color':<6}  "
        f"{'x0':>8} {'y0':>8}  "
        f"{'x1':>8} {'y1':>8}  "
        f"{'d0':>8} {'d1':>8} {'error':>8}"
    )
    print("-" * 90)

    for r in rows:
        ball_id, color, x0, y0, x1, y1, d0, d1, error = r
        print(
            f"{ball_id:>3}  {color:<6}  "
            f"{fmt(x0)} {fmt(y0)}  "
            f"{fmt(x1)} {fmt(y1)}  "
            f"{fmt(d0)} {fmt(d1)} {fmt(error)}"
        )

def main():
    if len(sys.argv) != 2:
        print("Usage: python step5reportResults.py base.json")
        sys.exit(1)

    base_file = sys.argv[1]
    measured_file = os.path.splitext(base_file)[0] + "_measured.json"

    # Load both files
    base_data = load_json(base_file)
    measured_data = load_json(measured_file)

    # Extract rows
    base_rows = extract_rows(base_data, "x0")
    measured_rows = extract_rows(measured_data, "x1")

    # Merge + compute
    merged_rows = merge_rows(base_rows, measured_rows)

    # Display
    print_table(merged_rows)

if __name__ == "__main__":
    main()