import math
from datetime import datetime


def check_df_value(out_df, expected, unique_row_ids, float_tolerance=float(0.00001)):
    """
    Compares a Spark Dataframe to a set of expected values, raising exceptions on detected differences.

    Args:
    out_df: A Spark Dataframe to check.
    expected: The expected values in the out_df. This is a list of rows, each row is a dictionary of column name to
        value.
    unique_row_ids: The columns that are unique to each row. These are used to match rows between out_df and expected.
    float_tolerance: The amount of acceptable difference between float values.
    """

    def extract_row_key(row, unique_row_ids):
        row_ids = []
        for k in unique_row_ids:
            if k not in row:
                raise Exception("Missing row id %s from output %s" % (k, row))
            row_ids.append(row[k])
        return row_ids

    out = out_df.collect()

    out_map = {}
    for r in out:
        k = extract_row_key(r, unique_row_ids)
        out_map[str(k)] = r

    for r in expected:
        k = extract_row_key(r, unique_row_ids)
        if str(k) not in out_map:
            raise Exception("Missing data for user %s" % k)

        out_row = out_map[str(k)]
        for attribute, value in r.items():
            if attribute not in out_row:
                raise Exception("Missing expected attribute %s from row %s" % (attribute, out_row))

            if type(out_row[attribute]) != type(value):
                raise Exception("Type mismatch for attribute %s for row %s is %s, expecting %s" % (
                    attribute, out_row, out_row[attribute], value))

            if isinstance(out_row[attribute], float):
                if not math.isclose(out_row[attribute], value, abs_tol=float_tolerance):
                    raise Exception(
                        "Attribute %s for row %s was %s, expected %s" % (attribute, out_row, out_row[attribute], value))
            elif isinstance(out_row[attribute], datetime) and out_row[attribute].tzinfo != value.tzinfo:
                input_value = out_row[attribute].replace(tzinfo=value.tzinfo)
                if input_value != value:
                    raise Exception(
                        "datetime attribute %s for row %s was %s, expected %s" % (
                            attribute, out_row, input_value, value))
            else:
                if out_row[attribute] != value:
                    raise Exception(
                        "Attribute %s for row %s was %s, expected %s" % (attribute, out_row, out_row[attribute], value))