from django.db.models import F

import pandas as pd
import pickle
from sklearn.preprocessing import OneHotEncoder
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.ensemble import GradientBoostingRegressor

from scheduler.models import *
from common.functions import distanceInKmByGPS


class AttendanceLearner:

    def __init__(self,season_id):
        self.season = Season.objects.filter(id=season_id).first()    
        self.rf_regressor = None
        self.predictions = None
        self.ids = []


    def save(self):
        with open('research/ml/models/{}.mod'.format(self.season.id), 'wb') as f:
            pickle.dump(self.rf_regressor, f)

    def load(self):
        with open('research/ml/models/{}.mod'.format(self.season.id), 'rb') as f:
            self.rf_regressor = pickle.load(f)

    def __get_pivot_tables(self):
        games = Game.objects.filter(season=self.season).exclude(historic_season=None)

        # create dataset
        df = pd.DataFrame.from_records(games.values())

        # remove outliers
        out_fields = ['attendance']
        for field in out_fields:
            q_low = df[field].quantile(0.01)
            q_hi  = df[field].quantile(0.99)
            df = df[(df[field] < q_hi) & (df[field] > q_low)]

        # pivots
        pivot_table = {}
        pivot_table['homeTeam_mean_attendance'] = df.pivot_table('attendance','homeTeam_id',aggfunc='mean')
        pivot_table['homeTeam_max_attendance'] = df.pivot_table('attendance','homeTeam_id',aggfunc='max')

        return pivot_table


    def __prepare_training_data(self):     
        games = Game.objects.filter(season=self.season).exclude(historic_season=None).annotate(
            home=F('homeTeam__shortname'),
            away=F('awayTeam__shortname'),
            home_lat=F('homeTeam__latitude'),
            home_lon=F('homeTeam__longitude'),
            home_attr=F('homeTeam__attractivity'),
            away_lat=F('awayTeam__latitude'),
            away_lon=F('awayTeam__longitude'),
            away_attr=F('awayTeam__attractivity'),
            home_country=F('homeTeam__country'),
            away_country=F('awayTeam__country'),
        )

        # create dataset
        df = pd.DataFrame.from_records(games.values())

        # data cleaning
        df['time'] = df['time'].replace('','0')
        df = df[df['attendance'] != 0]

        # remove outliers
        out_fields = ['attendance']
        for field in out_fields:
            q_low = df[field].quantile(0.01)
            q_hi  = df[field].quantile(0.99)
            df = df[(df[field] < q_hi) & (df[field] > q_low)]

        # pivots
        self.pivot_homeTeam_mean = df.pivot_table('attendance','homeTeam_id',aggfunc='mean')
        self.pivot_homeTeam_max = df.pivot_table('attendance','homeTeam_id',aggfunc='max')

        # add more features
        df['weekday'] = df.apply(lambda r: r['date'].weekday(), axis=1)
        df['day'] = df.apply(lambda r: r['date'].day, axis=1)
        df['month'] = df.apply(lambda r: r['date'].month, axis=1)
        df['year'] = df.apply(lambda r: r['date'].year, axis=1)
        df['distance'] = df.apply(lambda r: distanceInKmByGPS(r['home_lon'],r['home_lat'],r['away_lon'],r['away_lat']), axis=1)
        df['weekend'] = df.apply(lambda r: int(r['weekday'] in [6,7]), axis=1)
        df['winter_season'] = df.apply(lambda r: int(r['month'] in [1,2,3,10,11,12]), axis=1)
        df['home_base'] = df.apply(lambda r: self.pivot_homeTeam_mean.loc[r['homeTeam_id'],'attendance'], axis=1)
        df['stadium_size'] = df.apply(lambda r: self.pivot_homeTeam_max.loc[r['homeTeam_id'],'attendance'], axis=1)
        df['early'] = df.apply(lambda r: r['time'].replace(':','') < "1800", axis=1)
        df['before2010'] = df.apply(lambda r: r['historic_season'].split('-')[0] < "2010", axis=1)

        # one hot encoding
        ohe_fields = ['home_country']

        for field in ohe_fields:
            ohe = OneHotEncoder()
            transformed = ohe.fit_transform(df[[field]])
            df[ohe.categories_[0]] = transformed.toarray()

        # sort label to last index
        cols = list(df.columns)
        cols.append(cols.pop(cols.index('attendance')))
        df = df[cols]

        remove_columns = ['season_id', 'resultEntered', 'reversible', 'reschedule', 'homeGoals', 'awayGoals',
                        'homeGoals2', 'awayGoals2', 'homeGoals3', 'awayGoals3', 'home', 'away', 'date', 'time',
                        'id', 'homeTeam_id', 'awayTeam_id', 'historic_season',
                        'home_country','home_lat','home_lon','away_lat','away_lon','away_country','year']
        feature_cols = list(set(df.columns[:-1]) - set(remove_columns))
        label = 'attendance'

        self.X = df[feature_cols]
        self.y = df[label] 

    def __prepare_prediction_data(self):  
        games = Game.objects.filter(season=self.season,historic_season=None).annotate(
            home=F('homeTeam__shortname'),
            away=F('awayTeam__shortname'),
            home_lat=F('homeTeam__latitude'),
            home_lon=F('homeTeam__longitude'),
            home_attr=F('homeTeam__attractivity'),
            away_lat=F('awayTeam__latitude'),
            away_lon=F('awayTeam__longitude'),
            away_attr=F('awayTeam__attractivity'),
            home_country=F('homeTeam__country'),
            away_country=F('awayTeam__country'),
        )

        # create dataset
        df = pd.DataFrame.from_records(games.values())
        
        self.ids = list(df['id'].values)

        # pivots
        pivot_tables = self.__get_pivot_tables()

        # add more features
        df['weekday'] = df.apply(lambda r: r['date'].weekday(), axis=1)
        df['day'] = df.apply(lambda r: r['date'].day, axis=1)
        df['month'] = df.apply(lambda r: r['date'].month, axis=1)
        df['year'] = df.apply(lambda r: r['date'].year, axis=1)
        df['distance'] = df.apply(lambda r: distanceInKmByGPS(r['home_lon'],r['home_lat'],r['away_lon'],r['away_lat']), axis=1)
        df['weekend'] = df.apply(lambda r: int(r['weekday'] in [6,7]), axis=1)
        df['winter_season'] = df.apply(lambda r: int(r['month'] in [1,2,3,10,11,12]), axis=1)
        df['home_base'] = df.apply(lambda r: pivot_tables['homeTeam_mean_attendance'].loc[r['homeTeam_id'],'attendance'], axis=1)
        df['stadium_size'] = df.apply(lambda r: pivot_tables['homeTeam_max_attendance'].loc[r['homeTeam_id'],'attendance'], axis=1)
        df['early'] = df.apply(lambda r: r['time'].replace(':','') < "1800", axis=1)
        df['before2010'] = df.apply(lambda r: 0, axis=1)

        # one hot encoding
        ohe_fields = ['home_country']

        for field in ohe_fields:
            ohe = OneHotEncoder()
            transformed = ohe.fit_transform(df[[field]])
            df[ohe.categories_[0]] = transformed.toarray()

        # sort label to last index
        cols = list(df.columns)
        cols.append(cols.pop(cols.index('attendance')))
        df = df[cols]

        remove_columns = ['season_id', 'resultEntered', 'reversible', 'reschedule', 'homeGoals', 'awayGoals',
                        'homeGoals2', 'awayGoals2', 'homeGoals3', 'awayGoals3', 'home', 'away', 'date', 'time',
                        'id', 'homeTeam_id', 'awayTeam_id', 'historic_season',
                        'home_country','home_lat','home_lon','away_lat','away_lon','away_country','year']
        feature_cols = list(set(df.columns[:-1]) - set(remove_columns))
        label = 'attendance'

        print (feature_cols)
        self.X = df[feature_cols]
        self.y = df[label] 



    def train(self):
        self.__prepare_training_data()

        X_train, X_test, y_train, y_test = train_test_split(
            self.X, self.y, test_size=0.3, random_state=1)  

        # self.rf_regressor = RandomForestRegressor(n_estimators = 200 ,  random_state = 42)
        self.rf_regressor = GradientBoostingRegressor(n_estimators = 200 ,  random_state = 42)
        self.rf_regressor.fit(X_train,y_train)
        self.score = self.rf_regressor.score(X_test,y_test)
        

    def predict(self):
        self.__prepare_prediction_data()

        print(len(self.ids))
        print(len(self.rf_regressor.predict(self.X)))
        self.predictions = dict(zip(self.ids,self.rf_regressor.predict(self.X)))
        print(self.rf_regressor.predict(self.X))
        print(self.predictions)

    def predict_games(self,games):
        teamObjects = Team.objects.filter(season=self.season ).values()
        dayObjects = Day.objects.filter(season=self.season ).values()

        teams=[]
        t_name={}
        t_country={}
        t_lon = {}
        t_lat = {}
        t_attractivity ={}
        getTeamById={}
        for t in teamObjects:
            # print (t['name'],  t['id'])
            teams.append(t['id'])
            t_name[t['id']]= t['name']
            t_country[t['id']]= t['country']
            t_attractivity[t['id']]=t['attractivity']
            t_lon[t['id']]=t['longitude']
            t_lat[t['id']]=t['latitude']
            getTeamById[t['id']]=t['name']

        distance ={}
        distanceById ={}
        attractivity ={}
        travelDict= self.season.travelDict()
        for t1 in teamObjects:
            for t2 in teamObjects:
                # print (t1['name'],t2['name'],distance[t1['name'],t2['name']] ," -> ", travelDict[t1['id']][t2['id']]['distance'] )
                distance[t1['id'],t2['id']] = travelDict[t1['id']][t2['id']]['distance'] 
                attractivity[t1['id'],t2['id']] = int(100*t1['attractivity']*t2['attractivity']);

        wds= {0:'Mon', 1:'Tue', 2:'Wed', 3:'Thu', 4:'Fri', 5:'Sat', 6:'Sun'}
        dayObjects = Day.objects.filter(season=self.season).values()
        days = [ d['id'] for d in dayObjects ]
        d_round = { d['id'] : d['round'] for d in dayObjects }

        getDateTimeDay={}
        getNiceDay={}
        getWeekDay={}
        getMonth={}
        getRound={}
        days = [ d['id'] for d in dayObjects if d['round']>0 ]
        for d in dayObjects:
            dt = parse(d['day'])
            getDateTimeDay[d['id']] = dt
            getNiceDay[d['id']] = str(dt.strftime('%a, %b %d, %Y'))
            getWeekDay[d['id']] = str(wds[dt.weekday()])
            getMonth[d['id']] = dt.month

        def getFeatureVector(t1,t2,d):
            gm={'homeTeam_id': t1, 'awayTeam_id': t2, 'month': d.month, 
            'home_attractivity': t_attractivity[t1],  'away_attractivity': t_attractivity[t2], 
            'weekday': d.weekday, 'home_lat': t_lat[t1] , 'home_lon': t_lon[t1] , 
            'away_lat': t_lat[t2], 'away_lon': t_lon[t2], 'weekday': d.weekday in [1,2,3,4],
            # 'home_country': t_country[t1], 'away_country':  t_country[t2] ,
             'distance' : distance[t1,t2]}             
            return list(gm.values()) 

        games_train = Game.objects.filter(season=self.season).exclude(historic_season=None)
        self.X= []
        self.y =[]
        for game in games_train:
            self.X.append(getFeatureVector(game.homeTeam.id, game.awayTeam.id,game.date) )
            self.y.append(game.attendance)

        X_train, X_test, y_train, y_test = train_test_split(
            self.X, self.y, test_size=0.3, random_state=1)  

        # self.rf_regressor = RandomForestRegressor(n_estimators = 200 ,  random_state = 42)
        self.rf_regressor = GradientBoostingRegressor(n_estimators = 200 ,  random_state = 42)
        self.rf_regressor.fit(X_train,y_train)
        self.score = self.rf_regressor.score(X_test,y_test)
        print (self.score)

        self.X= []
        self.y =[]
        for (t1,t2,d) in games:
            self.X.append(getFeatureVector(t1,t2,getDateTimeDay[d]) )

        result = dict(zip(games,self.rf_regressor.predict(self.X)))

        for g in games: 
            result[g] = int (result[g])

        return result

            

