import psycopg2
import pandas as pd
from multiprocessing import Pool
from sqlalchemy import create_engine
import logging


class DbInitiator:

    def __init__(self, config):
        self.config = config
        self.p = Pool(self.config['processors'])
        self.con = psycopg2.connect(database=self.config['database'], user=self.config['user'],
                                    host=self.config['host'])
        self.cur = self.con.cursor()
        self.con_sql_alchemy = create_engine(self.config['url'])

    def __enter__(self):
        return self

    def __exit__(self):
        self.con_sql_alchemy.close()
        self.con.close()

    def init_db(self):
        """Initiates the database. Transform and populate data from all CVS located in input folder"""
        self.cur.execute("""CREATE TABLE IF NOT EXISTS public.teams(
                    team_id int GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
                    team_name text,
                    CONSTRAINT team_name_uniq UNIQUE (team_name)
                    )
                    WITH (OIDS = FALSE);
                    ALTER TABLE public.teams OWNER TO footdata;
                    """
                         )

        self.cur.execute("""CREATE TABLE IF NOT EXISTS public.countries(
                    country_id int GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
                    country_name text,
                    CONSTRAINT country_name_uniq UNIQUE (country_name)
                    )
                    WITH (OIDS = FALSE);
                    ALTER TABLE public.countries OWNER TO footdata;
                    """
                         )

        self.cur.execute("""CREATE TABLE IF NOT EXISTS public.leagues(
                    league_id int GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
                    league_name text,
                    CONSTRAINT league_name_uniq UNIQUE (league_name)
                    )
                    WITH (OIDS = FALSE);
                    ALTER TABLE public.leagues OWNER TO footdata;
                    """
                         )

        self.cur.execute("""CREATE TABLE IF NOT EXISTS public.matches(
                    match_id int GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
                    date timestamp without time zone,
                    home_team_id int REFERENCES teams (team_id),
                    away_team_id int REFERENCES teams (team_id),
                    fthg int,
                    ftag int,
                    ftr int,
                    hthg int,
                    htag int,
                    htr int,
                    b365h double precision,
                    b365d double precision,
                    b365a double precision,
                    league_id int REFERENCES leagues (league_id),
                    country_id int REFERENCES countries (country_id),
                    htftr int,
                    CONSTRAINT date_home_uniq UNIQUE (date, home_team_id)
                    )
                    WITH (OIDS = FALSE);
                    ALTER TABLE public.matches OWNER TO footdata;
                    """
                         )

        self.cur.execute(""" CREATE INDEX IF NOT EXISTS match_id ON matches (match_id);
                            CREATE INDEX IF NOT EXISTS home_team_id ON matches (home_team_id);
                            CREATE INDEX IF NOT EXISTS away_team_id ON matches (away_team_id);
                            CREATE INDEX IF NOT EXISTS date ON matches (date);""")
        self.con.commit()

    def init_feature_db(self, count):
        """
        :param count: number of last matches to generate features on
        :return: N/A
        """
        self.cur.execute("""
                        CREATE TABLE IF NOT EXISTS features{0}_matches (
                          match_id    int REFERENCES matches (match_id) ON UPDATE CASCADE ON DELETE CASCADE, 
                          CONSTRAINT match_side_previous_{0}_pkey PRIMARY KEY (match_id, previous_id, side)
                        );""".format(count))
        self.cur.execute('CREATE INDEX IF NOT EXISTS match_side_{0}_index ON matches_{0}_previous (match_id, side);'.
                         format(count))
        self.con.commit()


