"""Train a small EC-frequency baseline using only the training split."""
import argparse
from collections import Counter
import json
from pathlib import Path

import pandas as pd


def train(path, output):
    root = Path(__file__).parent
    task = json.loads((root / 'task.json').read_text())
    valid = set(json.loads((root / 'ec_vocabulary.json').read_text()))
    frame = pd.read_parquet(path)
    unique = frame.drop_duplicates(task['input_columns'] + ['ec_number'])
    if not set(unique.ec_number) <= valid:
        raise ValueError('Training data contains ECs outside the pinned vocabulary.')
    counts = Counter(unique.ec_number)
    ranked = sorted(counts, key=lambda ec: (-counts[ec], ec))
    if not ranked:
        raise ValueError('Training data is empty.')
    limit = 1  # Both tasks are scored on the first EC only.
    result = {'prediction': ranked[:limit], 'training_pairs': len(unique),
              'baseline': 'EC frequency; ignores the input'}
    target = Path(output)
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(json.dumps(result, indent=2) + '\n')
    print(f'Trained on {len(unique)} input–EC pairs; saved {target}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--train', required=True)
    parser.add_argument('--output', default='weights/baseline.json')
    args = parser.parse_args()
    train(args.train, args.output)