from typing import Union, Callable, List

import numpy as np

from pyspark.sql import Row
from programs.nodes.nodes import GeounitNode, getNodeAttr, hasNodeAttr, RAW, SYN, UNIT_SYN, GEOCODE
from programs.schema import schema as sk
from programs.schema.schemas.schemamaker import SchemaMaker

def makeHistRowsFromMultiSparse(node: Union[GeounitNode, dict],
                                schema: sk,
                                run_id: int = None,
                                row_recoder: Callable = None,
                                add_schema_name: bool = True,
                                microdata_field: str = 'priv', # priv means "protected via the differential privacy routines in this code base" variable to be renamed after P.L.94-171 production
                                geocode_dict = None) -> List[Row]:
    """
    Converts a node's raw and/or private histogram into a list of Spark Rows, where columns are the dimensions of the histogram
    and 'priv' and 'orig' are column names containing protected and raw counts respectively. If :microdata_field: is set ('priv'/'orig'), then
    instead of having count in that field, the corresponding number (i.e. equal to the count) is inserted in the Row list,
    and the field is discarded.
    :param node: GeounitNode or dict, the node, whose histograms are converted to list or Rows (to then convert to DataFrame)
    :param schema: histogram schema
    :param run_id: auxiliary field, to be keep track of the number of experimental run
    :param row_recoder: recoder class, if some recoding is desired between histogram variables and DataFrame (list of Rows) output columns
    :param add_schema_name: whether to mangle histogram variables name by adding schema name to that, so that they are not confused with desired output names in the recoder
    :param microdata_field: whether to proliferate rows to the histogram count (as opposed to just have a single row for each hist. var. combination with count in a column)
    :return:
    """

    data_present = []  # The list that tells whether the node has raw data, protected data or both. If yes, the data is in this list.
    if hasNodeAttr(node, RAW):
        orig = getNodeAttr(node, RAW).sparse_array
        data_present.append((orig, 'orig'))
    if hasNodeAttr(node, SYN):
        # priv means "protected via the differential privacy routines in this code base" variable to be renamed after P.L.94-171 production
        priv = getNodeAttr(node, SYN).sparse_array
        data_present.append((priv, 'priv'))

    # Collect (multiD) indices of all non-zero histogram counts, for one or both of private/raw datasets
    indices = []
    for dataset in data_present:
        indices = indices + dataset[0].indices.tolist()

    # Leave only unique indices, we will have a Row per each index (each index corresponds to a combination of histogram variables)
    all_nonzero_indices = np.unique(indices)

    #print(f'Node indices: {all_nonzero_indices}')
    #print(f'Node orig counts: {orig.data}')
    #print(f'Node .data: {unit_syn_priv.data}')

    rows = []

    if row_recoder is not None:
        if geocode_dict is not None:
            recoder = row_recoder(geocode_dict=geocode_dict)
        else:
            # To ensure backwards compatibility with old writers
            recoder = row_recoder()

    # Build the row list
    for ind in all_nonzero_indices:
        rowdict = {}

        # Add run_id
        if run_id is not None:
            rowdict['run_id'] = run_id

        # Add geocode
        rowdict[GEOCODE] = getNodeAttr(node, GEOCODE)

        # And protected count, or raw count, or both
        for dataset in data_present:
            rowdict[dataset[1]] = int(dataset[0][0, ind])

        # Add all the histogram variables from the schema
        cell = np.unravel_index(ind, schema.shape)
        for dim, level in enumerate(cell):
            dimcol = f"{schema.dimnames[dim]}"
            if add_schema_name:
                dimcol = dimcol + f"_{schema.name}"
            rowdict[dimcol] = str(level)

        priv_count = None

        # Produce microdata (i.e. remove the count column and instead have that number of identical rows)
        if microdata_field is not None:
            # Number of identical rows to add equal to the count (taking either protected count in field 'priv', or raw count in field 'orig')
            nrows2add = rowdict[microdata_field]
            # Remove the 'priv' or 'orig' count column, that has been used to calculate number of rows to add
            rowdict.pop(microdata_field, None)
        else:
            nrows2add = 1
            if 'priv' in rowdict:
                priv_count = rowdict['priv']

        # Recode the histogram variables to output variables
        if row_recoder is not None:
            rowdict = recoder.recode(rowdict)
            if priv_count is not None:
                rowdict['priv'] = priv_count


        row = Row(**rowdict)
        rows.extend([row] * nrows2add)  # If microdata, then add more than one row, but number of records equal to the count in the histogram

    return rows

