import argparse
import logging

import numpy as np
import pandas as pd
from omegaconf import OmegaConf

from womac.config import WomacConfig
from womac.womac import Womac

logger = logging.getLogger(__name__)


def main():
    """Command-line interface for running WOMAC."""
    # configure logging for CLI output
    logging.basicConfig(
        format="%(asctime)s %(levelname)s: %(message)s",
        datefmt="%H:%M:%S",
        level=logging.INFO,
    )
    parser = argparse.ArgumentParser(description="Run WOMAC scoring from the command line")
    parser.add_argument("--predictions", "-p", required=True, help="Path to predictions file (CSV or .npy)")
    parser.add_argument("--targets", "-t", required=True, help="Path to targets file (CSV or .npy)")
    parser.add_argument("--config", "-c", required=True, help="Path to YAML configuration file")
    parser.add_argument("--output", "-o", default=None, help="Optional path to save leaderboard CSV")
    parser.add_argument("--top_k", "-k", type=int, default=25, help="Log the top K experts indices.")
    args = parser.parse_args()

    def load_array(path):
        if path.endswith(".csv"):
            return pd.read_csv(path).to_numpy()  # expects first row as header
        elif path.endswith(".npy"):
            return np.load(path)
        else:
            raise ValueError(f"Unsupported file format: {path}")

    X = load_array(args.predictions)
    Y = load_array(args.targets).squeeze()

    # Load and validate config via OmegaConf
    raw_cfg = OmegaConf.load(args.config)
    womac_keys = ["min_responses_per_task", "min_responses_per_expert", "device", "seed"]
    womac_init = {key: raw_cfg.pop(key, None) for key in womac_keys}
    schema = OmegaConf.structured(WomacConfig)
    cfg = OmegaConf.merge(schema, raw_cfg)
    wcfg: WomacConfig = OmegaConf.to_object(cfg)

    # Run WOMAC scoring
    model = Womac(**womac_init)
    result = model.score_competition(X, Y, wcfg)

    # Log top-k or full leaderboard
    logger.info("\n%s", result.leaderboard(args.top_k))

    # Optionally save results
    if args.output:
        df = pd.DataFrame(
            {
                "rank": list(range(1, len(result.ranked_indices) + 1)),
                "expert_index": result.ranked_indices.tolist(),
                "score": result.ranked_scores.tolist(),
            }
        )
        df.to_csv(args.output, index=False)
