from udi_grammar_py import Chart, Op, rolling
from udi_grammar_py.spec import transfer_kwargs

def test_chart_initialization():
    chart = Chart()
    assert chart._spec == {}

def test_chart_add_source():
    chart = Chart()
    chart.source(source='test_source', name='test_name')
    assert chart._spec['source'] == [{'name': 'test_name', 'source': 'test_source'}]

def test_chart_add_multiple_sources():
    chart = Chart()
    chart.source('test_name_1', 'test_source_2').source('test_name_2', 'test_source_2')
    assert chart._spec['source'] == [{'name': 'test_name_1', 'source': 'test_source_2'}, {'name': 'test_name_2', 'source': 'test_source_2'}]

def test_chart_to_json():
    chart = Chart()
    chart.source('test_name', 'test_source')
    expected_json = '{"source": {"name": "test_name", "source": "test_source"}}'
    assert chart.to_json() == expected_json

def test_chart_to_dict():
    chart = Chart()
    chart.source('test_name', 'test_source')
    expected_dict = {'source': {'name': 'test_name', 'source': 'test_source'}}
    assert chart.to_dict() == expected_dict


def test_transfer_kwargs_empty():
    state = {}
    kwargs = {}
    transfer_kwargs({'field', 'type'}, state, kwargs)
    assert state == {}

def test_transfer_kwargs_valid_args():
    state = {}
    kwargs = {'field': 'age', 'type': 'quantitative'}
    transfer_kwargs({'field', 'type'}, state, kwargs)
    assert state == {'field': 'age', 'type': 'quantitative'}

def test_transfer_kwargs_invalid_args():
    state = {}
    kwargs = {'field': 'age', 'invalid_arg': 'value'}
    transfer_kwargs({'field', 'type'}, state, kwargs)
    assert state == {'field': 'age'}

def test_transfer_kwargs_in_name():
    state = {}
    kwargs = {'in_name': 'input_field'}
    transfer_kwargs({'in'}, state, kwargs)
    assert state == {'in': 'input_field'}
    
def test_transfer_kwargs_out_name():
    state = {}
    kwargs = {'out_name': 'input_field'}
    transfer_kwargs({'out'}, state, kwargs)
    assert state == {'out': 'input_field'}

# Tests charts originally from examples page.

# Bar Charts
def test_total_record_count():
    chart = Chart()
    chart.source('donors', './data/example_donors.csv') \
        .rollup({'count': Op.count()}) \
        .mark('bar') \
        .x(field='count', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "donors", "source": "./data/example_donors.csv"},
        "transformation": [
            {"rollup": {"count": {"op": "count"}}},
        ],
        "representation": {
            "mark": "bar",
            "mapping": {"encoding": "x", "field": "count", "type": "quantitative"},
        },
    }

def test_count_by_category():
    chart = Chart()
    chart.source('donors', './data/example_donors.csv') \
        .groupby('sex') \
        .rollup({'count': Op.count()}) \
        .mark('bar') \
        .x(field='sex', type='nominal') \
        .y(field='count', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "donors", "source": "./data/example_donors.csv"},
        "transformation": [
            {"groupby": "sex"},
            {"rollup": {"count": {"op": "count"}}},
        ],
        "representation": {
            "mark": "bar",
            "mapping": [
                {"encoding": "x", "field": "sex", "type": "nominal"},
                {"encoding": "y", "field": "count", "type": "quantitative"},
            ],
        },
    }

def test_aggregate_by_category():
    chart = Chart()
    chart.source('donors', './data/example_donors.csv') \
        .groupby('sex') \
        .rollup({'average weight': Op.mean('weight')}) \
        .mark('bar') \
        .x(field='sex', type='nominal') \
        .y(field='average weight', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "donors", "source": "./data/example_donors.csv"},
        "transformation": [
            {"groupby": "sex"},
            {"rollup": {"average weight": {"op": "mean", "field": "weight"}}},
        ],
        "representation": {
            "mark": "bar",
            "mapping": [
                {"encoding": "x", "field": "sex", "type": "nominal"},
                {"encoding": "y", "field": "average weight", "type": "quantitative"},
            ],
        },
    }

