from numpy import linalg, zeros, ones, hstack, asarray, vstack, array, mean, std
import itertools
import matplotlib.pyplot as plt
from datetime import datetime
import pandas as pd
import numpy as np
import matplotlib.dates as mdates
from sklearn.metrics import mean_squared_error
from math import sqrt
import warnings
import copy
import time
warnings.filterwarnings("ignore")

from PVPolyfit import preprocessing as preprocess
from PVPolyfit import utilities
from PVPolyfit import clustering as cluster
from PVPolyfit import kernel


def pvpolyfit(train_df, test_df, Y_tag, xs, I_tag, ghi_tag, cs_tag, highest_num_clusters, highest_degree, kernel_type, Y_high_filter, min_count_per_day, plot_graph = True, graph_type = 'regression', print_info = False):
    pvpoly = PVPolyfit(train_df, test_df, Y_tag, xs, I_tag, ghi_tag, cs_tag, print_info)

    # .prepare has its own filters in place, see docs for information
    pvpoly.prepare(10000, 8)

    rmse_list = []
    std_rmse_list = []
    pvpoly_objects = []
    for i in range(1, highest_num_clusters+1):
        pvpoly_iter = copy.deepcopy(pvpoly)
        try:
            pvpoly_iter.run(num_clusters = i, num_iterations = 1, degrees = list(range(1,highest_degree+1)), kernel_type = kernel_type)
            all_best_dfs, ultimate_days, avg_rmse, std_rmse = pvpoly_iter.evaluate(print_info = print_info)
            rmse_list.append(avg_rmse)
            std_rmse_list.append(std_rmse)
            pvpoly_objects.append(pvpoly_iter)

        except:
            break

    if len(rmse_list) == 0:
        raise Exception("No Output was produced. Check to make sure your inputs follow the documentation.")

    min_idx = np.argmin(rmse_list)

    if print_info:
        print("{} cluster(s) were used.".format(range(1,highest_num_clusters)[min_idx]))

    days_rmses, model_output = pvpoly_objects[min_idx].plot(graph_type = graph_type, print_info = print_info, plot_graph = plot_graph)

    return model_output, days_rmses