class Dao:
    def __init__(self, config):
        self.config = config
        self.p = Pool(self.config['processors'])
        self.con = psycopg2.connect(database=self.config['database'], user=self.config['user'],
                                    host=self.config['host'])
        self.cur = self.con.cursor()
        self.con_sql_alchemy = create_engine(self.config['url'])

    def __enter__(self):
        return self

    def __exit__(self):
        self.con_sql_alchemy.close()
        self.con.close()

    def get_countries(self):
        return pd.read_sql_query('SELECT * FROM countries;', self.con, index_col='country_id')

    def get_countries_unique(self, matches):
        return self.clean_df_db_duplicates(matches[['country_name']].drop_duplicates(), "countries",
                                           dup_cols=["country_name"])

    def write_countries(self, countries):
        countries.to_sql('countries', self.con_sql_alchemy, index_label='country_id', if_exists='append')

    def get_leagues(self):
        return pd.read_sql_query('SELECT * FROM leagues;', self.con, index_col='league_id')

    def get_leagues_unique(self, matches):
        return self.clean_df_db_duplicates(matches[['league_name']].drop_duplicates(), "leagues",
                                           dup_cols=["league_name"])

    def write_leagues(self, leagues):
        leagues.to_sql('leagues', self.con_sql_alchemy, index_label='league_id', if_exists='append')

    def get_teams(self):
        return pd.read_sql_query('SELECT * FROM teams;', self.con, index_col='team_id')

    def get_teams_unique(self, teams):
        return self.clean_df_db_duplicates(teams, "teams", dup_cols=["team_name"])

    def write_teams(self, teams):
        teams.to_sql('teams', self.con_sql_alchemy, index_label='team_id', if_exists='append')

    def get_matches(self):
        return pd.read_sql_query('SELECT * FROM matches;', self.con, index_col='match_id')

    def get_matches_unique(self, matches):
        return self.clean_df_db_duplicates(matches, "matches", dup_cols=["date", "home_team_id"], log=True)

    def write_matches(self, matches):
        matches.to_sql('matches', self.con_sql_alchemy, index_label='match_id', if_exists='append')

    def write_features(self, features, table_name):
        features.to_sql(table_name, self.con_sql_alchemy, index_label='match_id', if_exists='append')

    def get_matches_without_features(self, table_name):

        exists_query = """SELECT EXISTS (SELECT relname FROM pg_class WHERE relname = %(table_name)s);"""
        if pd.read_sql_query(exists_query, self.con_sql_alchemy, params={"table_name": table_name}).loc[0]['exists']:

            matches = self.get_matches()

            return self.clean_df_db_duplicates(matches, table_name, dup_cols=["match_id"], log=True)

        else:
            return self.get_matches()

    def get_features(self, table_name):
        return pd.read_sql(table_name, self.con_sql_alchemy, index_col='match_id')

    def get_previous_matches(self, team_id, date, count):
        """
        Get last &count matches of both teams.
        """

        query = """SELECT * FROM matches 
        WHERE (home_team_id = %(team_id)s OR away_team_id = %(team_id)s) AND %(date)s > matches.date
        ORDER BY date DESC 
        LIMIT %(count)s;"""

        params = {'team_id': team_id, 'date': date, 'count': count}

        return pd.read_sql_query(query, self.con, index_col='match_id', params=params)

    def clean_df_db_duplicates(self, df, table_name, dup_cols=[],
                               filter_continuous_col=None, filter_categorical_col=None, log=False):
        """
        Remove rows from a dataframe that already exist in a database
        Required:
            df : dataframe to remove duplicate rows from
            engine: SQLAlchemy engine object
            tablename: tablename to check duplicates in
            dup_cols: list or tuple of column names to check for duplicate row values
        Optional:
            filter_continuous_col: the name of the continuous data column for BETWEEEN min/max filter
                                   can be either a datetime, int, or float data type
                                   useful for restricting the database table size to check
            filter_categorical_col : the name of the categorical data column for Where = value check
                                     Creates an "IN ()" check on the unique values in this column
        Returns
            Unique list of values from dataframe compared to database table
        """
        args = 'SELECT %s FROM %s' % (', '.join(['"{0}"'.format(col) for col in dup_cols]), table_name)
        args_contain_filter, args_cat_filter = None, None
        if filter_continuous_col is not None:
            if df[filter_continuous_col].dtype == 'datetime64[ns]':
                args_contain_filter = """ "%s" BETWEEN Convert(datetime, '%s')
                                                  AND Convert(datetime, '%s')""" % (filter_continuous_col,
                                                                                    df[filter_continuous_col].min(),
                                                                                    df[filter_continuous_col].max())

        if filter_categorical_col is not None:
            args_cat_filter = ' "%s" in(%s)' % (filter_categorical_col,
                                                ', '.join(["'{0}'".format(value) for value in
                                                           df[filter_categorical_col].unique()]))

        if args_contain_filter and args_cat_filter:
            args += ' Where ' + args_contain_filter + ' AND' + args_cat_filter
        elif args_contain_filter:
            args += ' Where ' + args_contain_filter
        elif args_cat_filter:
            args += ' Where ' + args_cat_filter

        df.drop_duplicates(dup_cols, keep='last', inplace=True)

        database_data = pd.read_sql(args, self.con_sql_alchemy)

        if len(database_data.index) > 0:
            df = pd.merge(df, pd.read_sql(args, self.con_sql_alchemy), how='left', on=dup_cols, indicator=True)
        else:
            return df

        if log:
            df_with_error_records = df[df['_merge'] == 'both'][dup_cols]
            if df_with_error_records.size > 0:
                pd.set_option('display.max_rows', None)
                pd.set_option('display.max_columns', None)
                logging.warning("The following records already have duplicates in table {0} on columns {1}: \n {2}".
                                format(table_name, dup_cols, df[df['_merge'] == 'both'][dup_cols].to_string))
                pd.reset_option('display.max_rows')
                pd.reset_option('display.max_columns')

        df = df[df['_merge'] == 'left_only']
        df.drop(['_merge'], axis=1, inplace=True)
        return df