def test_combining_data_sources():
    chart = Chart()
    chart.source('donors', './data/example_donors.csv') \
        .source('samples', './data/example_samples.csv') \
        .join(on=['id', 'donor_id'], in_name=['donors', 'samples'], out_name='donor_sample_combined') \
        .groupby('sex') \
        .rollup({'sample count': Op.count()}) \
        .mark('bar') \
        .x(field='sex', type='nominal') \
        .y(field='sample count', type='quantitative')

    assert chart.to_dict() == {
        "source": [
            {"name": "donors", "source": "./data/example_donors.csv"},
            {"name": "samples", "source": "./data/example_samples.csv"},
        ],
        "transformation": [
            {
                "in": ["donors", "samples"],
                "join": {"on": ["id", "donor_id"]},
                "out": "donor_sample_combined"
            },
            {"groupby": "sex"},
            {"rollup": {"sample count": {"op": "count"}}},
        ],
        "representation": {
            "mark": "bar",
            "mapping": [
                {"encoding": "x", "field": "sex", "type": "nominal"},
                {"encoding": "y", "field": "sample count", "type": "quantitative"},
            ],
        },
    }

def test_single_stacked_bar_chart():
    chart = Chart()
    chart.source('samples', './data/example_samples.csv') \
        .groupby('organ') \
        .rollup({'count': Op.count()}) \
        .representation().mark('bar') \
        .x(field='count', type='quantitative') \
        .color(field='organ', type='nominal')

    assert chart.to_dict() == {
        "source": {"name": "samples", "source": "./data/example_samples.csv"},
        "transformation": [
            {"groupby": "organ"},
            {"rollup": {"count": {"op": "count"}}},
        ],
        "representation": {
            "mark": "bar",
            "mapping": [
                {"encoding": "x", "field": "count", "type": "quantitative"},
                {"encoding": "color", "field": "organ", "type": "nominal"},
            ],
        },
    }

def test_single_stacked_bar_chart_relative():
    chart = Chart()
    chart.source('samples', './data/example_samples.csv') \
        .groupby('organ') \
        .rollup({'frequency': Op.frequency()}) \
        .mark('bar') \
            .x(field='frequency', type='quantitative') \
            .color(field='organ', type='nominal')

    assert chart.to_dict() == {
        "source": {"name": "samples", "source": "./data/example_samples.csv"},
        "transformation": [
            {"groupby": "organ"},
            {"rollup": {"frequency": {"op": "frequency"}}},
        ],
        "representation": {
            "mark": "bar",
            "mapping": [
                {"encoding": "x", "field": "frequency", "type": "quantitative"},
                {"encoding": "color", "field": "organ", "type": "nominal"},
            ],
        },
    }

def test_multiple_stacked_bar_charts():
    chart = Chart()
    chart.source('samples', './data/example_samples.csv') \
        .groupby(['organ', 'organ_condition']) \
        .rollup({'count': Op.count()}) \
        .mark('bar') \
            .x(field='count', type='quantitative') \
            .y(field='organ', type='nominal') \
            .color(field='organ_condition', type='nominal')

    assert chart.to_dict() == {
        "source": {"name": "samples", "source": "./data/example_samples.csv"},
        "transformation": [
            {"groupby": ["organ", "organ_condition"]},
            {"rollup": {"count": {"op": "count"}}},
        ],
        "representation": {
            "mark": "bar",
            "mapping": [
                {"encoding": "x", "field": "count", "type": "quantitative"},
                {"encoding": "y", "field": "organ", "type": "nominal"},
                {"encoding": "color", "field": "organ_condition", "type": "nominal"},
            ],
        },
    }