class PVPolyfit:
    '''
    .______   ____    ____ .______     ______    __      ____    ____  _______  __  .___________.
    |   _  \  \   \  /   / |   _  \   /  __  \  |  |     \   \  /   / |   ____||  | |           |
    |  |_)  |  \   \/   /  |  |_)  | |  |  |  | |  |      \   \/   /  |  |__   |  | `---|  |----`
    |   ___/    \      /   |   ___/  |  |  |  | |  |       \_    _/   |   __|  |  |     |  |     
    |  |         \    /    |  |      |  `--'  | |  `----.    |  |     |  |     |  |     |  |     
    | _|          \__/     | _|       \______/  |_______|    |__|     |__|     |__|     |__|     

    An object, PVPolyfit, created for the creation of an accurate regression of Output depending on the two covariates

    PARAMETERS
    ----------

        train_df: df
            holds training data with columns and index specified below
        test_df: df
            holds testing data with columns and index specified below
        Y_tag: str
            column name of output tag
        xs: list of str
            list of column names for two covariates
        ghi_tag: str
            column name of GHI input
        cs_tag: str
            column name of clearsky GHI generated by pvlib simulation (link below)

    USER MUST INPUT DF's WITH FOLLOWING COLUMNS:

    |       Description      |  Original Use Case  |   Model Purpose          |
    |------------------------|---------------------|--------------------------|
    | Output,       Y_tag:   | DC Power            | Target for regression    |
    | xs:              x1:   | POA Irradiance      | Covariate for regression |
    |                  x2:   | Ambient Temperature | Covariate for regression |
    | Measured GHI, ghi_tag  | GHI (irradiance)    | Day classification       |
    | PVLib Clearsky, cs_tag | Simulated GHI       | Day classification       |

    PVLib has a good tutorial to generate clearsky data: https://pvlib-python.readthedocs.io/en/stable/generated/pvlib.location.Location.get_clearsky.html
    '''

    def __init__(self, train_df, test_df, Y_tag, xs, I_tag, ghi_tag, cs_tag, print_info):
        self.train_df = train_df
        self.test_df = test_df
        self.Y_tag = Y_tag
        self.xs = xs
        self.I_tag = I_tag
        self.ghi_tag = ghi_tag
        self.cs_tag = cs_tag
        self.print_info = print_info


        self.num_clusters = 0
        self.num_iterations = 0
        self.degrees = []

        self.cut_results = []
        self.test_cut_results = []
        self.ordered_pair_list = []
        self.test_ordered_pair_list = []
        self.combined_test_cut_results = []
        self.test_km_labels = []
        self.avg_rmse = 0


        # all_best_dfs[Degree][Day][Columns: 'Y', 'mins', 'maxs']
        self.all_best_dfs = []

        # model_day_couts[Degree][Best model's [Train Counts, Test Counts]]
        self.model_day_counts = []

        # ultimate_days[Day][i, ind]
        self.ultimate_days = []

    def prepare(self, Y_high_filter, min_count_per_day):
        ''' Preprocess and classify days in DataFrame '''
        self.train_df = preprocess.data_preprocessing(self.train_df, self.xs, self.Y_tag, self.I_tag, self.cs_tag, Y_high_filter, self.print_info)
        classification, k, MF = preprocess.classify_weather_day_GM_Tina(self.train_df, self.cs_tag, self.ghi_tag)

        self.train_df['day_type'] = classification

        # cuts train_df into daily DF's
        # also returns a filtered train_df which cuts out a day if its length is too small
        index_list, day_hour_list, self.cut_results, self.train_df = utilities.find_and_break_days_or_hours(self.train_df, True, min_count_per_day = min_count_per_day, frequency = 'days', print_info = self.print_info)

        middles_dates, hours_kpi = utilities.get_weighted_middle_of_day_and_calculate_float_since_noon(self.cut_results, self.ghi_tag, self.cs_tag)
        ### For each day, compile frequencies
        ### For each day, output the # times each class is triggered
        self.ordered_pair_list, freq_df = cluster.create_conglomerated_vectors_for_clustering_algorithm(self.cut_results, hours_kpi, day_hour_list, self.Y_tag)

        self.test_df = preprocess.data_preprocessing(self.test_df, self.xs, self.Y_tag, self.I_tag, self.cs_tag, Y_high_filter, self.print_info)
        test_classification, test_k, test_MF = preprocess.classify_weather_day_GM_Tina(self.test_df, self.cs_tag, self.ghi_tag)
        self.test_df['day_type'] = test_classification
        test_index_list, test_day_hour_list, self.test_cut_results, self.test_df = utilities.find_and_break_days_or_hours(self.test_df, True, min_count_per_day = min_count_per_day, frequency = 'days', print_info = self.print_info)
        test_middles_dates, test_hours_kpi = utilities.get_weighted_middle_of_day_and_calculate_float_since_noon(self.test_cut_results, self.ghi_tag, self.cs_tag)
        self.test_ordered_pair_list, test_freq_df = cluster.create_conglomerated_vectors_for_clustering_algorithm(self.test_cut_results, test_hours_kpi, test_day_hour_list, self.Y_tag)

    def run(self, num_clusters = 6, num_iterations = 1, degrees = list(range(1,10)), kernel_type = 'polynomial'):
        '''
        Iterates through Degrees
        For each Degree, iterates n times
        Returns best model for each input day

        Parameters:
            num_clusters: int, default 6
                number of clusters used in clustering algorithm, synonymous with number of 'types of days'

            num_iterations: int, default 1
                number of times algorithm loops, indicates volatility of algorithm (usually very small, so default = 1)

            degrees: list of ints
                range of degrees that polynomial kernel iterates through

            kernel_type: str
                type of regression kernel to be used
                OPTIONS: polynomial - a(AB)+
        '''

        self.num_clusters = num_clusters
        self.num_iterations = num_iterations
        self.degrees = degrees
        self.kernel_type = kernel_type

        self.all_best_dfs = []
        self.model_day_counts = []
        for degree in self.degrees:
            P_se_list = []
            combined_P_list = []
            combined_day_counts = []

            # 1. Run the code an n number of times
            for i in range(self.num_iterations):

                # clusters and adds 'model_num' column to cut_results & test_cut_results
                train_kmeans_dfs, test_kmeans_dfs, self.test_km_labels, self.cut_results, self.test_cut_results, train_model_day_count, test_model_day_count = cluster.cluster_ordered_pairs_and_return_df_of_days_in_cluster(self.cut_results, self.test_cut_results, self.ordered_pair_list, self.test_ordered_pair_list, kmeans_num_clusters = self.num_clusters, print_info = self.print_info)
                saved_models = cluster.save_model_for_each_cluster(train_kmeans_dfs, degree, self.Y_tag, self.xs, self.kernel_type)
                self.kmeans_Y_lists = kernel.process_test_data_through_models(test_kmeans_dfs, saved_models, self.test_km_labels, self.xs)

                # 2. For each iteration, save the modelled P and colors (based on model used)
                combined_P_list.append(self.kmeans_Y_lists)
                self.combined_test_cut_results.append(self.test_cut_results)
                combined_day_counts.append([train_model_day_count, test_model_day_count])
                P_se_km = kernel.EvaluateModel(array(self.test_df[self.Y_tag].tolist()),array(self.kmeans_Y_lists)).rmse()
                P_se_list.append(P_se_km)

            # 3. Gather the minimum and maximum for each index, save in two lists
            mins = []
            maxs = []
            for i in range(len(self.test_df.index)):
                min = 9999
                max = -9999
                for j in range(len(combined_P_list)):
                    if (combined_P_list[j][i] < min):
                        min = combined_P_list[j][i]
                    if (combined_P_list[j][i] > max):
                        max = combined_P_list[j][i]
                mins.append(min)
                maxs.append(max)

            best_index = np.argmin(P_se_list)
            best_model = combined_P_list[best_index]
            best_df = pd.DataFrame()

            best_df['Y'] = best_model
            best_df['mins'] = mins
            best_df['maxs'] = maxs

            best_df.index = self.test_df.index
            _,_,dfg,_ = utilities.find_and_break_days_or_hours(best_df, False, min_count_per_day = 0, frequency = 'days', print_info = self.print_info)
            self.all_best_dfs.append(dfg)
            self.model_day_counts.append(combined_day_counts[best_index])

    def evaluate(self, print_info = True):
        '''
            Determine rmse for each day for each degree
            and return index of best model for each day
        '''
        # iterate by day
        all_rmse = []
        self.ultimate_days = []
        for i in range(len(self.all_best_dfs[0])):
            min = 9999
            ind = 0
            # iterate by degree
            for j in range(len(self.all_best_dfs)):
                iterating_rmse = kernel.EvaluateModel(array(self.test_cut_results[i][self.Y_tag].tolist()),array(self.all_best_dfs[j][i]['Y'].tolist())).rmse()
                if abs(iterating_rmse) < abs(min):
                    min = iterating_rmse
                    ind = j
            if print_info:
                print("{} index: {}, degrees len: {}".format(len(self.all_best_dfs), ind, len(self.degrees)))
                print("Day {} chooses degree {} with {}".format(i, self.degrees[ind], min))
            all_rmse.append(min)
            self.ultimate_days.append([i, ind])

        self.avg_rmse = np.array(all_rmse).mean()
        return self.all_best_dfs, self.ultimate_days, self.avg_rmse, np.array(all_rmse).std()

    def plot(self, graph_type = 'regression', print_info = True, plot_graph = False):
        iter_rmses = []
        if graph_type == 'regression':
            colors = ['red', 'blue', 'green', 'orange', 'purple', 'brown', 'gold', 'pink', 'gray', 'cyan', 'darkgreen', 'burleywood', 'lawngreen', 'cornflowerblue']*5
            model_outputs = []
            for i in range(len(self.all_best_dfs[0])):
                model_number = self.test_km_labels[i]
                color = colors[model_number]

                ind = self.ultimate_days[i][1]

                Y_output_daily = self.all_best_dfs[ind][i]['Y'].tolist()
                model_outputs.append(Y_output_daily)
                day_index = self.all_best_dfs[ind][i].index.tolist()
                day_maxes = self.all_best_dfs[ind][i]['maxs'].tolist()
                day_mins = self.all_best_dfs[ind][i]['mins'].tolist()
                day_meas = array(self.test_cut_results[i][self.Y_tag].tolist())
                dt_index = pd.to_datetime(day_index)

                if plot_graph:
                    plt.plot(dt_index, day_meas, 'k')
                    plt.plot(dt_index, Y_output_daily, color)
                    plt.fill_between(dt_index, day_maxes, day_mins, facecolor = color)
                    plt.xlabel("time")
                    plt.ylabel("Watts")
                    plt.xticks(rotation=60)
                    plt.title('Modelled Multiple Day Types (by color)')


                uncer = (array(Y_output_daily)-day_meas)#/(day_meas))
                calc_rmse = sqrt(mean_squared_error(day_meas, array(Y_output_daily)))
                iter_rmses.append(calc_rmse)

                if print_info:
                    print("[{}]:".format(datetime.strptime(day_index[0], '%m/%d/%Y %H:%M:%S %p').strftime('%Y-%m-%d')))
                    print("\trmse: {:.4f}, error: {:.4f}".format(calc_rmse, uncer.mean()))
            if plot_graph:
                plt.show()
        return iter_rmses, model_outputs

    def plot_classification_map(self):
        return

    def model_information(self):
        return
