"""Create zips_grouped table and add multicolumn index to zips table

Revision ID: 3b9dffe063a6
Revises: be4dcb9eede6
Create Date: 2020-08-03 11:48:55.068356

"""

# revision identifiers, used by Alembic.
revision = '3b9dffe063a6'
down_revision = 'be4dcb9eede6'
branch_labels = None
depends_on = None

from alembic import op
import sqlalchemy as sa


def upgrade(engine_name):
    globals()["upgrade_%s" % engine_name]()


def downgrade(engine_name):
    globals()["downgrade_%s" % engine_name]()





def upgrade_data_broker():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('zips_grouped',
        sa.Column('created_at', sa.DateTime(), nullable=True),
        sa.Column('updated_at', sa.DateTime(), nullable=True),
        sa.Column('zips_grouped_id', sa.Integer(), nullable=False),
        sa.Column('zip5', sa.Text(), nullable=True),
        sa.Column('state_abbreviation', sa.Text(), nullable=True),
        sa.Column('county_number', sa.Text(), nullable=True),
        sa.Column('congressional_district_no', sa.Text(), nullable=True),
        sa.PrimaryKeyConstraint('zips_grouped_id')
    )
    op.create_index(op.f('ix_zips_grouped_zip5'), 'zips_grouped', ['zip5'], unique=False)
    op.create_index('ix_zips_zip5_state_abbreviation_county_number', 'zips', ['zip5', 'state_abbreviation', 'county_number'], unique=False)

    # Populating the new table
    op.execute("""
        INSERT INTO zips_grouped (zip5, state_abbreviation, county_number)
        SELECT zip5, state_abbreviation, county_number
        FROM zips
        GROUP BY zip5, state_abbreviation, county_number;
        
        WITH district_counts AS (
            SELECT zip5, COUNT(DISTINCT zips.congressional_district_no) AS cd_count
            FROM zips
            GROUP BY zip5)
        UPDATE zips_grouped
        SET created_at = NOW(),
            updated_at = NOW(),
            congressional_district_no = CASE WHEN cd_count <> 1
                                             THEN '90'
                                             END
        FROM district_counts AS dc
        WHERE dc.zip5 = zips_grouped.zip5;
        
        UPDATE zips_grouped
        SET congressional_district_no = zips.congressional_district_no
        FROM zips
        WHERE zips_grouped.congressional_district_no IS NULL
            AND zips.zip5 = zips_grouped.zip5;
    """)
    # ### end Alembic commands ###


def downgrade_data_broker():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index('ix_zips_zip5_state_abbreviation_county_number', table_name='zips')
    op.drop_index(op.f('ix_zips_grouped_zip5'), table_name='zips_grouped')
    op.drop_table('zips_grouped')
    # ### end Alembic commands ###