def test_multiple_stacked_bar_charts_relative():
    chart = Chart()
    chart.source('samples', './data/example_samples.csv') \
        .groupby('organ', out_name='groupCounts') \
        .rollup({'organ_count': Op.count()}) \
        .groupby(['organ', 'organ_condition']) \
        .rollup({'organ_and_condition_count': Op.count()}) \
        .join(on='organ', in_name=['samples', 'groupCounts'], out_name='datasets') \
        .derive({'frequency': 'd.organ_and_condition_count / d.organ_count'}) \
        .mark('bar') \
            .x(field='frequency', type='quantitative') \
            .y(field='organ', type='nominal') \
            .color(field='organ_condition', type='nominal')

    assert chart.to_dict() == {
        "source": {"name": "samples", "source": "./data/example_samples.csv"},
        "transformation": [
            {"groupby": "organ", "out": "groupCounts"},
            {"rollup": {"organ_count": {"op": "count"}}},
            {"groupby": ["organ", "organ_condition"]},
            {"rollup": {"organ_and_condition_count": {"op": "count"}}},
            {
                "in": ["samples", "groupCounts"],
                "join": {"on": "organ"},
                "out": "datasets"
            },
            {"derive": {"frequency": "d.organ_and_condition_count / d.organ_count"}},
        ],
        "representation": {
            "mark": "bar",
            "mapping": [
                {"encoding": "x", "field": "frequency", "type": "quantitative"},
                {"encoding": "y", "field": "organ", "type": "nominal"},
                {"encoding": "color", "field": "organ_condition", "type": "nominal"},
            ],
        },
    }

# Circular Charts
def test_pie_chart():
    chart = Chart()
    chart.source('samples', './data/example_samples.csv') \
        .transformation() \
        .groupby('organ') \
        .rollup({'frequency': Op.frequency()})
    
    chart.representation().mark('arc') \
         .color(field='organ', type='nominal') \
         .theta(field='frequency', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "samples", "source": "./data/example_samples.csv"},
        "transformation": [
            {"groupby": "organ"},
            {"rollup": {"frequency": {"op": "frequency"}}},
        ],
        "representation": {
            "mark": "arc",
            "mapping": [
                {"encoding": "color", "field": "organ", "type": "nominal"},
                {"encoding": "theta", "field": "frequency", "type": "quantitative"},
            ],
        },
    }

def test_donut_chart():
    chart = Chart()
    chart.source('samples', './data/example_samples.csv') \
        .transformation() \
        .groupby('organ') \
        .rollup({'frequency': Op.frequency()})
    
    chart.representation().mark('arc') \
         .color(field='organ', type='nominal') \
         .theta(field='frequency', type='quantitative') \
         .radius2(value=60)

    assert chart.to_dict() == {
        "source": {"name": "samples", "source": "./data/example_samples.csv"},
        "transformation": [
            {"groupby": "organ"},
            {"rollup": {"frequency": {"op": "frequency"}}},
        ],
        "representation": {
            "mark": "arc",
            "mapping": [
                {"encoding": "color", "field": "organ", "type": "nominal"},
                {"encoding": "theta", "field": "frequency", "type": "quantitative"},
                {"encoding": "radius2", "value": 60},
            ],
        },
    }

# scatterplots
def test_basic_scatter_plot():
    chart = Chart()
    chart.source('donors', './data/example_donors.csv') \
        .representation() \
        .mark('point') \
        .y(field='height', type='quantitative') \
        .x(field='weight', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "donors", "source": "./data/example_donors.csv"},
        "representation": {
            "mark": "point",
            "mapping": [
                {"encoding": "y", "field": "height", "type": "quantitative"},
                {"encoding": "x", "field": "weight", "type": "quantitative"},
            ],
        },
    }

def test_scatter_plot_with_categories():
    chart = Chart()
    chart.source('donors', './data/example_donors.csv') \
        .representation() \
        .mark('point') \
        .y(field='height', type='quantitative') \
        .x(field='weight', type='quantitative') \
        .color(field='sex', type='nominal') \
        .shape(field='sex', type='nominal')

    assert chart.to_dict() == {
        "source": {"name": "donors", "source": "./data/example_donors.csv"},
        "representation": {
            "mark": "point",
            "mapping": [
                {"encoding": "y", "field": "height", "type": "quantitative"},
                {"encoding": "x", "field": "weight", "type": "quantitative"},
                {"encoding": "color", "field": "sex", "type": "nominal"},
                {"encoding": "shape", "field": "sex", "type": "nominal"},
            ],
        },
    }

