from google.colab import drive
drive.mount('/content/drive')
folder_path = "/content/drive/MyDrive/machine_learning/ps_conv/"

predict_folder_path = "/content/ps_conv/predict/"
!mkdir -p /content/ps_conv/predict/

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import datetime
import math
from sklearn import linear_model

train_raw_df = pd.read_csv(folder_path + "train_data.csv")
test_raw_df = pd.read_csv(folder_path + "test_data.csv")
submission_raw_df = pd.read_csv(folder_path + "submission.csv")

qs = [0.01, 0.1, 0.5, 0.9, 0.99]#予測4分点
products = ['ice1', 'ice2', 'ice3', 'oden1', 'oden2', 'oden3', 'oden4', 'hot1', 'hot2', 'hot3', 'dessert1',
       'dessert2', 'dessert3', 'dessert4', 'dessert5', 'drink1', 'drink2', 'drink3', 'drink4', 'drink5', 'drink6', 'alcol1', 'alcol2', 'alcol3',
       'snack1', 'snack2', 'snack3', 'bento1', 'bento2', 'bento3', 'bento4',  'tild1', 'tild2', 'men1', 'men2', 'men3', 'men4', 'men5', 'men6']

#date追加 trainは1-3月が2022, testは全て2022(曜日みると2018, 2019の方がよかったっぽい)
def dates(is_train):
  date_list = train_raw_df["date"] if is_train else test_raw_df["date"]
  X_date_split = date_list.str.split("/", expand=True)#4/12など
  X_date_split.columns = ["month", "day"]
  X_date_split["month"] = X_date_split["month"].astype(int)
  X_date_split["year"] = "2022"
  if is_train:
    X_date_split.loc[X_date_split["month"]>3, "year"] = "2021"

  X_date_split["month"] = X_date_split["month"].astype(str)  
  X_date_split["day"] = X_date_split["day"].astype(str)
  X_date_split["date"] = X_date_split["year"] + X_date_split["month"].str.zfill(2) + X_date_split["day"].str.zfill(2)
  return pd.to_datetime(X_date_split["date"])

def preprocess(is_train):
  df = train_raw_df.copy() if is_train else test_raw_df.copy()
  df["date"] = dates(is_train=is_train)
  df["day"] = df["date"].dt.day
  df["weekday"] = df["date"].dt.weekday
  df["rain_day"] = np.where(df["rain"]>0,1,0)
  df["sunny_and_rain"] = np.where((df["rain_day"]==0) & (df["rain_day"].shift(-1)==1),1,0)
  df["period_in_week"] = np.where(df["day"] > 7, 1, 0)
  df["period_in_week"] = np.where(df["day"] > 21, 2, df["period_in_week"])
  return df

train_df = preprocess(is_train=True)
test_df = preprocess(is_train=False)

def percentiles_from(mean, std):
  std99 = 2.33
  std90 = 1.29
  adj99 = 1.25 #高めの数値を異常値として比較的削っているので調整（値適当）
  adj90 = 1.23 #高めの数値を異常値として比較的削っているので調整（値適当）
  p_01 = (mean - std99 * std)
  p_10 = (mean - std90 * std)
  p_50 = (mean)
  p_90 = (mean + std90 * std * adj90)
  p_99 = (mean + std99 * std * adj99)
  return [p_01, p_10, p_50, p_90, p_99]
  #p_01, p_99 = stats.norm.interval(alpha=0.98, loc=mean, scale=std)
  #p_10, p_90 = stats.norm.interval(alpha=0.80, loc=mean, scale=std)

def calc_percentiles(sr):
  mean = sr.mean()
  std = sr.std()
  return percentiles_from(mean=mean, std=std)

#曜日ごと計算
def weekday_calc(df, weekday, column):
  df = df[df["weekday"]==weekday]
  return calc_percentiles(df[column])

