#  ___________________________________________________________________________
#
#  Prescient
#  Copyright 2020 National Technology & Engineering Solutions of Sandia, LLC
#  (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S.
#  Government retains certain rights in this software.
#  This software is distributed under the Revised BSD License.
#  ___________________________________________________________________________

import csv
from typing import TypeVar, Callable, Iterable, Sequence, Any, Tuple, Union, Dict

RecordType = TypeVar('RecordType')
KeyType = TypeVar('KeyType')
Dialect = Union[str, csv.Dialect]

class CsvReporter:
    ''' Writes a row of data to a CSV file for each record object provided.

        An instance of this class writes data in CSV format to a provided
        file (or file-like object).  Each time write_record() is called,
        data is extracted from the supplied record object and a single row
        is written to the CSV file.  How data is extracted from the record
        is determined by callable objects passed as arguments when the 
        CsvReporter is created.
    '''

    @staticmethod
    def from_dict(fileobj, 
                  header_to_value_map: Dict[str, Callable[[RecordType], Any]],
                  write_headers: bool = True,
                  dialect: Dialect = 'excel'):
        ''' Create a CsvReporter from a dictionary of header to value method

            This method creates a new CsvReporter from a supplied dictionary,
            with one dictionary entry per column in the output.  For each column,
            the the key is the column header and the value is a function that 
            is passed a record object and returns the column's value.

            Parameters
            ----------
            fileobj
                A file or other object with a write() method that data will be written to
            header_to_value_map
                A dictionary that maps column names to a function that provides a value for the column
            dialect
                A dialect or dialect name, as described in documentation for csv.writer().

            Returns
            -------
            Returns a new CsvReporter.
        '''
        headers = header_to_value_map.keys() if write_headers else None
        row_extractor = lambda rec: (f(rec) for f in header_to_value_map.values())
        return CsvReporter(fileobj, headers, row_extractor, dialect)


    def __init__(self, fileobj, 
                 headers: Iterable[str], 
                 row_data_provider: Callable[[RecordType], Iterable[Any]],
                 dialect: Dialect = 'excel'):
        ''' Constructor

            Parameters
            ----------
            fileobj
                A file or other object with a write() method that data will be written to
            headers
                Strings that will be the first line written to the CSV.  May be None if
                a header row is not wanted.
            row_data_provider
                A callable that takes a record as an argument and returns an array of
                values that will be written to the file as a row of CSV.
            dialect
                A dialect or dialect name, as described in documentation for csv.writer().
        '''
        self._writer = csv.writer(fileobj, dialect)
        if headers is not None:
            self._writer.writerow(headers)
        self._row_data_provider = row_data_provider

    def write_record(self, record: RecordType):
        ''' Write a row to the CSV, with data from the supplied record

            This method converts a record object to an array of values by calling the row_from_record
            iterable supplied to the constructor.  The resulting array is then written to the file.

            Parameters
            ----------
            record
                A data record that will be converted to a line of CSV and written to the file
        '''
        row = self._row_data_provider(record)
        self._writer.writerow(row)

        
class CsvMultiRowReporter:
    ''' Writes multiple rows of data to a CSV file for each record object provided.

        An instance of this class writes data in CSV format to a provided
        file (or file-like object).  Each time write_record() is called,
        a set of row keys is generated for the record, and row of data is written
        to the CSV for each row key.  Row keys are generated by a Callable supplied
        to the constructor, and the row of data for each key is generated by a
        second Callable, also supplied to the constructor.
    '''

    @staticmethod
    def from_dict(fileobj, 
                  key_provider: Callable[[RecordType], Iterable[KeyType]],
                  header_to_value_map: Dict[str, Callable[[RecordType, KeyType], Any]],
                  dialect: Dialect = 'excel'):
        ''' Create a CsvMultiRowReporter from a dictionary that maps each column header to a
            function that provides the column's value.

            This method creates a new CsvMultiRowReporter from a supplied dictionary,
            with one dictionary entry per column in the output.  For each column,
            the the key is the column header and the value is a function that 
            supplies the column's value.  These value providing functions are passed
            the original record and the row's key as arguments.

            Parameters
            ----------
            fileobj
                A file or other object with a write() method that data will be written to
            header_to_value_map
                A dictionary that maps column names to a function that provides a value for the column
            key_provider
                A callable that takes a record as an argument and returns an iterable of
                rowkeys.  One line of CSV data will be written to the file for each returned
                rowkey.
            dialect
                A dialect or dialect name, as described in documentation for csv.writer().

            Returns
            -------
            Returns a new CsvMultiRowReporter.
        '''
        headers = header_to_value_map.keys()
        row_extractor = lambda rec,key: (f(rec,key) for f in header_to_value_map.values())
        return CsvMultiRowReporter(fileobj, headers, key_provider, row_extractor, dialect)


    def __init__(self, fileobj, 
                 headers: Sequence[str], 
                 key_provider: Callable[[RecordType], Iterable[KeyType]],
                 row_data_provider: Callable[[Tuple[RecordType, KeyType]], Iterable[Sequence[Any]]],
                 dialect: Dialect = 'excel'):
        ''' Constructor

            Parameters
            ----------
            fileobj
                A file or other object with a write() method that data will be written to
            headers
                Strings that will be the first line written to the CSV.  May be None if
                a header row is not wanted.
            key_provider
                A callable that takes a record as an argument and returns an iterable of
                rowkeys.  One line of CSV data will be written to the file for each returned
                rowkey.
            row_data_provider
                A callable that takes a record and a row key as arguments and returns an 
                array of values that will be written to the file as a row of CSV data 
                for the record and the rowkey.
            dialect
                A dialect or dialect name, as described in documentation for csv.writer().
        '''
        self._writer = csv.writer(fileobj)
        if headers is not None:
            self._writer.writerow(headers)
        self._key_provider = key_provider
        self._row_data_provider = row_data_provider

    def write_record(self, record: RecordType):
        ''' Write multiple rows of CSV data extracted from a single supplied record

            This method converts a record object to any number of rows of CSV and writes them to the file.
            The number of rows is determined by passing the record to the key_provider callable
            supplied to the constructor.  For each key returned by that method, the row_data_provider method
            is called, passing the original record and the row key as arguments.  The array returned from 
            that call is then written to the file in CSV format.

            Parameters
            ----------
            record
                A data record that will be converted to multiple lines of CSV and written to the file
        '''
        rows = (self._row_data_provider(record, key) for key in self._key_provider(record))
        self._writer.writerows(rows)


if __name__ == "__main__":
    class C():
        def __init__(self, a, b):
            self.a = a
            self.b = b

    recs = [C(1,2), C(3,4), C(6,5), C(10,2)]
    headers = ['Column B', 'Column A']
    one_row = lambda rec: (rec.b, rec.a)
    
    with open('temp.csv', 'w', newline='') as f:
        writer = CsvReporter(f, headers, one_row)
        for rec in recs:
            writer.write_record(rec)

    keys = lambda rec: (rec.a, rec.a + 1)
    row_from_rowkey = lambda rec, key: (key, rec.a)
    with open('temp2.csv', 'w', newline='') as f:
        writer = CsvMultiRowReporter(f, headers, keys, row_from_rowkey)
        for rec in recs:
            writer.write_record(rec)

    with open('temp3.csv', 'w', newline='') as f:
        writer = CsvReporter.from_dict(f, {'Column A': lambda t: t.a,
                                           'Column Q': lambda t: t.a,
                                           'Column Z': lambda t: t.a,
                                           'Column H': lambda t: t.a,
                                           'Column B': lambda t: t.b})
        for rec in recs:
            writer.write_record(rec)