def test_bubble_plot():
    chart = Chart()
    chart.source('donors', './data/example_donors.csv') \
        .representation() \
        .mark('point') \
        .y(field='height', type='quantitative') \
        .x(field='weight', type='quantitative') \
        .size(field='age', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "donors", "source": "./data/example_donors.csv"},
        "representation": {
            "mark": "point",
            "mapping": [
                {"encoding": "y", "field": "height", "type": "quantitative"},
                {"encoding": "x", "field": "weight", "type": "quantitative"},
                {"encoding": "size", "field": "age", "type": "quantitative"},
            ],
        },
    }
    
# heatmaps

# Heatmaps
def test_heatmap_of_categories():
    chart = Chart()
    chart.source('samples', './data/example_samples.csv') \
        .transformation() \
        .groupby(['organ', 'organ_condition']) \
        .rollup({'count': Op.count()})
    
    chart.representation().mark('rect') \
         .x(field='organ_condition', type='nominal') \
         .y(field='organ', type='nominal') \
         .color(field='count', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "samples", "source": "./data/example_samples.csv"},
        "transformation": [
            {"groupby": ["organ", "organ_condition"]},
            {"rollup": {"count": {"op": "count"}}},
        ],
        "representation": {
            "mark": "rect",
            "mapping": [
                {"encoding": "x", "field": "organ_condition", "type": "nominal"},
                {"encoding": "y", "field": "organ", "type": "nominal"},
                {"encoding": "color", "field": "count", "type": "quantitative"},
            ],
        },
    }

def test_heatmap_with_aggregation():
    chart = Chart()
    chart.source('samples', './data/example_samples.csv') \
        .groupby(['organ', 'organ_condition']) \
        .rollup({'average_weight': Op.mean('weight')}) \
        .mark('rect') \
        .x(field='organ_condition', type='nominal') \
        .y(field='organ', type='nominal') \
        .color(field='average_weight', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "samples", "source": "./data/example_samples.csv"},
        "transformation": [
            {"groupby": ["organ", "organ_condition"]},
            {"rollup": {"average_weight": {"op": "mean", "field": "weight"}}},
        ],
        "representation": {
            "mark": "rect",
            "mapping": [
                {"encoding": "x", "field": "organ_condition", "type": "nominal"},
                {"encoding": "y", "field": "organ", "type": "nominal"},
                {"encoding": "color", "field": "average_weight", "type": "quantitative"},
            ],
        },
    }

def test_heatmap_with_multiple_aggregations():
    chart = Chart()
    chart.source('samples', './data/example_samples.csv') \
        .transformation() \
        .groupby(['organ', 'organ_condition']) \
        .rollup({'average_weight': Op.mean('weight'), 'count': Op.count()})
    
    chart.representation().mark('rect') \
        .x(field='organ_condition', type='nominal') \
        .y(field='organ', type='nominal') \
        .color(field='average_weight', type='quantitative') \
        .size(field='count', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "samples", "source": "./data/example_samples.csv"},
        "transformation": [
            {"groupby": ["organ", "organ_condition"]},
            {"rollup": {"average_weight": {"op": "mean", "field": "weight"}, "count": {"op": "count"}}},
        ],
        "representation": {
            "mark": "rect",
            "mapping": [
                {"encoding": "x", "field": "organ_condition", "type": "nominal"},
                {"encoding": "y", "field": "organ", "type": "nominal"},
                {"encoding": "color", "field": "average_weight", "type": "quantitative"},
                {"encoding": "size", "field": "count", "type": "quantitative"},
            ],
        },
    }

# distribution plots

def test_histogram():
    chart = Chart()
    chart.source('penguins', './data/penguins.csv') \
        .transformation() \
        .binby(field='bill_length_mm', bins=10, nice=True, output={'bin_start': 'start', 'bin_end': 'end'}) \
        .rollup({'count': Op.count()})
    
    chart.representation().mark('rect') \
         .x(field='start', type='quantitative') \
         .x2(field='end', type='quantitative') \
         .y(field='count', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "penguins", "source": "./data/penguins.csv"},
        "transformation": [
            {"binby": {"field": "bill_length_mm", "bins": 10, "nice": True, "output": {"bin_start": "start", "bin_end": "end"}}},
            {"rollup": {"count": {"op": "count"}}},
        ],
        "representation": {
            "mark": "rect",
            "mapping": [
                {"encoding": "x", "field": "start", "type": "quantitative"},
                {"encoding": "x2", "field": "end", "type": "quantitative"},
                {"encoding": "y", "field": "count", "type": "quantitative"},
            ],
        },
    }

