import pickle

import ipywidgets as widgets
import matplotlib.pyplot as plt
import numpy as np
import qgrid
from IPython.display import display
from cloudmesh.common.console import Console
from cloudmesh.timeseries.source import Data
from cloudmesh.timeseries.transform import Normalize
# noinspection PyUnresolvedReferences
from ipywidgets import *
from pandas import DataFrame
from cloudmesh.timeseries.source import Data110
from cloudmesh.timeseries.source import Source


class City110(Data):

    """
    Usage:

    city110=City110()
    columns = ["State", 'City', "Lat", "Long", "INSURANCE"]
    city110.drop(columns)

    city110.data

    """
    def __init__(self):
        Data.__init__(self)
        city110 = CityData(name="Static_110US_Cities.csv")
        city110.data = city110.data.set_index('FIPS')
        cities = city110.data.index
        city315 = CityData(name="Static_315USCities_Pop.csv")
        city315.data = city315.data[cities]
        city315.data = city315.data.transpose()

        for field in ['Insurance',
                      'pop_density_2010',
                      'poverty_percent',
                      'svi_minority',
                      "svi_overall",
                      "senior_percent",
                      "black_percent",
                      "PVI",
                      "svi_ses",
                      "svi_households",
                      "svi_housing_transport"]:
            city315.data[field] = city110.data[field]

        self.data = city315.data

class CityData(Data):

    def __init__(self, name=None):

        if name is None:
            raise ValueError("name is not specified")

        Data.__init__(self)

        self.name = name

        print(f"Load: {name}")
        self.load(self.name)



class City(Data):

    def __init__(self, name="Static_315USCities_Pop.csv"):
        Data.__init__(self)

        self.name = name

        print(f"Load: {name}")
        self.load(self.name)

    def show(self):
        table = qgrid.show_grid(self.data.transpose())
        return table

    def normalize_all(self):
        columns = \
            ['Population', 'CASTHMA', 'HIGHCHOL',
             'DIABETES', 'OBESITY', 'CANCER', 'STROKE', 'MHLTH', 'CSMOKING',
             'CHOLSCREEN', 'INSURANCE', 'CHD', 'CHECKUP', 'KIDNEY', 'BINGE', 'LPA',
             'ARTHRITIS', 'BPMED', 'PHLTH', 'BPHIGH', 'COPD', 'Nhosp', 'Nbeds',
             'Estbeds', 'Percentseniors', 'Percentblacks', 'Percenthispanics']
        self.data = Normalize.rows(self.data, columns)
        return self

    def plot(self, columns):
        for c in columns:
            plt.plot(self.data.loc[c, :].tolist())

    def location(fips):
        return self.data[fips]['Lat'], self.data[fips]['Lat']