def makeHistRowsFromUnitMultiSparse(node: Union[GeounitNode, dict],
                                schema: sk,
                                rows,
                                run_id: int = None,
                                row_recoder: Callable = None,
                                add_schema_name: bool = True,
                                microdata_field: str = 'priv', # 'priv' means "protected via the differential privacy routines in this code base" variable to be renamed after P.L.94-171 production
                                geocode_dict = None) -> List[Row]:
    """
    Converts a node's raw and/or private histogram into a list of Spark Rows, where columns are the dimensions of the histogram
    and 'priv' and 'orig' are column names containing private and raw counts respectively. If :microdata_field: is set ('priv'/'orig'), then
    instead of having count in that field, the corresponding number (i.e. equal to the count) is inserted in the Row list,
    and the field is discarded.
    :param node: GeounitNode or dict, the node, whose histograms are converted to list or Rows (to then convert to DataFrame)
    :param schema: histogram schema
    :param run_id: auxiliary field, to be keep track of the number of experimental run
    :param row_recoder: recoder class, if some recoding is desired between histogram variables and DataFrame (list of Rows) output columns
    :param add_schema_name: whether to mangle histogram variables name by adding schema name to that, so that they are not confused with desired output names in the recoder
    :param microdata_field: whether to proliferate rows to the histogram count (as opposed to just have a single row for each hist. var. combination with count in a column)
    :return:
    """

    unit_syn_priv = getNodeAttr(node, UNIT_SYN).sparse_array

    # Collect (multiD) indices of all non-zero histogram counts, for one or both of private/raw datasets
    indices = unit_syn_priv.indices.tolist()

    # Leave only unique indices, we will have a Row per each index (each index corresponds to a combination of histogram variables)
    all_nonzero_indices = np.unique(indices)

    if row_recoder is not None:
        if geocode_dict is not None:
            recoder = row_recoder(geocode_dict=geocode_dict)
        else:
            # To ensure backwards compatibility with old writers
            recoder = row_recoder()

    # Build the row list
    for ind in all_nonzero_indices:
        rowdict = {}

        # Add geocode
        rowdict[GEOCODE] = getNodeAttr(node, GEOCODE)

        # And protected count, or raw count, or both
        rowdict["priv"] = int(unit_syn_priv[0, ind])

        # Add all the histogram variables from the schema
        cell = np.unravel_index(ind, schema.shape)
        for dim, level in enumerate(cell):
            dimcol = f"{schema.dimnames[dim]}"
            rowdict[dimcol] = str(level)


        # Produce microdata (i.e. remove the count column and instead have that number of identical rows)
        if microdata_field is not None:
            # Number of identical rows to add equal to the count (taking either private count in field 'priv', or raw count in field 'orig')
            nrows2add = rowdict[microdata_field]
            # Remove the 'priv' or 'orig' count column, that has been used to calculate number of rows to add
            rowdict.pop(microdata_field, None)
        else:
            nrows2add = 1

        #Only include vacant/GQ rows. TODO: parameterize or refactor this
        if int(rowdict['tenvacgq']) <= 3:
            continue

        # Recode the histogram variables to output variables
        rowdict = recoder.recode(rowdict)

        row = Row(**rowdict)
        rows.extend([row] * nrows2add)  # If microdata, then add more than one row, but number of records equal to the count in the histogram

    return rows

