import pandas as pd
import jinja2

from collections import OrderedDict
from json import dumps
from IPython.html.widgets import interact
from IPython.html import widgets
from IPython.display import display, display_pretty, Javascript, HTML
from IPython.utils.traitlets import Any, Bool, Dict, List, Unicode
from threading import Lock
from urllib2 import urlopen

sub_est_2012_df = pd.read_csv(
    urlopen('http://www.census.gov/popest/data/cities/totals/2012/files/SUB-EST2012.csv'),
    encoding='latin-1',
    dtype={'STATE': 'str', 'COUNTY': 'str', 'PLACE': 'str'}
)

sub_est_2012_df_by_county = sub_est_2012_df[sub_est_2012_df.SUMLEV == 50]
sub_est_2012_df_by_state = sub_est_2012_df_by_county.groupby(['STATE']).sum()

# Alternatively we could have just taken the summary rows for the states

# sub_est_2012_df_by_state = sub_est_2012_df[sub_est_2012_df.SUMLEV == 40]

# Taken from http://www.census.gov/geo/reference/ansi_statetables.html

state = pd.read_csv(urlopen('http://www.census.gov/geo/reference/docs/state.txt'), sep='|', dtype={'STATE': 'str'})

state.drop(
    ['STATENS'],
    inplace=True, axis=1
)

sub_est_2012_df_by_state = pd.merge(sub_est_2012_df_by_state, state, left_index=True, right_on='STATE')
sub_est_2012_df_by_state.drop(
    ['SUMLEV', 'COUSUB', 'CONCIT', 'ESTIMATESBASE2010', 'POPESTIMATE2010', 'POPESTIMATE2011'],
    inplace=True, axis=1
)

# Taken from http://www.census.gov/popest/data/state/asrh/2012/SC-EST2012-AGESEX-CIV.html

sc_est2012_agesex_civ_df = pd.read_csv(
    urlopen('http://www.census.gov/popest/data/state/asrh/2012/files/SC-EST2012-AGESEX-CIV.csv'),
    encoding='latin-1',
    dtype={'SUMLEV': 'str'}
)

sc_est2012_agesex_civ_df_sumlev040 = sc_est2012_agesex_civ_df[
    (sc_est2012_agesex_civ_df.SUMLEV == '040') &
    (sc_est2012_agesex_civ_df.SEX != 0) &
    (sc_est2012_agesex_civ_df.AGE != 999)
]
sc_est2012_agesex_civ_df_sumlev040.drop(
    ['SUMLEV', 'NAME', 'ESTBASE2010_CIV', 'POPEST2010_CIV', 'POPEST2011_CIV'],
    inplace=True, axis=1
)
sc_est2012_agesex_civ_df_sumlev040['STATE'] = sc_est2012_agesex_civ_df_sumlev040['STATE'].apply(lambda x: '%02d' % (x,))

sc_est2012_sex = sc_est2012_agesex_civ_df_sumlev040.groupby(['STATE', 'REGION', 'DIVISION', 'SEX'], as_index=False)[['POPEST2012_CIV']].sum()
sc_est2012_sex = pd.merge(sc_est2012_sex, state, left_on='STATE', right_on='STATE')

sc_est2012_age = sc_est2012_agesex_civ_df_sumlev040.groupby(['STATE', 'REGION', 'DIVISION', 'AGE'], as_index=False)[['POPEST2012_CIV']].sum()

age_buckets = pd.cut(sc_est2012_age.AGE, range(0,100,20))

sc_est2012_age = sc_est2012_age.groupby(['STATE', 'REGION', 'DIVISION', age_buckets], as_index=False)['POPEST2012_CIV'].sum()
sc_est2012_age = pd.merge(sc_est2012_age, state, left_on='STATE', right_on='STATE')

region_codes = {
    0: 'United States Total',
    1: 'Northeast',
    2: 'Midwest',
    3: 'South',
    4: 'West'
}
division_codes = {
    0: 'United States Total',
    1: 'New England',
    2: 'Middle Atlantic',
    3: 'East North Central',
    4: 'West North Central',
    5: 'South Atlantic',
    6: 'East South Central',
    7: 'West South Central',
    8: 'Mountain',
    9: 'Pacific'
}

%%javascript
require.config({
    paths: {
        d3: '//cdnjs.cloudflare.com/ajax/libs/d3/3.4.8/d3.min'
    }
});

display(HTML("""
<style>
.bar {
  fill: steelblue;
}
.bar:hover {
  fill: brown;
}
.axis {
  font: 10px sans-serif;
}
.axis path,
.axis line {
  fill: none;
  stroke: #000;
}
.x.axis path {
  display: none;
}
</style>
<div id="chart_d3"/>
"""))