class City2(Data):

    def set_value_columns(self, exclude=['City', 'State']):
        self.value_columns = self.data.index.to_list()
        for x in exclude:
            self.value_columns.remove(x)
        print(self.value_columns)

    def __init__(self, name="Static_110US_Cities.csv"):
        Data.__init__(self)
        if name is not None:
            self.name = name

            print(f"Load: {name}")

            self.load(self.name)
            self.initialize()

    def initialize(self):
        self.data['Nbeds_percent'] = self.data['Nbeds'] / self.data['Population']
        self.data['Population_log'] = np.log(self.data['Population'])
        self.data.index = self.data["FIPS"]
        self.data = self.data.drop(columns=['FIPS'])
        self.data = self.data.transpose()
        self.fips = self.data
        self.static = self.data.transpose()
        self.qgrid_selected_cities_train = None
        self.qgrid_selected_cities_predict = None

        self.qgrid_selected_attributes = None
        self.selected_cities_train = None
        self.selected_cities_predict = None
        self.selected_attributes = None

        self.set_value_columns()

        # self.value_columns = \
        #    ['Population',
        #     'Insurance',
        #     'Nbeds',
        #     'Nbeds_percent',
        #     'senior_percent',
        #     'black_percent',
        #     'pop_density_2010',
        #     'poverty_percent',
        #     'svi_minority',
        #     'svi_overall']

    def select(self, kind="cities"):
        if kind == "train":
            self.qgrid_selected_cities_train = qgrid.show_grid(self.static)
            return self.qgrid_selected_cities_train
        if kind == "predict":
            self.qgrid_selected_cities_predict = qgrid.show_grid(self.static)
            return self.qgrid_selected_cities_predict

        elif kind == "attributes":

            names = list(self.data.index)
            names.remove("City")
            names.remove("State")

            names = DataFrame({"attributes": names})

            names = names.set_index("attributes")
            self.qgrid_selected_attributes = qgrid.show_grid(names)

            return self.qgrid_selected_attributes
        else:
            raise ValueError("kind not supported")

    def selected(self, kind="train"):
        if kind == "train":
            return self.qgrid_selected_cities_train.get_selected_df()
        elif kind == "predict":
            return self.qgrid_selected_cities_predict.get_selected_df()
        elif kind == "attributes":
            return self.qgrid_selected_attributes.get_selected_df()
        else:
            raise ValueError("kind not supported")

    def not_selected(self):
        all_data = self.static.copy
        s = self.selected()
        raise NotImplementedError

    def normalize_all(self):
        self.data = Normalize.columns(self.data.transpose(), which=self.value_columns)
        return self

    def values(self):
        df = self.data.copy()
        df = df[self.value_columns]

        return df

    def plot(self, columns):
        for c in columns:
            plt.plot(self.data.loc[c, :].tolist())

    """
    def normalize_by_column(self, column):
        df = self.data.copy()
        for c in self.value_columns:
            df[c] = df[c] / df[column]
        return df
    """

    def write(self, filename, info=False):

        data = [self.selected("train"),
                self.selected("predict"),
                self.selected("attributes")]
        pickle.dump(data, open(filename, "wb"))
        self.read(filename, info=info)

    def read(self, filename, info=True):
        self.selected_cities_train, self.selected_cities_predict, self.selected_attributes = \
            pickle.load(open(filename, "rb"))
        if info:
            print()
            print("filename =", filename)
            self.experiment()
            print()

    def experiment(self):
        train = self.selected_cities_train.index.to_list()
        predict = self.selected_cities_predict.index.to_list()
        intersection = list(set(train).intersection(predict))
        if len(intersection) > 0:
            print()
            Console.red("Warning: The following elements are in training and prediction")
            Console.red(f"fips_overlap={intersection}")
            print()
        print("fips_train =", train)
        print("fips_predict =", predict)
        print("attributes =", self.selected_attributes.index.to_list())

    def do_selection(self, experiment):
        train = self.select("train")
        predict = self.select("predict")
        attributes = self.select("attributes")
        Console.red(f"\n\nSelect the Cities for the training experiment {experiment}")
        display(train)
        Console.red(f"\n\nSelect the Cities for the prediction experiment {experiment}")
        display(predict)
        Console.red(f"\n\nSelect the Attributes for the experiment {experiment}")
        display(attributes)
        print()
        Console.red(f"Press the save button to save experiment {experiment}")

        button = widgets.Button(description=f"Save")
        output = widgets.Output()
        display(button, output)

        def save_experiment(b):
            self.write(experiment)
            self.read(experiment)
            with output:
                print("Saving ...")
                print()
                print("filename =", experiment)
                self.experiment()
                print()

        # cities.selected_attributes
        # cities.selected_cities

        button.on_click(save_experiment)


class City3(Data):

    def fields(self, names):
        static = self.data.loc[names, :]
        return static

    def set_value_columns(self, exclude=['City', 'State', 'Lat', 'Long']):
        self.value_columns = self.data.index.to_list()
        for x in exclude:
            self.value_columns.remove(x)

    def drop(self, columns):
        self.static = self.static.drop(columns=columns)
        self.data = self.static.transpose()

    def __init__(self):
        City2.__init__(self, name=None)

        self.static = Data110.read_data("Static_110US_Cities.csv")
        self.static = self.static.set_index("FIPS")
        static_315 = Data110.read_data("Static_315USCities_Pop.csv").transpose()

        ignore_in_merge = ['State', 'City', 'Population', 'Nbeds', 'Percentseniors']
        merge = ['CASTHMA', 'HIGHCHOL',
                 'DIABETES', 'OBESITY', 'CANCER', 'STROKE', 'MHLTH', 'CSMOKING',
                 'CHOLSCREEN', 'INSURANCE', 'CHD', 'CHECKUP', 'KIDNEY', 'BINGE', 'LPA',
                 'ARTHRITIS', 'BPMED', 'PHLTH', 'BPHIGH', 'COPD', 'Nhosp', 'Nbeds',
                 'Estbeds', 'Percentseniors', 'Percentblacks', 'Percenthispanics']

        a = set(self.static.index.to_list())
        b = set(static_315.index.to_list())
        intersect = a.intersection(b)

        static_315 = static_315.transpose()[intersect].transpose()

        self.static = self.static.sort_index(axis=0)
        static_315 = static_315.sort_index(axis=0)
        static_315 = static_315.drop(columns=['State', 'City', 'Population', 'Nbeds'])

        self.static = self.static.join(static_315)

        self.static['Nbeds_per1000'] = self.static['Nbeds'] / self.static['Population'] * 1000
        # self.static['Population_log'] = np.log(self.static['Population'])

        # reorder
        self.static = self.static[['State', 'City', 'Lat', 'Long', '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']]
        self.data = self.static.transpose()
        # self.initialize()

        self.set_value_columns()

        # self.value_columns = self.data.index.to_list()

        # for x in ['Long', 'Lat', 'City', 'State']:
        #    self.value_columns.remove(x)
        # self.static2 = self.static2.drop(columns=ignore_in_merge)