def test_kde_density_plot():
    chart = Chart()
    chart.source("penguins", "./data/penguins.csv").kde(
        field="bill_length_mm",
        samples=100,
        output={"sample": "bill_length_mm", "density": "density"},
    ).mark("area") \
    .x(field="bill_length_mm", type="quantitative") \
    .y(field="density", type="quantitative")

    assert chart.to_dict() == {
        "source": {"name": "penguins", "source": "./data/penguins.csv"},
        "transformation": [
            {"kde": {"field": "bill_length_mm", "samples": 100, "output": {"sample": "bill_length_mm", "density": "density"}}},
        ],
        "representation": {
            "mark": "area",
            "mapping": [
                {"encoding": "x", "field": "bill_length_mm", "type": "quantitative"},
                {"encoding": "y", "field": "density", "type": "quantitative"},
            ],
        },
    }

def test_kde_density_plot_grouped():
    chart = Chart()
    chart.source('penguins', './data/penguins.csv') \
        .groupby('species') \
        .kde(field='bill_length_mm', samples=100, output={
            'sample': 'bill_length_mm',
            'density': 'density'
        }) \
        .mark('area') \
            .x(field='bill_length_mm', type='quantitative') \
            .y(field='density', type='quantitative') \
            .color(field='species', type='nominal') \
            .opacity(value=0.25) \
        .mark('line') \
            .x(field='bill_length_mm', type='quantitative') \
            .y(field='density', type='quantitative') \
            .color(field='species', type='nominal')

    assert chart.to_dict() == {
        "source": {"name": "penguins", "source": "./data/penguins.csv"},
        "transformation": [
            {"groupby": "species"},
            {"kde": {"field": "bill_length_mm", "samples": 100, "output": {"sample": "bill_length_mm", "density": "density"}}},
        ],
        "representation": [
            {
                "mark": "area",
                "mapping": [
                    {"encoding": "x", "field": "bill_length_mm", "type": "quantitative"},
                    {"encoding": "y", "field": "density", "type": "quantitative"},
                    {"encoding": "color", "field": "species", "type": "nominal"},
                    {"encoding": "opacity", "value": 0.25},
                ],
            },
            {
                "mark": "line",
                "mapping": [
                    {"encoding": "x", "field": "bill_length_mm", "type": "quantitative"},
                    {"encoding": "y", "field": "density", "type": "quantitative"},
                    {"encoding": "color", "field": "species", "type": "nominal"},
                ],
            },
        ],
    }

def test_empirical_cdf():
    chart = Chart()
    chart.source('penguins', './data/penguins.csv') \
        .transformation() \
        .orderby('bill_length_mm') \
        .derive({'total': 'count()'}) \
        .derive({'percentile': rolling('count() / d.total')})
    
    chart.representation().mark('line') \
         .x(field='bill_length_mm', type='quantitative') \
         .y(field='percentile', type='quantitative')

    assert chart.to_dict() == {
        "source": {"name": "penguins", "source": "./data/penguins.csv"},
        "transformation": [
            {"orderby": "bill_length_mm"},
            {"derive": {"total": "count()"}},
            {"derive": {"percentile": rolling("count() / d.total")}},
        ],
        "representation": {
            "mark": "line",
            "mapping": [
                {"encoding": "x", "field": "bill_length_mm", "type": "quantitative"},
                {"encoding": "y", "field": "percentile", "type": "quantitative"},
            ],
        },
    }

