"""abreka exp diff — side-by-side comparison of two experiments."""

from __future__ import annotations

import click
from rich.console import Console
from rich.table import Table

from abreka.config import find_abreka_dir, load_config
from abreka.persistence import load_experiment, normalize_exp_id


@click.command("diff")
@click.argument("id1")
@click.argument("id2")
def diff_cmd(id1: str, id2: str):
    """Side-by-side comparison of two experiments."""
    abreka_dir = find_abreka_dir()
    config = load_config(abreka_dir)
    id1 = normalize_exp_id(id1)
    id2 = normalize_exp_id(id2)
    exp1 = load_experiment(abreka_dir, id1)
    exp2 = load_experiment(abreka_dir, id2)

    maximize = config.experiment.metric_direction == "maximize"

    console = Console()
    table = Table(title=f"Diff: {id1} vs {id2}", show_header=True, header_style="bold")
    table.add_column("Field", style="bold")
    table.add_column(f"Exp {id1}")
    table.add_column(f"Exp {id2}")

    table.add_row("Status", exp1.status.value, exp2.status.value)
    table.add_row("Hypothesis", exp1.hypothesis, exp2.hypothesis)
    table.add_row("Method", exp1.method, exp2.method)
    table.add_row("Parents", ", ".join(exp1.parents), ", ".join(exp2.parents))

    # Collect all metric keys
    all_splits = sorted(set(list(exp1.metrics.keys()) + list(exp2.metrics.keys())))
    for split in all_splits:
        keys1 = exp1.metrics.get(split, {})
        keys2 = exp2.metrics.get(split, {})
        all_keys = sorted(set(list(keys1.keys()) + list(keys2.keys())))
        for key in all_keys:
            v1 = keys1.get(key)
            v2 = keys2.get(key)
            s1 = f"{v1:.6f}" if v1 is not None else "-"
            s2 = f"{v2:.6f}" if v2 is not None else "-"
            # Color the better value green, worse red
            if v1 is not None and v2 is not None and v1 != v2:
                better_is_1 = (v1 > v2) == maximize
                s1 = f"[green]{s1}[/green]" if better_is_1 else f"[red]{s1}[/red]"
                s2 = f"[green]{s2}[/green]" if not better_is_1 else f"[red]{s2}[/red]"
            table.add_row(f"{split}:{key}", s1, s2)

    console.print(table)