class CityDF:

    def set_value_columns(self, exclude=['City', 'State']):
        self.value_columns = self.data.index.to_list()
        for x in exclude:
            self.value_columns.remove(x)
        print(self.value_columns)

    @property
    def static(self):
        return self.data.transpose()

    @property
    def fips(self):
        return self.data

    def __init__(self, df):
        Data.__init__(self)
        self.data = df
        self.table = {
            "train": None,
            "predict": None,
            "attribute": None
        }
        self._selected = {
            "train": None,
            "predict": None,
            "attribute": None
        }

    def select(self, kind="cities"):
        if kind in ["train", "predict"]:
            self.table[kind] = qgrid.show_grid(self.static)
            return self.table[kind]
        elif kind == "attributes":
            names = list(self.data.index)
            names.remove("City")
            names.remove("State")

            names = DataFrame({"attributes": names})

            names = names.set_index("attributes")
            self.table["attributes"] = qgrid.show_grid(names)

            return self.table["attributes"]
        else:
            raise ValueError("kind not supported")

    def selected(self, kind="train"):
        return self.table[kind].get_selected_df()

    def normalize_all(self):
        self.data = Normalize.columns(self.data.transpose(), which=self.value_columns)
        return self

    def plot(self, columns):
        for c in columns:
            plt.plot(self.data.loc[c, :].tolist())

    def write(self, filename, info=False):

        data = [self.selected("train"),
                self.selected("predict"),
                self.selected("attributes")]
        pickle.dump(data, open(filename, "wb"))
        self.read(filename, info=info)

    def read(self, filename, info=True):
        self._selected["train"], self._selected["predict"], self._selected["attributes"] = \
            pickle.load(open(filename, "rb"))
        if info:
            print()
            print("filename =", filename)
            self.experiment()
            print()

    def experiment(self):
        train = self._selected["train"].index.to_list()
        predict = self._selected["predict"].index.to_list()
        intersection = list(set(train).intersection(predict))
        if len(intersection) > 0:
            print()
            Console.red("Warning: The following elements are in training and prediction")
            Console.red(f"fips_overlap={intersection}")
            print()
        print("fips_train =", train)
        print("fips_predict =", predict)
        print("attributes =", self._selected["attributes"].index.to_list())

    def do_selection(self, experiment):
        train = self.select("train")
        predict = self.select("predict")
        attributes = self.select("attributes")
        Console.red(f"\n\nSelect the Cities for the training experiment {experiment}")

        display(train)
        Console.green("please scroll for more ...")
        print()

        Console.red(f"\n\nSelect the Cities for the prediction experiment {experiment}")
        display(predict)
        Console.green("please scroll for more ...")
        print()

        Console.red(f"\n\nSelect the Attributes for the experiment {experiment}")
        display(attributes)
        Console.green("please scroll for more ...")
        print()

        print()
        Console.red(f"Press the save button to save experiment {experiment}")

        button = widgets.Button(description=f"Save")
        output = widgets.Output()
        display(button, output)

        def save_experiment(b):
            self.write(experiment)
            self.read(experiment)
            with output:
                print("Saving ...")
                print()
                print("filename =", experiment)
                self.experiment()
                print()

        # cities.selected_attributes
        # cities.selected_cities

        button.on_click(save_experiment)