sub_est_2012_df_by_state_template = jinja2.Template(
"""
// Based on http://bl.ocks.org/mbostock/3885304

require(["d3"], function(d3) {
    var data = []

    {% for row in data %}
    data.push({ 'state': '{{ row[4] }}', 'population': {{ row[1] }} });
    {% endfor %}

    d3.select("#chart_d3 svg").remove()

    var margin = {top: 20, right: 20, bottom: 30, left: 40},
        width = 800 - margin.left - margin.right,
        height = 400 - margin.top - margin.bottom;

    var x = d3.scale.ordinal()
        .rangeRoundBands([0, width], .25);

    var y = d3.scale.linear()
        .range([height, 0]);

    var xAxis = d3.svg.axis()
        .scale(x)
        .orient("bottom");

    var yAxis = d3.svg.axis()
        .scale(y)
        .orient("left")
        .ticks(10)
        .tickFormat(d3.format('.1s'));
        
    var svg = d3.select("#chart_d3").append("svg")
        .attr("width", width + margin.left + margin.right)
        .attr("height", height + margin.top + margin.bottom)
        .append("g")
        .attr("transform", "translate(" + margin.left + "," + margin.top + ")");

    x.domain(data.map(function(d) { return d.state; }));
    y.domain([0, d3.max(data, function(d) { return d.population; })]);

    svg.append("g")
        .attr("class", "x axis")
        .attr("transform", "translate(0," + height + ")")
        .call(xAxis);

    svg.append("g")
        .attr("class", "y axis")
        .call(yAxis)
        .append("text")
        .attr("transform", "rotate(-90)")
        .attr("y", 6)
        .attr("dy", ".71em")
        .style("text-anchor", "end")
        .text("Population");

    svg.selectAll(".bar")
        .data(data)
        .enter().append("rect")
        .attr("class", "bar")
        .attr("x", function(d) { return x(d.state); })
        .attr("width", x.rangeBand())
        .attr("y", function(d) { return y(d.population); })
        .attr("height", function(d) { return height - y(d.population); });
});
"""
)
display(Javascript(sub_est_2012_df_by_state_template.render(
    data=sub_est_2012_df_by_state.sort(['POPESTIMATE2012'], ascending=False)[:5].itertuples()))
)

class MultipleSelectWidget(widgets.DOMWidget):
    _view_name = Unicode('MultipleSelectView', sync=True)
    
    value = List(sync=True)
    values = Dict(sync=True)
    values_order = List(sync=True)
    description = Unicode(sync=True)

    def __init__(self, *args, **kwargs):
        self.value_lock = Lock()

        self.values = kwargs.get('values', [])
        self.value = kwargs.get('value', [])
        self.values_order = kwargs.get('values_order', [])
            
        widgets.DOMWidget.__init__(self, *args, **kwargs)

%%javascript
require(["widgets/js/widget"], function(WidgetManager){
    var MultipleSelectView = IPython.DOMWidgetView.extend({
        initialize: function(parameters) {
            this.model.on('change',this.update,this);
            this.options = parameters.options;
            this.child_views = [];
            // I had to override DOMWidgetView's initialize to set model.views otherwise
            // multiple views would get attached to the model
            this.model.views = [this];
        },
        
        render : function(){
            this.$el
                .addClass('widget-hbox');
            this.$label = $('<div />')
                .appendTo(this.$el)
                .addClass('widget-hlabel')
                .hide();
            this.$listbox = $('<select/>')
                .addClass('widget-listbox')
                .attr('multiple', '')
                .attr('size', 6)
                .appendTo(this.$el);
            this.$el_to_style = this.$listbox;
            this.update();
        },
        
        update : function(options){
            if (typeof(options) === 'undefined' || options.updated_view != this) {
                var values = this.model.get('values');
                var values_order = this.model.get('values_order');
                
                var that = this;
                _.each(values_order, function(key, index) {
                    if (that.$listbox.find('option[key="' + key + '"]').length === 0) {
                        $('<option />')
                            .text(values[key])
                            .attr('key', key)
                            .appendTo(that.$listbox)
                            .on('click', $.proxy(that.handle_click, that));
                    }                    
                });
                
                var value = this.model.get('value') || [];
                
                this.$listbox.find('option').each(function(index, element) {
                    var key = $(element).attr('key');
                    
                    if (key in values) {
                        if (value.indexOf(key) != -1) {
                            $(element).prop('selected', true);
                        }
                    } else {
                        $(element).remove();
                    }
                });

                var description = this.model.get('description');
                if (description.length === 0) {
                    this.$label.hide();
                } else {
                    this.$label.text(description);
                    this.$label.show();
                }
            }
            return MultipleSelectView.__super__.update.apply(this);
        },

        handle_click: function (event) {
            var value = $(event.target).parent().children('option:selected').map(function() { return $(this).attr('key') }).get()
            
            this.model.set('value', value, {updated_view: this});
            this.touch();
        },    
    });
    WidgetManager.register_widget_view('MultipleSelectView', MultipleSelectView);
});