from typing import Union, Callable, List

import numpy as np

from pyspark.sql import Row
from programs.nodes.nodes import GeounitNode, getNodeAttr, hasNodeAttr, RAW, SYN, UNIT_SYN, GEOCODE
from programs.schema import schema as sk
from programs.schema.schemas.schemamaker import SchemaMaker

def makeHistRowsFromMultiSparse(node: Union[GeounitNode, dict],
                                schema: sk,
                                run_id: int = None,
                                row_recoder: Callable = None,
                                add_schema_name: bool = True,
                                microdata_field: str = 'priv', # priv means "protected via the differential privacy routines in this code base" variable to be renamed after P.L.94-171 production
                                geocode_dict = None) -> List[Row]:
    """
    Converts a node's raw and/or private histogram into a list of Spark Rows, where columns are the dimensions of the histogram
    and 'priv' and 'orig' are column names containing protected and raw counts respectively. If :microdata_field: is set ('priv'/'orig'), then
    instead of having count in that field, the corresponding number (i.e. equal to the count) is inserted in the Row list,
    and the field is discarded.
    :param node: GeounitNode or dict, the node, whose histograms are converted to list or Rows (to then convert to DataFrame)
    :param schema: histogram schema
    :param run_id: auxiliary field, to be keep track of the number of experimental run
    :param row_recoder: recoder class, if some recoding is desired between histogram variables and DataFrame (list of Rows) output columns
    :param add_schema_name: whether to mangle histogram variables name by adding schema name to that, so that they are not confused with desired output names in the recoder
    :param microdata_field: whether to proliferate rows to the histogram count (as opposed to just have a single row for each hist. var. combination with count in a column)
    :return:
    """

    data_present = []  # The list that tells whether the node has raw data, protected data or both. If yes, the data is in this list.
    if hasNodeAttr(node, RAW):
        orig = getNodeAttr(node, RAW).sparse_array
        data_present.append((orig, 'orig'))
    if hasNodeAttr(node, SYN):
        # priv means "protected via the differential privacy routines in this code base" variable to be renamed after P.L.94-171 production
        priv = getNodeAttr(node, SYN).sparse_array
        data_present.append((priv, 'priv'))

    # Collect (multiD) indices of all non-zero histogram counts, for one or both of private/raw datasets
    indices = []
    for dataset in data_present:
        indices = indices + dataset[0].indices.tolist()

    # Leave only unique indices, we will have a Row per each index (each index corresponds to a combination of histogram variables)
    all_nonzero_indices = np.unique(indices)

    #print(f'Node indices: {all_nonzero_indices}')
    #print(f'Node orig counts: {orig.data}')
    #print(f'Node .data: {unit_syn_priv.data}')

    rows = []

    if row_recoder is not None:
        if geocode_dict is not None:
            recoder = row_recoder(geocode_dict=geocode_dict)
        else:
            # To ensure backwards compatibility with old writers
            recoder = row_recoder()

    # Build the row list
    for ind in all_nonzero_indices:
        rowdict = {}

        # Add run_id
        if run_id is not None:
            rowdict['run_id'] = run_id

        # Add geocode
        rowdict[GEOCODE] = getNodeAttr(node, GEOCODE)

        # And protected count, or raw count, or both
        for dataset in data_present:
            rowdict[dataset[1]] = int(dataset[0][0, ind])

        # Add all the histogram variables from the schema
        cell = np.unravel_index(ind, schema.shape)
        for dim, level in enumerate(cell):
            dimcol = f"{schema.dimnames[dim]}"
            if add_schema_name:
                dimcol = dimcol + f"_{schema.name}"
            rowdict[dimcol] = str(level)

        priv_count = None

        # Produce microdata (i.e. remove the count column and instead have that number of identical rows)
        if microdata_field is not None:
            # Number of identical rows to add equal to the count (taking either protected count in field 'priv', or raw count in field 'orig')
            nrows2add = rowdict[microdata_field]
            # Remove the 'priv' or 'orig' count column, that has been used to calculate number of rows to add
            rowdict.pop(microdata_field, None)
        else:
            nrows2add = 1
            if 'priv' in rowdict:
                priv_count = rowdict['priv']

        # Recode the histogram variables to output variables
        if row_recoder is not None:
            rowdict = recoder.recode(rowdict)
            if priv_count is not None:
                rowdict['priv'] = priv_count


        row = Row(**rowdict)
        rows.extend([row] * nrows2add)  # If microdata, then add more than one row, but number of records equal to the count in the histogram

    return rows

