"""stats_functions

Revision ID: b2c02f04f5a2
Revises: c195308b054d
Create Date: 2018-03-26 10:47:39.364510

"""

# revision identifiers, used by Alembic.
revision = 'b2c02f04f5a2'
down_revision = 'c195308b054d'
branch_labels = None
depends_on = None

from alembic import op
import sqlalchemy as sa


def upgrade():
    op.execute('''
create type stat as (as_of date, val double PRECISION);
        ''')

    op.execute('''
create type _ts_stat as (as_of TIMESTAMP WITH TIME ZONE, val DOUBLE PRECISION);
        ''')

    op.execute('''
create function cast_ts_stat_as_stat(_ts_stat) returns stat
IMMUTABLE
STRICT
LANGUAGE sql
as $$
  select ($1).as_of::date, ($1).val
$$;
        ''')

    op.execute('''
create cast (_ts_stat as stat) WITH FUNCTION cast_ts_stat_as_stat(_ts_stat)
        ''')



    op.execute('''
create function to_stat_date(date) returns integer
    immutable
    strict
    language sql
as $$
select (extract(epoch from $1) / (24*60*60))::INTEGER
$$
        ''')

    op.execute('''
create function stat_date_to_date(integer) returns date
    immutable
    strict
    language sql
as $$
select to_timestamp($1 * (24*60*60))::date;
$$
        ''')

    op.execute('''
CREATE or replace FUNCTION cast_stat_as_jsonb(stat)
  RETURNS jsonb
IMMUTABLE
STRICT
LANGUAGE SQL
AS $$
SELECT jsonb_build_array(to_stat_date(($1).as_of), ($1).val);
$$;
        ''')

    op.execute('''
CREATE CAST (stat AS jsonb) with function cast_stat_as_jsonb(stat);
        ''')

    op.execute('''
create function jsonb_to_stat(jsonb) returns stat
	immutable
	strict
	language sql
as $$
select stat_date_to_date(($1->>0)::integer) as as_of, ($1->>1)::float as val;
$$    
        ''')


    op.execute('''
    create function nearest_stat_array(jsonb, date, integer) returns stat
    	immutable
    	strict
    	language sql
    as $$
    select
      jsonb_to_stat(last(r order by (r->>0)::int))
    from jsonb_array_elements($1) r
    where (r->>0)::INTEGER <= to_stat_date($2) 
      and (r->>0)::INTEGER > to_stat_date($2) - $3
    $$
        ''')

    op.execute('''
CREATE OR REPLACE FUNCTION nearest_stat_array(jsonb, int, int)
  returns stat
immutable
strict
language plpgsql
AS $$
DECLARE
  maxdate int;
BEGIN
  maxdate := (select max((r->>0)::int) from jsonb_array_elements($1) r);
  return nearest_stat_array($1, to_timestamp((maxdate-$2)*24*60*60)::date, $3);
END;
$$;
    ''')


def downgrade():
    op.execute('''drop function if exists nearest_stat_array(jsonb, date, integer)''')
    op.execute('''drop function if exists jsonb_to_stat(jsonb)''')
    op.execute('''drop function if exists cast_stat_as_jsonb(stat) cascade''')
    op.execute('''drop function if exists stat_date_to_date(integer) cascade''')
    op.execute('''drop function if exists to_stat_date(date) cascade''')
    op.execute('''drop function if exists cast_ts_stat_as_stat(_ts_stat) cascade''')
    op.execute('''drop type if exists _ts_stat''')
    op.execute('''drop type if exists stat''')