#!/usr/bin/env python
# coding: utf-8

"""
Usage:
  predict [--daysin=DAYSIN] [--n=N] [--epochs=EPOCHS] [--output=OUTPUT] [--data=DATADIR] [--fields=FIELDS] [--cuda=OFF]

Options:
  --daysin=DAYSIN    Days as input [default: 3].
  --n=N              number of experiments [default: 10]
  --epochs=EPOCHS    Number of epochs [default: 100]
  --output=OUTPUT    Name of the outputdir [default: output-debug-1]
  --fields=FIELDS    Name of the fields [default: None]
  --data=DATADIR     LOcation of the datadir [default: ./data]
  --cuda=OFF         TRue or False if you like to use cuda [default: False]
"""

from docopt import docopt
import pickle

from multiprocessing import Pool
from keras.models import Sequential
from keras.layers import LSTM
from keras.layers import Dense

from cloudmesh.timeseries.prepare import Prepare
from cloudmesh.timeseries.graphics import Graphics
from cloudmesh.timeseries.result import Result
from cloudmesh.timeseries.result import checkpoint_variables
from cloudmesh.timeseries.result import prediction
from cloudmesh.timeseries.result import supress_keras_warnings

import pandas as pd
import numpy as np
from cloudmesh.common.console import Console
import sys
from cloudmesh.common.util import banner
from cloudmesh.timeseries.source import Source
from tqdm import tqdm
from tqdm.keras import TqdmCallback

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = "2"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