#処理済みのdfと最小二乗法のx, yを渡すとpercentileが返る
def percentiles_calc_by_linear(df, calc_col, target_col):
    clf = linear_model.LinearRegression()
    clf.fit(df[calc_col].values.reshape(-1,1), df[target_col])
    mn = clf.predict(np.array(row[calc_col]).reshape(1,-1))[0]
    diff = df[target_col] - df[calc_col] * clf.coef_ #これ計算しなくても出た気がする
    return percentiles_from(mean=mn, std=diff.std())

#異常値を削除
def drop_outliers(df, ols):
  ols = [o for o in ols if o in df.index]
  if len(ols) > 0:
    df = df.drop(ols)
  return df

def ice_calc(num, row):
    col = "ice" + str(num)
    outliers = {
        1: [341], 2:[], 3:[], 
    }
    df = train_df.copy()
    df = drop_outliers(df, outliers[num])

    #高温度で曲がるので計算の範囲内に落とす
    df = df[df["highest"] > 8]
    df = df[df["highest"] < 23] #187行

    return percentiles_calc_by_linear(df, "highest", col)

def oden_calc(num, row):
    col = "oden" + str(num)
    outliers = {
        1: [236, 15, 187, 178, 70], 2:[298, 307, 33], 3:[239, 276, 40, 43], 4:[325], 
    }
    df = train_df.copy()
    df = drop_outliers(df, outliers[num])
    df = df[df[col]!=0]

    #温度で角度変わる
    temp_threshold = 10 if num == 2 else 9
    if row["lowest"] > temp_threshold:
      df = df[df["lowest"] > (temp_threshold - 1)] #131行
    else:
      df = df[df["lowest"] <= (temp_threshold + 1)] #122行

    return percentiles_calc_by_linear(df, "lowest", col)

def hot_calc(num, row):
    col = "hot" + str(num)
    df = train_df.copy()
    
    #高温度で曲がるので計算の範囲内に落とす
    df = df[df["lowest"] < 14.5] #182行

    if row["weekday"] in [0,1,4]:
      df = df[df["weekday"].isin([0, 1, 4])]
    else:
      df = df[df["weekday"].isin([2, 3, 5, 6])]

    return percentiles_calc_by_linear(df, "lowest", col)

def dessert_calc(num, row):
  col = "dessert" + str(num)
  df = train_df.copy()
  #16,26日にピークがあり、そこから数日かけて下がっていく
  days = []
  if row["day"] in [16,26]:
    days = [16,26]
  elif row["day"] in [17,27]:
    days = [17,27]
  elif row["day"] in [18,28]:
    days = [18,28]
  elif row["day"] in [19,29]:
    days = [19,29]
  else:
    days = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,20,21,22,23,24,25,30,31]
  
  df = df[df["day"].isin(days)]
  if row["weekday"] in [3,6]:
    df = df[df["weekday"].isin([3,6])]
  else:
    df = df[df["weekday"].isin([0,1,2,4,5])]
  return calc_percentiles(df[col])

def drink1234_calc(num, row):
    col = "drink" + str(num)
    outliers = { #2はもう少し落とせそう
        1: [], 2:[303, 296, 289, 212, 177, 58, 37, 79, 142], 3:[87, 163, 30, 72, 58],  4:[]
    }
    df = train_df.copy()
    df = drop_outliers(df, outliers[num])

    #drink1,2,3はrainで処理を分ける
    if num != 4:
      if row["rain"] > 2:
        df = df[df["rain"]>2]
      else:
        df = df[df["rain"]<=2]

    #15度未満は平均値
    if row["highest"] < 15:
      df = df[df["highest"]<15]
      return calc_percentiles(df[col])
    #15度以上は最小二乗法
    else:
      df = df[df["highest"]>=15]
      return percentiles_calc_by_linear(df, "highest", col)

