"""
Write results in various formats.
"""

# Copyright (c) 2025, Carnegie Mellon University.  All Rights Reserved.

from TeIn import TeIn

class Output:

    # -------------- Methods (alphabetical) ---------------- #

    def __init__(self, parameters):
        self._type = parameters['type']
        self._outputPath = parameters['outputPath']
        self._outputLength = parameters['outputLength']


    def execute(self, batch):
        """
        Write the output about the batch to a file

        batch: A dict of {qid: {'qstring': qstring,
                                'ranking': [(score, externalId) ...]}
                          ... }
        """
        if self._type == 'trec_eval':
            teIn = TeIn(self._outputPath, self._outputLength)
            for qid in batch:
                teIn.appendQuery(qid, batch[qid]['ranking'], 'reference')
            teIn.close()
#
# You will need to add another format for HW5.
#
        else:
            raise Exception('Error: Unknown Output format')
            
        return(batch)