def display_chart_d3(data, show_javascript, states, div):
    sub_est_2012_df_by_state_template = jinja2.Template(
    """
    // Based on http://www.recursion.org/d3-for-mere-mortals/

    require(["d3"], function(d3) {
        var data = []

        {% for row in data %}
        data.push({ 'state': '{{ row[4] }}', 'population': {{ row[1] }} });
        {% endfor %}

        d3.select("#chart_d3_interactive svg").remove()

        var margin = {top: 20, right: 20, bottom: 30, left: 40},
            width = 800 - margin.left - margin.right,
            height = 400 - margin.top - margin.bottom;

        var x = d3.scale.ordinal()
            .rangeRoundBands([0, width], .25);

        var y = d3.scale.linear()
            .range([height, 0]);

        var xAxis = d3.svg.axis()
            .scale(x)
            .orient("bottom");

        var yAxis = d3.svg.axis()
            .scale(y)
            .orient("left")
            .ticks(10)
            .tickFormat(d3.format('.1s'));

        var svg = d3.select("#chart_d3_interactive").append("svg")
            .attr("width", width + margin.left + margin.right)
            .attr("height", height + margin.top + margin.bottom)
            .append("g")
            .attr("transform", "translate(" + margin.left + "," + margin.top + ")");

        x.domain(data.map(function(d) { return d.state; }));
        y.domain([0, d3.max(data, function(d) { return d.population; })]);

        svg.append("g")
            .attr("class", "x axis")
            .attr("transform", "translate(0," + height + ")")
            .call(xAxis);

        svg.append("g")
            .attr("class", "y axis")
            .call(yAxis)
            .append("text")
            .attr("transform", "rotate(-90)")
            .attr("y", 6)
            .attr("dy", ".71em")
            .style("text-anchor", "end")
            .text("Population");

        svg.selectAll(".bar")
            .data(data)
            .enter().append("rect")
            .attr("class", "bar")
            .attr("x", function(d) { return x(d.state); })
            .attr("width", x.rangeBand())
            .attr("y", function(d) { return y(d.population); })
            .attr("height", function(d) { return height - y(d.population); });
    });
    """
    )
    rendered_template = sub_est_2012_df_by_state_template.render(
        data=data[data['STUSAB'].map(lambda v: v in states)].itertuples()
    )
    
    if show_javascript:
        display(widgets.PopupWidget(children=(widgets.HTMLWidget(value='<div style="width:600px; height: 400px; overflow:scroll;"><pre>{}</pre></div>'.format(rendered_template)),)))
        
    display(Javascript(rendered_template))

values = {
    record['STUSAB']: "{0} - {1}".format(record['STUSAB'], record['STATE_NAME']) for record in state[['STUSAB', 'STATE_NAME']].sort('STUSAB').to_dict(outtype='records')
}
i = interact(
    display_chart_d3,
    data=widgets.fixed(sub_est_2012_df_by_state),
    show_javascript=widgets.CheckboxWidget(value=False),
    states=MultipleSelectWidget(
        value=['CA', 'NY'],
        values=values,
        values_order=sorted(values.keys())
    ),
    div=widgets.HTMLWidget(value='<div id="chart_d3_interactive"></div>')
)

%%javascript
require.config({
    paths: {
        chartjs: '//cdnjs.cloudflare.com/ajax/libs/Chart.js/0.2.0/Chart.min'
    }
});