if __name__ == '__main__':
    arguments = docopt(__doc__, version='predict v1.0')
    if arguments['--fields'] in ['None', "", "[]"]:
        arguments["--fields"] = []
    else:
        arguments["--fields"] = arguments["--fields"].split(",")
    if arguments["--cuda"] in ["False", "0", "false"]:
        arguments["--cuda"] = False

    print(arguments)

    if not arguments["--cuda"]:
        os.environ["CUDA_VISIBLE_DEVICES"] = ""
    else:
        print ("CUDA not supportsd in this script")
        sys.exit(1)

    Source.directory = arguments["--data"]

    Graphics.scale(1)
    Graphics.face()
    Console.init()

    supress_keras_warnings()

    debug = True
    start_date = "2020-03-09"
    end_date = "2020-05-25"
    dropout = 0.0   # not yet usedactivations = ["tanh"]
    activations = ["tanh"]
    number_of_workers = int(arguments["--n"])

    n_nodes = 32
    batch_size = 110
    epochs = int(arguments["--epochs"])
    verbose = 0
    # step_length_range=range(1, 15)
    d = int(arguments["--daysin"])
    step_length_range = range(d, d + 1)

    checkpoint = True

    logdir = "log-test"
    os.system(f"rm -rf {logdir}")

    output = arguments["--output"]

    os.system(f"mkdir -p {output}")
    os.system(f"mkdir -p {output}/individual")
    os.system(f"mkdir -p {output}/cumulative")

    field_tests = [
        arguments["--fields"]
        #    ["CHD", "CSMOKING", "CANCER"]
        #    ["CHD", "CSMOKING", "CANCER", "svi_minority", "Estbeds"]
        #    [],
        #    ['Population'],
        #    ['Insurance'],
        #    ['Nbeds'],
        #    ['Nbeds_per1000'],
        #    ['Nhosp'],
        #    ['Estbeds'],
        #    ['senior_percent'],
        #    ['black_percent'],
        #    ['Percentblacks'],
        #    ['Percenthispanics'],
        #    ['pop_density_2010'],
        #    ['poverty_percent'],
        #    ['svi_minority'],
        #    ['svi_overall'],
        #    ['CASTHMA'],
        #    ['HIGHCHOL'],
        #    ['DIABETES'],
        #    ['OBESITY'],
        #    ['CANCER'],
        #    ['STROKE'],
        #    ['MHLTH'],
        #    ['CSMOKING'],
        #    ['CHOLSCREEN'],
        #    ['INSURANCE'],
        #    ['CHD'],
        #    ['CHECKUP'],
        #    ['KIDNEY'],
        #    ['BINGE'],
        #    ['LPA'],
        #    ['ARTHRITIS'],
        #    ['BPMED'],
        #    ['PHLTH'],
        #    ['BPHIGH'],
        #    ['COPD']
    ]

    # data = Prepare(source="magic", start_date=start_date, end_date=end_date)
    data = Prepare(start_date=start_date, end_date=end_date)

    data.restrict_cities(kind="death")

    print(data.cases.data.columns.shape)
    print(data.death.data.columns.shape)
    print(data.city.data.columns.shape)

    data.normalize()
    data.normalize_city()


    N = len(field_tests) * len(step_length_range) * number_of_workers
    N_pbar = tqdm(total=N)

    # In[ ]:

    for fields in field_tests:

        in_fields = fields + ['Cases', 'Death']
        out_fields = ['Cases', 'Death']

        print(in_fields)


        def lstm_run(i, data, in_fields, out_fields, n_steps_in, n_nodes, epochs, batch_size, verbose):

            np.random.seed(int.from_bytes(os.urandom(4), byteorder='little'))
            print (f" {i} prepare samples")
            X_predict, y_predict = data.samples(in_fields, out_fields, n_steps_in=n_steps_in)
            print(f"{i} prepare samplesdone")

            X_train, y_train = data.shuffle()

            n_timesteps, n_features, n_outputs = X_train.shape[1], X_train.shape[2], y_train.shape[1]

            if debug:
                print(i,
                      "n_timesteps:", n_timesteps,
                      "n_features:", n_features,
                      "n_outputs:", n_outputs,
                      "batch_size:", batch_size,
                      "n_nodes:", n_nodes,
                      "epochs:", epochs)

            # define model
            model = Sequential()

            model.add(LSTM(n_nodes,
                           activation='tanh',
                           input_shape=(n_timesteps, n_features)))
            # model.add(Dense(int(n_nodes/2), activation= 'relu' ))
            model.add(Dense(n_outputs))
            model.compile(loss='mse', optimizer='adam')

            print (f"{i} model compile done")
            # model.summary()
            if debug:
                the_callbacks = [
                    TqdmCallback(),
                ]

                # fit network
                model.fit(X_train, y_train,
                          epochs=epochs,
                          batch_size=batch_size,
                          verbose=verbose,
                          # validation_split=0.2,
                          callbacks=the_callbacks
                          )
            else:
                # fit network
                model.fit(X_train, y_train,
                          epochs=epochs,
                          batch_size=batch_size,
                          verbose=verbose,
                          # validation_split=0.2,
                          # callbacks=the_callbacks
                          )
            print (f"{i} fit done")


            X = X_predict
            y = y_predict

            print(f"{i} predict")
            result, rmse_cases, rmse_death = prediction(X, y, model, n_timesteps, n_features, out_fields, verbose)
            length = int(len(X) / batch_size)
            print(f"{i} predict done")

            sfields = "-".join(in_fields)

            c = Result(data, result, length, n_timesteps, "Cases", fields=sfields, i=i,
                       directory=f"{output}/individual/{sfields}")
            d = Result(data, result, length, n_timesteps, "Death", fields=sfields, i=i,
                       directory=f"{output}/individual/{sfields}")

            print(f"{i} result done")

            s_cases = c.summarize_error()
            s_death = d.summarize_error()
            if debug:
                Console.ok(f"Output in={str(in_fields)}, out: {str(out_fields)}")
                print("  RMSE cases:             ", round(rmse_cases, 4))
                print("  RMSE death:             ", round(rmse_death, 4))
                print("  Cumulative error cases: ", s_cases)
                print("  Cumulative error death: ", s_death)
                c.summarize_plot(dots=False)
                d.summarize_plot(dots=False)


            if checkpoint:
                sfields = "-".join(in_fields)
                print(f"{i} checkpoint {output}/individual/{sfields}/{n_steps_in}")

                checkpoint_variables(filename=f"{output}/individual/{sfields}/{n_steps_in}",
                                     model=model,
                                     variables=locals())
                print(f"{i} checkpoint done")

            global N
            N = N - 1
            return rmse_cases, rmse_death, s_cases, s_death


        total = None

        for n_steps_in in step_length_range:

            banner(f"days in {n_steps_in}")
            sfields = str("-".join(in_fields))

            if debug:
                banner(f"Length={n_steps_in} left={N} fields={sfields}")
            N_pbar.update(number_of_workers)
            parameters = []

            for r in range(0, number_of_workers):
                #
                # PREPARE PREDICT SAMPLES
                #
                parameter = (
                    r,
                    data,
                    in_fields,
                    out_fields,
                    n_steps_in,
                    n_nodes,
                    epochs,
                    batch_size,
                    verbose
                )
                parameters.append(parameter)

            with Pool(number_of_workers) as p:
                results = p.starmap(lstm_run, parameters)
                p.close()
                p.join()

            N = N - number_of_workers

            df = pd.DataFrame(results)
            df.columns = ["rmse_cases", "rmse_death", "cum_error_cases", "cum_error_death"]
            df["n_steps_in"] = n_steps_in
            # f = list(in_fields)
            # f.remove('Cases')
            # f.remove('Death')
            df['fields'] = sfields
            # if debug:
            #    display(df);
            if total is None:
                total = df.copy()
            else:
                frames = [total, df]
                total = pd.concat(frames)
            if checkpoint:
                os.system(f"mkdir -p {output}/cumulative")
                pickle.dump(total, open(f"{output}/cumulative/{sfields}-{n_steps_in}.pkl", "wb"))
            # if debug:
            #    display(total);
        # cleanup intermediate files
        for n_steps_in in step_length_range[:-1]:
            os.system(f"rm {output}/cumulative/{sfields}-{n_steps_in}.pkl")

    # for i in range(0,number_of_workers):
    #    banner(f"run {i}")
    #    rmse_cases, rmse_death,  s_cases, s_death = lstm_run(data, in_fields, out_fields,
    #    n_steps_in, n_nodes, epochs, batch_size, verbose)

    N_pbar.close()
