import sys
import cv2
import random
import time
import pandas as pd
import pickle
import numpy as np

from pathlib import Path
HERE = Path(__file__).parent.absolute()

# <---- 'train.ipynb' で学習方法を変えた場合は
#        predict() 内にもコピペしてください
def predict(eeg_data):

    MODEL_FILE = HERE / "data" / "models" / "00000000000000_0000.sav"
    # <---- 'train.ipynb' で保存したモデルのファイル

    # <---- 以降，'train.ipynb' と同様です
    #       （eeg_data の形式も，'train.ipynb' 内の eeg_data と同じものです）

    print(eeg_data)

    dataset = pd.DataFrame(
        columns = [
            "0_beta", "1_beta", "2_beta", "3_beta",
            "4_beta", "5_beta", "6_beta", "7_beta",
            "0_alpha", "1_alpha", "2_alpha", "3_alpha",
            "4_alpha", "5_alpha", "6_alpha", "7_alpha",
            "0_theta", "1_theta", "2_theta", "3_theta",
            "4_theta", "5_theta", "6_theta", "7_theta",
        ]
    )

    N = len(eeg_data)
    DT = 1 / 250

    F = np.fft.fft(eeg_data)
    F_abs = np.abs(F)
    freq = np.fft.fftfreq(N, d = DT)

    spectrum = pd.DataFrame(
        data = F_abs,
        columns = ["0", "1", "2", "3", "4", "5", "6", "7"],
        index = freq
    )
    spectrum.index.name = "frequency"

    spectrum = spectrum.iloc[:int(N/2)]
    spectrum = np.divide(spectrum, spectrum.max().max())

    spectrum_beta = spectrum.query("13 <= frequency < 32")
    spectrum_alpha = spectrum.query("8 <= frequency < 13")
    spectrum_theta = spectrum.query("4 <= frequency < 8")

    beta_mean = spectrum_beta.mean()
    alpha_mean = spectrum_alpha.mean()
    theta_mean = spectrum_theta.mean()

    tmp = beta_mean.tolist() + alpha_mean.tolist() + theta_mean.tolist()
    dataset.loc[ len(dataset) ] = tmp

    dataset = dataset.values

    model = pickle.load(MODEL_FILE.open(mode = "rb"))
    # <---- モデルの読みこみ

    predicted = model.predict(dataset)
    # <---- 読みこんだモデルに，dataset を入れて予測

    return int( predicted[0] )
    # <---- 0 or 1 or 2 の int 型で返るようにしてください

def main():

    input("Make sure that 'stream_eeg.py' is running. [press ENTER to continue...]")
    time.sleep(2)

    hand_to_win = {
        0: 2,
        1: 0,
        2: 1
    }

    code2hand = {
        0: "rock",
        1: "scissors",
        2: "paper"
    }
    
    for i in range(10):

        cv2.namedWindow("img", cv2.WINDOW_AUTOSIZE)

        img = cv2.imread( str(HERE / "src" / "jan.jpg") )
        cv2.imshow("img", img)
        cv2.waitKey(1)

        time.sleep(1)

        img = cv2.imread( str(HERE / "src" / "ken.jpg") )
        cv2.imshow("img", img)
        cv2.waitKey(1)

        time.sleep(1)

        img = cv2.imread( str(HERE / "src" / "blank.jpg") )
        cv2.imshow("img", img)
        cv2.waitKey(1)

        start = time.time()
        time.sleep(3)
        end = time.time()

        tmp = (HERE / "data" / "raw").glob("*.csv")
        tmp = max( [int(f.stem) for f in list(tmp)] )
        eeg_data = pd.read_csv(HERE / "data" / "raw" / f"{tmp}.csv", index_col = "time")

        tmp = eeg_data.query(f"{start} < time < {end}")
        if tmp.iloc[0].var() < tmp.iloc[1].var():
            target_data = tmp.iloc[1::2]
        else:
            target_data = tmp.iloc[0::2]

        predicted_your_hand = int( predict(target_data) )

        img = cv2.imread( str(HERE / "src" / f"{code2hand[hand_to_win[predicted_your_hand]]}.jpg") )
        cv2.imshow("img", img)
        cv2.waitKey(1)

        time.sleep(3)

    cv2.destroyAllWindows()
    input("Kill 'stream_eeg.py'. [press ENTER to continue...]")

if __name__ == "__main__":

    main()