"""add_budget_to_activity_territory

Revision ID: 4d142d0dc132
Revises: be788af9004f
Create Date: 2020-01-20 09:10:20.549283

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '4d142d0dc132'
down_revision = 'be788af9004f'
branch_labels = None
depends_on = None


def upgrade():
    bind = op.get_bind()

    op.add_column('ActivityTerritory', sa.Column('budget_spend', sa.Numeric(precision=12, scale=2), nullable=True))

    query_result = bind.execute(''' SELECT id, budget_spend FROM "Activity" ''').fetchall()

    for activity_id, budget in query_result:
        number_of_territories = bind.execute(
        '''
            SELECT COUNT(*)
            FROM "ActivityTerritory"
            WHERE activity_id = {}
        '''.format(activity_id)
        ).scalar()

        bind.execute(
        '''
            UPDATE "ActivityTerritory"
            SET "budget_spend" = {0}
            WHERE activity_id = {1}
        '''.format(budget / number_of_territories, activity_id)
        )

    op.alter_column('ActivityTerritory', 'budget_spend', nullable=False)


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('ActivityTerritory', 'budget_spend')
    # ### end Alembic commands ###