import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import pickle
from tensorflow import keras

from cloudmesh.timeseries.metric import Metric

def supress_keras_warnings():
    import tensorflow as tf
    import logging
    logger = tf.get_logger()
    logger.setLevel(logging.ERROR)

def get_variables(names=None):
    """
    returns all the variables as list from locals() with that name

    checkpoint_variables(..., variable=get_variables(names=[' ', ' '])
    """
    raise NotImplementedError

def checkpoint_variables(filename=None,
                         variables=None,
                         model=None,
                         exclude=["the_callbacks", 'model', 'TqdmCallback']):
    directory = os.path.dirname(filename)
    os.system(f"mkdir -p {directory}")
    if model is not None:
        model.save(f"{filename}.mdl")
    for key in exclude:
        if key in variables:
            del variables[key]
    pickle.dump(variables, open(f"{filename}.pkl", "wb"))

def load_checkpoint_variables(filename=None):
    directory = os.path.dirname(filename)
    variables = pickle.load(open(f"{filename}.pkl", "rb"))

    model_file = f"{filename}.mdl"

    if os.path.exists(filename):
        model = keras.models.load_model('path/to/location')
    else:
        model = None
    return variables, model

def prediction(X, y, model, n_timesteps, n_features, out_fields, verbose=0):
    predict = []
    for i in range(0, len(X)):
        xx = X[i]
        yy = y[i]
        xx = xx.reshape(1, n_timesteps, n_features)
        yhat = model.predict(xx, verbose=verbose)
        v = yhat[0].tolist() + yy.tolist()
        predict.append(v)

    result = pd.DataFrame(predict)

    result.columns = [f"p_{x}" for x in out_fields] + [f"r_{x}" for x in out_fields]

    rmse_cases = ((result.p_Cases - result.r_Cases) ** 2).mean() ** .5
    rmse_death = ((result.p_Death - result.r_Death) ** 2).mean() ** .5

    return result, rmse_cases, rmse_death


class Result:

    def __init__(self, data, result, length, timesteps, kind="Cases", fields="", i=0, directory="."):
        # kind can also be "Death"
        self.data = data
        self.result = result
        self.kind = kind
        self.length = length
        self.timesteps = timesteps
        self.fields = fields
        self.i = i
        self.directory = directory

    def filename(self):
        return f"{self.directory}/{self.timesteps}-{self.length}-{self.kind}-{self.fields}-{self.i}.pdf"

    def index(self, fips):
        return self.data.cities.index(fips)

    def fips(self, i):
        return self.data.cities[i]

    def error_by_index(self, i):
        fips, r = self.select(i)
        pred = f"p_{self.kind}"
        real = f"r_{self.kind}"
        e = Metric.get(r[real], r[pred], i)
        return e

    def select_by_fips(self, fips):
        i = self.index(fips)
        return self.select(i)

    def select(self, i):
        t0 = i * self.length
        t1 = t0 + self.length
        r = self.result[t0:t1].copy()
        r.index = list(range(0, self.length))
        fips = self.data.cities[i]
        return fips, r

    def plot_by_fips(self, fips):
        i = self.data.cities.index(fips)
        self.plot_by_index(i)

    def _save_plot(self, plt, notebook=False):

        if notebook:
            plt.show()

        name = self.filename()
        print(f"save figure ---> {name}")

        directory = os.path.dirname(name)
        os.system(f"mkdir -p {directory}")

        plt.savefig(name)
        plt.close()

    def plot_by_index(self, i, dots=True, fill=True, notebook=False):
        fips, r = self.select(i)
        predict = r[f"p_{self.kind}"]
        real = r[f"r_{self.kind}"]
        # c_error = round(self.error_by_index(i).transpose()["cum_error"].values[0], 2)
        c_error = round(Metric.cumulative_error(predict, real), 2)

        error = np.absolute(real - predict)

        x = list(range(0, len(predict)))
        if dots:
            plt.scatter(x, predict)
            plt.scatter(x, real)
            # plt.scatter(x,error, color="red")

        plt.plot(predict, label=f'{fips} prediction')
        plt.plot(real, label=f'{fips} real')
        plt.plot(error, label=f'{fips} error', color="red")
        if fill:
            plt.fill_between(x, predict, real, alpha=0.1, color="grey")
            plt.fill_between(x, error, alpha=0.05, color="red")

        name = self.data.city.data[fips]["City"]
        population = self.data.city.data[fips]["Population"]
        # l = len(real)

        plt.title(f"{name}, FIPS={fips}, Population={population}, Length={self.timesteps}, Error={c_error}%")
        plt.xlabel("Days")
        plt.ylabel(self.kind)
        plt.grid(True)
        plt.legend()
        self._save_plot(plt, notebook)


    def summarize_plot(self, dots=True, fill=True, notebook=False):
        kind = self.kind
        real = np.zeros(self.length)
        predict = np.zeros(self.length)

        for fips in self.data.cities:
            f, r = self.select_by_fips(fips)
            real = real + r[f"r_{kind}"]
            predict = predict + r[f"p_{kind}"]

        error = np.absolute(real - predict)
        c_error = round(Metric.cumulative_error(predict, real), 2)

        # print (error[0])
        # cum_error = np.zeros(self.length, dtype=float)
        # cum_real = np.zeros(self.length, dtype=float)

        # previous_error = 0.0
        # for i in range(0, self.length):
        #    cum_error[i] = previous_error + error[i]
        #    previous_error = cum_error[i]
        # print (cum_error)

        x = list(range(0, len(predict)))

        if dots:
            plt.scatter(x, predict)
            plt.scatter(x, real)

        plt.plot(predict, label=f'prediction')
        plt.plot(real, label=f'real')
        plt.plot(error, label=f'error', color="red")
        # plt.plot(cum_error, label=f'cum_error', color="red")
        if fill:
            plt.fill_between(x, predict, real, alpha=0.1, color="grey")
            plt.fill_between(x, error, alpha=0.05, color="red")

        plt.title(f"Length={self.timesteps}, Cumulative Results {kind},  error={c_error}%")
        plt.xlabel("Days")
        plt.ylabel(self.kind)
        plt.grid(True)
        plt.legend()
        self._save_plot(plt, notebook)

    def summarize_error(self):
        kind = self.kind
        real = np.zeros(self.length)
        predict = np.zeros(self.length)

        for fips in self.data.cities:
            f, r = self.select_by_fips(fips)
            real = real + r[f"r_{kind}"]
            predict = predict + r[f"p_{kind}"]

        c_error = round(Metric.cumulative_error(predict, real), 2)

        return c_error