def display_chart_chartjs(sc_est2012_sex, show_javascript, show, div):
    sc_est2012_sex_template = jinja2.Template(
    """
    require(["chartjs"], function() {
        var data = {
            labels : {{ labels }},
            datasets : [
            {
                fillColor : "rgba(220,220,220,0.5)",
                strokeColor : "rgba(220,220,220,1)",
                data : {{ dataset_male }}
            },
            {
                fillColor : "rgba(151,187,205,0.5)",
                strokeColor : "rgba(151,187,205,1)",
                data : {{ dataset_female }}
            }
            ]
        }
    
    
        var ctx = $('#chart_chartjs')[0].getContext('2d');
        new Chart(ctx).Bar(data,{});
    });
    """
    )
    
    if show == 'by_region':
        data_frame = sc_est2012_sex.groupby(['REGION', 'SEX'], as_index=False)[['POPEST2012_CIV']].sum()
        labels = [region_codes[code] for code in data_frame['REGION'].drop_duplicates()]
    elif show == 'by_division':
        data_frame = sc_est2012_sex.groupby(['DIVISION', 'SEX'], as_index=False)[['POPEST2012_CIV']].sum()
        labels = [division_codes[code] for code in data_frame['DIVISION'].drop_duplicates()]
    elif show == 'by_state':
        data_frame = sc_est2012_sex.groupby(['STATE', 'SEX'], as_index=False)[['POPEST2012_CIV']].sum()
        data_frame = pd.merge(data_frame, state, left_on='STATE', right_on='STATE')
        labels = data_frame['STATE_NAME'].drop_duplicates().tolist()

    dataset_male = data_frame[data_frame.SEX == 1]['POPEST2012_CIV'].tolist()
    dataset_female = data_frame[data_frame.SEX == 2]['POPEST2012_CIV'].tolist()

    rendered_template = sc_est2012_sex_template.render(
        labels=dumps(labels),
        dataset_male=dumps(dataset_male),
        dataset_female=dumps(dataset_female),
    )

    if show_javascript:
        display(widgets.PopupWidget(children=(widgets.HTMLWidget(value='<div style="width:600px; height: 400px; overflow:scroll;"><pre>{}</pre></div>'.format(rendered_template)),)))
    
    display(Javascript(rendered_template))

i = interact(
    display_chart_chartjs,
    sc_est2012_sex=widgets.fixed(sc_est2012_sex),
    show_javascript=widgets.CheckboxWidget(value=False),
    show=widgets.DropdownWidget(
        values={'By Region':'by_region', 'By Division': 'by_division', 'By State': 'by_state'},
        value='by_region'
    ),
    div=widgets.HTMLWidget(value='<canvas width=800 height=400 id="chart_chartjs"></canvas>')
)

%%javascript
require.config({
    paths: {
        highcharts: '//cdnjs.cloudflare.com/ajax/libs/highcharts/4.0.1/highcharts'
    }
});

def display_chart_highcharts(sc_est2012_age, show_javascript, show, div):
    sc_est2012_age_template = jinja2.Template(
    """
    require(["highcharts"], function() {
        $('#chart_highcharts').highcharts({
            chart: {
                type: 'column'
            },
            title: {
                text: 'Population By Age'
            },
            subtitle: {
                text: 'Source: www.census.gov'
            },
            xAxis: {
                categories: {{ categories }}
            },
            yAxis: {
                min: 0,
                title: {
                    text: 'Population'
                }
            },
            plotOptions: {
                column: {
                    pointPadding: 0.2,
                    borderWidth: 0
                }
            },
            series: {{ series }}
        });
    });
    """
    )
    
    if show == 'by_region':
        data_frame = sc_est2012_age.groupby(['REGION', 'AGE'], as_index=False)[['POPEST2012_CIV']].sum()
        categories = [region_codes[code] for code in data_frame['REGION'].drop_duplicates()]
    elif show == 'by_division':
        data_frame = sc_est2012_age.groupby(['DIVISION', 'AGE'], as_index=False)[['POPEST2012_CIV']].sum()
        categories = [division_codes[code] for code in data_frame['DIVISION'].drop_duplicates()]
        
    series_names = data_frame['AGE'].drop_duplicates().tolist()
    series = [{
        'name': series_name, 'data': data_frame[data_frame.AGE == series_name]['POPEST2012_CIV'].tolist()
    } for series_name in series_names]
    
    rendered_template = sc_est2012_age_template.render(
        categories=dumps(categories),
        series=dumps(series)
    )

    if show_javascript:
        display(widgets.PopupWidget(children=(widgets.HTMLWidget(value='<div style="width:600px; height: 400px; overflow:scroll;"><pre>{}</pre></div>'.format(rendered_template)),)))

    display(Javascript(rendered_template))

i = interact(
    display_chart_highcharts,
    sc_est2012_age=widgets.fixed(sc_est2012_age),
    show_javascript=widgets.CheckboxWidget(value=False),
    show=widgets.DropdownWidget(
        values={'By Region':'by_region', 'By Division': 'by_division'},
        value='by_region'
    ),
    div=widgets.HTMLWidget(value='<div id="chart_highcharts"></div>')
)