def makeHistRowsFromUnitMultiSparse(node: Union[GeounitNode, dict],
                                schema: sk,
                                rows,
                                run_id: int = None,
                                row_recoder: Callable = None,
                                add_schema_name: bool = True,
                                microdata_field: str = 'priv', # 'priv' means "protected via the differential privacy routines in this code base" variable to be renamed after P.L.94-171 production
                                geocode_dict = None) -> List[Row]:
    """
    Converts a node's raw and/or private histogram into a list of Spark Rows, where columns are the dimensions of the histogram
    and 'priv' and 'orig' are column names containing private and raw counts respectively. If :microdata_field: is set ('priv'/'orig'), then
    instead of having count in that field, the corresponding number (i.e. equal to the count) is inserted in the Row list,
    and the field is discarded.
    :param node: GeounitNode or dict, the node, whose histograms are converted to list or Rows (to then convert to DataFrame)
    :param schema: histogram schema
    :param run_id: auxiliary field, to be keep track of the number of experimental run
    :param row_recoder: recoder class, if some recoding is desired between histogram variables and DataFrame (list of Rows) output columns
    :param add_schema_name: whether to mangle histogram variables name by adding schema name to that, so that they are not confused with desired output names in the recoder
    :param microdata_field: whether to proliferate rows to the histogram count (as opposed to just have a single row for each hist. var. combination with count in a column)
    :return:
    """

    unit_syn_priv = getNodeAttr(node, UNIT_SYN).sparse_array

    # Collect (multiD) indices of all non-zero histogram counts, for one or both of private/raw datasets
    indices = unit_syn_priv.indices.tolist()

    # Leave only unique indices, we will have a Row per each index (each index corresponds to a combination of histogram variables)
    all_nonzero_indices = np.unique(indices)

    if row_recoder is not None:
        if geocode_dict is not None:
            recoder = row_recoder(geocode_dict=geocode_dict)
        else:
            # To ensure backwards compatibility with old writers
            recoder = row_recoder()

    # Build the row list
    for ind in all_nonzero_indices:
        rowdict = {}

        # Add geocode
        rowdict[GEOCODE] = getNodeAttr(node, GEOCODE)

        # And protected count, or raw count, or both
        rowdict["priv"] = int(unit_syn_priv[0, ind])

        # Add all the histogram variables from the schema
        cell = np.unravel_index(ind, schema.shape)
        for dim, level in enumerate(cell):
            dimcol = f"{schema.dimnames[dim]}"
            rowdict[dimcol] = str(level)


        # Produce microdata (i.e. remove the count column and instead have that number of identical rows)
        if microdata_field is not None:
            # Number of identical rows to add equal to the count (taking either private count in field 'priv', or raw count in field 'orig')
            nrows2add = rowdict[microdata_field]
            # Remove the 'priv' or 'orig' count column, that has been used to calculate number of rows to add
            rowdict.pop(microdata_field, None)
        else:
            nrows2add = 1

        #Only include vacant/GQ rows. TODO: parameterize or refactor this
        if int(rowdict['tenvacgq']) <= 3:
            continue

        # Recode the histogram variables to output variables
        rowdict = recoder.recode(rowdict)

        row = Row(**rowdict)
        rows.extend([row] * nrows2add)  # If microdata, then add more than one row, but number of records equal to the count in the histogram

    return rows