#'drink5', 'drink6', 
def drink56_calc(num, row):
  col = "drink" + str(num)
  df = train_df.copy()
  df = df[df["highest"]>9.2] #testのhighestが9.3-だから
  shreshold = 2
  if row["rain"] > shreshold:
    df = df[df["rain"] > shreshold]
  else:
    df = df[df["rain"] <= shreshold]
  return calc_percentiles(df[col])

def alcol_calc(num, row):
  col = "alcol" + str(num)
  outliers = {
      1: [246,], 2:[143], 3:[60, 176, 211]
  }
  df = train_df.copy()
  df = drop_outliers(df, outliers[num])
  shreshold = 8
  if row["rain"] > shreshold:
    df = df[df["rain"]>shreshold]
  else:
    df = df[df["rain"]<=shreshold]
  return weekday_calc(df, row["weekday"], col)

def snack_calc(num, row):
  col = "snack" + str(num)
  outliers = {
      1: [268,310,77], 2:[5], 3:[]
  }
  df = train_df.copy()
  df = drop_outliers(df, outliers[num])
  return weekday_calc(df, row["weekday"], col)

def bento_calc(num, row):
  col = "bento" + str(num)
  outliers = {
      1: [133,154,63], 2:[151,235], 3:[144, 280, 266], 4:[295, 27, 20], 
  }
  df = train_df.copy()
  df = drop_outliers(df, outliers[num])
  return weekday_calc(df, row["weekday"], col)

def tild_calc(num, row):
  col = "tild" + str(num)
  outliers = {
      1: [], 2:[35],
  }
  df = train_df.copy()
  df = drop_outliers(df, outliers[num])
  df = df[df["sunny_and_rain"]==row["sunny_and_rain"]]
  if row["sunny_and_rain"] == 1:
    if row["weekday"] == 6: #日数が少ないので合算する
      df = df[df["weekday"]==6]
    else:
      df = df[df["weekday"]!=6]
    return calc_percentiles(df[col])
  else:
    return weekday_calc(df, row["weekday"], col)

def men_calc(num, row):
  col = "men" + str(num)
  outliers = {
      1: [283, 293], 2:[62], 3:[], 4:[], 5:[84], 6:[27, 185],  
  }
  df = train_df.copy()
  df = drop_outliers(df, outliers[num])
  df = df[df["sunny_and_rain"]==row["sunny_and_rain"]]
  if row["sunny_and_rain"] == 1:
    if row["weekday"] in [0, 1, 4]: #日数が少ないので合算する
      df = df[df["weekday"].isin([0, 1, 4])]
    else:
      df = df[df["weekday"].isin([2, 3, 5, 6])]
    return calc_percentiles(df[col])
  else:
    return weekday_calc(df, row["weekday"], col)

results = []

for idx, row in test_df.iterrows():
  row_result = []
  for i in range(1,4):
    row_result.extend(ice_calc(i, row))
  for i in range(1,5):
    row_result.extend(oden_calc(i, row))
  for i in range(1,4):
    row_result.extend(hot_calc(i, row))
  for i in range(1,6):
    row_result.extend(dessert_calc(i, row))
  for i in range(1,5):
    row_result.extend(drink1234_calc(i, row))
  for i in [5,6]:
    row_result.extend(drink56_calc(i, row))
  for i in range(1,4):
    row_result.extend(alcol_calc(i, row))
  for i in range(1,4):
    row_result.extend(snack_calc(i, row))
  for i in range(1,5):
    row_result.extend(bento_calc(i, row))
  for i in range(1,3):
    row_result.extend(tild_calc(i, row))
  for i in range(1,7):
    row_result.extend(men_calc(i, row))
  results.append(row_result)

results = pd.DataFrame(results)
results = results.where(results > 0, 0)
results = results.round()
cols = []
for pr in products:
  cols.extend([pr + "_" + str(q) for q in qs])
results.columns = cols
results

sub = submission_raw_df["id"]
sub = pd.concat([sub, results], axis=1)
sub.to_csv(predict_folder_path + "probspace_conv_simple.csv", header=True, index=None)