def test_empirical_cdf_grouped():
    chart = Chart()
    chart.source('penguins', './data/penguins.csv') \
        .transformation() \
        .orderby('bill_length_mm') \
        .groupby('species') \
        .derive({'total': 'count()'}) \
        .derive({'percentile': rolling('count() / d.total')})
    
    chart.representation().mark('line') \
         .x(field='bill_length_mm', type='quantitative') \
         .y(field='percentile', type='quantitative') \
         .color(field='species', type='nominal')

    assert chart.to_dict() == {
        "source": {"name": "penguins", "source": "./data/penguins.csv"},
        "transformation": [
            {"orderby": "bill_length_mm"},
            {"groupby": "species"},
            {"derive": {"total": "count()"}},
            {"derive": {"percentile": rolling("count() / d.total")}},
        ],
        "representation": {
            "mark": "line",
            "mapping": [
                {"encoding": "x", "field": "bill_length_mm", "type": "quantitative"},
                {"encoding": "y", "field": "percentile", "type": "quantitative"},
                {"encoding": "color", "field": "species", "type": "nominal"},
            ],
        },
    }

# Table Tests
def test_table_default():
    chart = Chart()
    chart.source('donors', './data/donors.csv') \
        .representation() \
        .mark('row') \
        .text(field='*', mark='text')
    assert chart.to_dict() == {
        "source": {"name": "donors", "source": "./data/donors.csv"},
        "representation": {
            "mark": "row",
            "mapping": {"encoding": "text", "field": "*", "mark": "text"},
        },
    }

def test_table_2():
    chart = Chart()
    chart.source('donors', './data/donors.csv') \
        .mark('row') \
        .color(field='weight_value', mark='rect', type='quantitative') \
        .size(field='height_value', mark='point', type='quantitative') \
        .text(field='*', mark='text')

    assert chart.to_dict() == {
        "source": {
            "name": "donors",
            "source": "./data/donors.csv",
        },
        "representation": {
            "mark": "row",
            "mapping": [
                {
                    "mark": "rect",
                    "field": "weight_value",
                    "encoding": "color",
                    "type": "quantitative",
                },
                {
                    "mark": "point",
                    "field": "height_value",
                    "encoding": "size",
                    "type": "quantitative",
                },
                {"mark": "text", "field": "*", "encoding": "text"},
            ],
        },
    }

def test_table_repeated_encoding():
    chart = Chart()
    chart.source('donors', './data/donors.csv') \
        .representation() \
        .mark('row') \
        .color(field='weight_value', mark='rect', type='quantitative') \
        .color(field='height_value', mark='point', type='quantitative') \
        .text(field='*', mark='text')
    assert chart.to_dict() == {
        "source": {
            "name": "donors",
            "source": "./data/donors.csv",
        },
        "representation": {
            "mark": "row",
            "mapping": [
                {
                    "mark": "rect",
                    "field": "weight_value",
                    "encoding": "color",
                    "type": "quantitative",
                },
                {
                    "mark": "point",
                    "field": "height_value",
                    "encoding": "color",
                    "type": "quantitative",
                },
                {"mark": "text", "field": "*", "encoding": "text"},
            ],
        },
    }


def test_layered_table_example_1():
    chart = Chart()
    chart.source('donors', './data/donors.csv') \
        .mark('row') \
        .color(field='*', mark='rect') \
        .mark('row') \
        .color(field='*', mark='text', value='white')

    assert chart.to_dict() == {
        "source": {
            "name": "donors",
            "source": "./data/donors.csv",
        },
        "representation": [
            {
                "mark": "row",
                "mapping": {"mark": "rect", "field": "*", "encoding": "color"},
            },
            {
                "mark": "row",
                "mapping": {
                    "mark": "text",
                    "field": "*",
                    "value": "white",
                    "encoding": "color",
                },
            },
        ],
    }

def test_layered_table_example_2():
    chart = Chart()
    chart.source('donors', './data/donors.csv') \
        .representation() \
        .mark('row') \
        .size(field='*', mark='point') \
        .color(field='*', mark='point')

    assert chart.to_dict() == {
        "source": {
            "name": "donors",
            "source": "./data/donors.csv",
        },
        "representation": {
            "mark": "row",
            "mapping": [
                {"mark": "point", "field": "*", "encoding": "size"},
                {"mark": "point", "field": "*", "encoding": "color"},
            ],
        },
    }