# priv means "protected via the differential privacy routines in this code base" variable to be renamed after P.L.94-171 production Converts a node's raw and/or private histogram into a list of Spark Rows, where columns are the dimensions of the histogram     and 'priv' and 'orig' are column names containing protected and raw counts respectively. If :microdata_field: is set ('priv'/'orig'), then     instead of having count in that field, the corresponding number (i.e. equal to the count) is inserted in the Row list,     and the field is discarded.     :param node: GeounitNode or dict, the node, whose histograms are converted to list or Rows (to then convert to DataFrame)     :param schema: histogram schema     :param run_id: auxiliary field, to be keep track of the number of experimental run     :param row_recoder: recoder class, if some recoding is desired between histogram variables and DataFrame (list of Rows) output columns     :param add_schema_name: whether to mangle histogram variables name by adding schema name to that, so that they are not confused with desired output names in the recoder     :param microdata_field: whether to proliferate rows to the histogram count (as opposed to just have a single row for each hist. var. combination with count in a column)     :return: # The list that tells whether the node has raw data, protected data or both. If yes, the data is in this list. # priv means "protected via the differential privacy routines in this code base" variable to be renamed after P.L.94-171 production # Collect (multiD) indices of all non-zero histogram counts, for one or both of private/raw datasets # Leave only unique indices, we will have a Row per each index (each index corresponds to a combination of histogram variables) #print(f'Node indices: {all_nonzero_indices}') #print(f'Node orig counts: {orig.data}') #print(f'Node .data: {unit_syn_priv.data}') # To ensure backwards compatibility with old writers # Build the row list # Add run_id # Add geocode # And protected count, or raw count, or both # Add all the histogram variables from the schema # Produce microdata (i.e. remove the count column and instead have that number of identical rows) # Number of identical rows to add equal to the count (taking either protected count in field 'priv', or raw count in field 'orig') # Remove the 'priv' or 'orig' count column, that has been used to calculate number of rows to add # Recode the histogram variables to output variables # If microdata, then add more than one row, but number of records equal to the count in the histogram # 'priv' means "protected via the differential privacy routines in this code base" variable to be renamed after P.L.94-171 production Converts a node's raw and/or private histogram into a list of Spark Rows, where columns are the dimensions of the histogram     and 'priv' and 'orig' are column names containing private and raw counts respectively. If :microdata_field: is set ('priv'/'orig'), then     instead of having count in that field, the corresponding number (i.e. equal to the count) is inserted in the Row list,     and the field is discarded.     :param node: GeounitNode or dict, the node, whose histograms are converted to list or Rows (to then convert to DataFrame)     :param schema: histogram schema     :param run_id: auxiliary field, to be keep track of the number of experimental run     :param row_recoder: recoder class, if some recoding is desired between histogram variables and DataFrame (list of Rows) output columns     :param add_schema_name: whether to mangle histogram variables name by adding schema name to that, so that they are not confused with desired output names in the recoder     :param microdata_field: whether to proliferate rows to the histogram count (as opposed to just have a single row for each hist. var. combination with count in a column)     :return: # Collect (multiD) indices of all non-zero histogram counts, for one or both of private/raw datasets # Leave only unique indices, we will have a Row per each index (each index corresponds to a combination of histogram variables) # To ensure backwards compatibility with old writers # Build the row list # Add geocode # And protected count, or raw count, or both # Add all the histogram variables from the schema # Produce microdata (i.e. remove the count column and instead have that number of identical rows) # Number of identical rows to add equal to the count (taking either private count in field 'priv', or raw count in field 'orig') # Remove the 'priv' or 'orig' count column, that has been used to calculate number of rows to add #Only include vacant/GQ rows. TODO: parameterize or refactor this # Recode the histogram variables to output variables # If microdata, then add more than one row, but number of records equal to the count in the histogram