#!/usr/bin/env python3
"""Combine CP-even and CP-odd artifacts into labeled binary artifacts."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np

from validation.artifacts import ValidationArtifact, load_artifact, save_artifact


def _combine(even: ValidationArtifact, odd: ValidationArtifact) -> ValidationArtifact:
    node_shift = len(even.node_features_flat)
    edge_shift = len(even.edge_src_flat)
    sample_id = np.concatenate(
        [
            np.char.add("ttH_CP_even:", even.sample_id),
            np.char.add("ttH_CP_odd:", odd.sample_id),
        ]
    )
    labels = np.concatenate(
        [
            np.zeros(even.event_count, dtype=np.int64),
            np.ones(odd.event_count, dtype=np.int64),
        ]
    )
    node_offsets = np.concatenate(
        [even.node_offsets, odd.node_offsets[1:] + node_shift]
    )
    edge_offsets = np.concatenate(
        [even.edge_offsets, odd.edge_offsets[1:] + edge_shift]
    )
    return ValidationArtifact(
        sample_id=sample_id,
        labels=labels,
        folds=np.concatenate([even.folds, odd.folds]),
        weights=np.concatenate([even.weights, odd.weights]),
        globals=np.concatenate([even.globals, odd.globals]),
        node_features_flat=np.concatenate(
            [even.node_features_flat, odd.node_features_flat]
        ),
        node_offsets=node_offsets,
        edge_src_flat=np.concatenate([even.edge_src_flat, odd.edge_src_flat]),
        edge_dst_flat=np.concatenate([even.edge_dst_flat, odd.edge_dst_flat]),
        edge_features_flat=np.concatenate(
            [even.edge_features_flat, odd.edge_features_flat]
        ),
        edge_offsets=edge_offsets,
        manifest={
            "task": "ttH_cp_even_vs_cp_odd",
            "label_map": {"CP_even": 0, "CP_odd": 1},
            "selection": "none",
            "sources": {"CP_even": even.manifest, "CP_odd": odd.manifest},
        },
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--even", type=Path, required=True)
    parser.add_argument("--odd", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    artifact = _combine(load_artifact(args.even), load_artifact(args.odd))
    save_artifact(artifact, args.output)
    print(f"wrote {args.output} ({artifact.event_count} events)")
    return 0


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