import numpy as np
import pandas as pd
import holidays
import warnings
from numba import jit
from copy import copy

warnings.simplefilter(action='ignore', category=pd.errors.PerformanceWarning)

@jit(nopython=True)
def apply_rolling_table(x):
    """

    :param x:
    :return:
    """
    op = x[0, 0]
    max_cl = max(x[:, 1])
    min_cl = min(x[:, 1])
    d = max(abs(op - max_cl) / op, abs(op - min_cl) / op) * 100
    return d

class FeatureCreator:
    """

    """
    def __init__(self, data_as_df):
        """

        :param data_as_df:
        """
        self.eps = 1e-5
        self.df = data_as_df
        self.hours_ma = [20, 30, 40, 50, 60, 70, 80, 100, 125, 150, 175]
        self.hours_rsi = [20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130,
                          140, 150, 160, 175]
        self.hours_macd = [[12, 26, 9], [16, 32, 12], [20, 38, 15],
                           [24, 44, 18], [28, 50, 21], [32, 56, 24],
                           [36, 62, 27], [40, 68, 30], [44, 74, 33],
                           [48, 80, 36], [52, 86, 39], [56, 92, 42],
                           [60, 98, 45], [64, 104, 48], [68, 110, 51],
                           [72, 106, 54], [76, 112, 57], [80, 118, 60], [84, 124, 63]]
        self.hours_atr = [20, 30, 40, 50, 60, 70, 80, 90, 100, 120, 140, 160]
        self.hours_general_points = [80, 120, 160, 200, 240, 280]
        self.hours_trens = [15, 20, 25, 30, 40, 50, 75, 100, 125, 150]
        self.hours_ao = [[5, 34], [10, 44], [15, 54], [20, 64], [25, 74], [30, 84]]
        self.hours_aligator = [[3, 5, 8, 13],
                               [5, 8, 13, 21],
                               [8, 13, 21, 34],
                               [13, 21, 34, 55],
                               [21, 34, 55, 89], ]

    def diff_and_div(self, df):
        """

        :param df:
        :return:
        """
        df['(c-o)/o'] = (df['close'] - df['open']) / df['open']
        df['(l-o)/o'] = (df['low'] - df['open']) / df['open']
        df['(h-o)/o'] = (df['high'] - df['open']) / df['open']

        df['(c-o)/c'] = (df['close'] - df['open']) / df['close']
        df['(c-l)/c'] = (df['close'] - df['low']) / df['close']
        df['(c-h)/c'] = (df['close'] - df['high']) / df['close']

        df['(c-l)/(o-l)'] = (df['close'] - df['low']) / (df['open'] - df['low'] + self.eps)
        df['(h-c)/(h-o)'] = (df['high'] - df['close']) / (df['high'] - df['open'] + self.eps)
        df['(c-o)/(h-l)'] = (df['close'] - df['open']) / (df['high'] - df['low'] + self.eps)

        df['(h-max(o,c))/(min(o,c)-l)'] = (df['high'] - np.max([df['open'], df['close']], axis=0)) / \
                                          (np.min([df['open'], df['close']], axis=0) - df['low'] + self.eps)

        df['(h-max(o,c))/abs(c-o)'] = (df['high'] - np.max([df['open'], df['close']], axis=0)) / \
                                      abs(df['close'] - df['open'] + self.eps)
        df['(min(o,c)-l)/abs(c-o)'] = (np.min([df['open'], df['close']], axis=0) - df['low']) / \
                                      abs(df['close'] - df['open'] + self.eps)

        return df

    @staticmethod
    def gap(df):
        """

        :param df:
        :return:
        """
        df['gap'] = np.append([0], df['open'].to_numpy()[1:] - df['close'].to_numpy()[:-1])
        df['gap/o'] = df['gap'] / df['open']

        return df

    def ma(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_ma:
            df[f'ma_close_{hour}h'] = df['close'].rolling(window=hour, min_periods=1).mean()
            df[f'(ma_close_{hour}h-c)/c'] = (df[f'ma_close_{hour}h'] - df['close']) / df['close']

        df['oc'] = df['close'] - df['open']
        df['m_p'] = (df['open'] + df['close']) / 2 - df['close']
        df['v_p'] = df['high'] - df['low']

        df['percent_oc_c'] = df['oc'] / df['close']
        df['percent_m_p_c'] = df['m_p'] / df['close']
        df['percent_v_p_c'] = df['v_p'] / df['close']
        df['percent_oc_o'] = df['oc'] / df['open']
        df['percent_m_p_o'] = df['m_p'] / df['open']
        df['percent_v_p_o'] = df['v_p'] / df['open']

        for hour in self.hours_ma:
            df[f'ma_oc_{hour}h'] = df['oc'].rolling(window=hour, min_periods=1).mean()
            df[f'ma_percent_oc_{hour}h'] = df['percent_oc_c'].rolling(window=hour, min_periods=1).mean()
            df[f'ma_oc_{hour}h/o'] = df[f'ma_oc_{hour}h'] / df['open']
            df[f'ma_oc_{hour}h/c'] = df[f'ma_oc_{hour}h'] / df['close']

        return df

    def ma_subdata(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_ma:
            df[f'ma_close_s1_{hour}h_sd'] = df['close_s1'].rolling(window=hour, min_periods=1).mean()
            df[f'(ma_close_s1_{hour}h-c)/c_sd'] = (df[f'ma_close_s1_{hour}h_sd'] - df['close_s1']) / (
                        df['close_s1'] + self.eps)
            df[f'ma_close_s2_{hour}h_sd'] = df['close_s2'].rolling(window=hour, min_periods=1).mean()
            df[f'(ma_close_s2_{hour}h-c)/c_sd'] = (df[f'ma_close_s2_{hour}h_sd'] - df['close_s2']) / (
                        df['close_s2'] + self.eps)
            df[f'ma_close_s3_{hour}h_sd'] = df['close_s3'].rolling(window=hour, min_periods=1).mean()
            df[f'(ma_close_s3_{hour}h-c)/c_sd'] = (df[f'ma_close_s3_{hour}h_sd'] - df['close_s3']) / (
                        df['close_s3'] + self.eps)
            df[f'ma_close_s4_{hour}h_sd'] = df['close_s4'].rolling(window=hour, min_periods=1).mean()
            df[f'(ma_close_s4_{hour}h-c)/c_sd'] = (df[f'ma_close_s4_{hour}h_sd'] - df['close_s4']) / (
                        df['close_s4'] + self.eps)

        return df

    def rsi(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_rsi:
            delta = df['close'].diff()
            delta.fillna(0, inplace=True)
            up = delta.clip(lower=0)
            down = -1 * delta.clip(upper=0)
            ema_up = up.ewm(com=hour, min_periods=1, adjust=False).mean()
            ema_down = down.ewm(com=hour, min_periods=1, adjust=False).mean()
            rs = ema_up / (ema_down + self.eps)

            rsi = 100 - (100 / (1 + rs))
            df[f'rsi_{hour}'] = rsi / 100
        return df

    def macd(self, df):
        """

        :param df:
        :return:
        """
        for hour_i, hour_j, hour_k in self.hours_macd:
            exp1 = df['close'].ewm(span=hour_i, min_periods=1, adjust=False).mean()
            exp2 = df['close'].ewm(span=hour_j, min_periods=1, adjust=False).mean()
            macd = exp1 - exp2
            exp3 = macd.ewm(span=hour_k, min_periods=1, adjust=False).mean()
            df[f'macd_{hour_i}_{hour_j}_{hour_k}'] = macd
            df[f'macd_signal_line_{hour_i}_{hour_j}_{hour_k}'] = exp3
            df[f'macd_diff_{hour_i}_{hour_j}_{hour_k}'] = macd - exp3 + self.eps
        return df

    def atr(self, df):
        """

        :param df:
        :return:
        """
        high_low = df['high'] - df['low']
        for hour in self.hours_atr:
            high_close = np.abs(df['high'] - df['close'].shift())
            low_close = np.abs(df['low'] - df['close'].shift())
            ranges = pd.concat([high_low, high_close, low_close], axis=1)
            true_range = np.max(ranges, axis=1)
            atr = true_range.rolling(window=hour, min_periods=1).sum() / hour

            df[f'atr_{hour}'] = atr
        return df

    def general_points(self, df):
        """

        :param df:
        :return:
        """
        df['support_80'] = df['close'].rolling(window=80, min_periods=1).min()
        df['support_80_mean'] = df['close'].rolling(window=80, min_periods=1).mean()
        df['resist_80'] = df['close'].rolling(window=80, min_periods=1).max()
        df['resist_80_mean'] = df['close'].rolling(window=80, min_periods=1).mean()
        df['support_80-close'] = df['support_80'] - df['close'] + self.eps
        df['support_80_mean-close'] = df['support_80_mean'] - df['close'] + self.eps
        df['resist_80-close'] = df['resist_80'] - df['close'] + self.eps
        df['resist_80_mean-close'] = df['resist_80_mean'] - df['close'] + self.eps

        for hour in self.hours_general_points:
            next_hour = hour + 40
            df[f'support_{next_hour}'] = df[f'support_{hour}'].rolling(window=hour, min_periods=1).min()
            df[f'support_{next_hour}_mean'] = df[f'support_{hour}'].rolling(window=hour, min_periods=1).mean()
            df[f'resist_{next_hour}'] = df[f'resist_{hour}'].rolling(window=hour, min_periods=1).max()
            df[f'resist_{next_hour}_mean'] = df[f'resist_{hour}'].rolling(window=hour, min_periods=1).mean()

            if hour >= 80:
                df[f'support_{next_hour}'].fillna(df[f'support_{hour}'], inplace=True)
                df[f'support_{next_hour}_mean'].fillna(df[f'support_{hour}_mean'], inplace=True)
                df[f'resist_{next_hour}'].fillna(df[f'resist_{hour}'], inplace=True)
                df[f'resist_{next_hour}_mean'].fillna(df[f'resist_{hour}_mean'], inplace=True)

            df[f'support_{next_hour}-close'] = df[f'support_{next_hour}'] - df['close'] + self.eps
            df[f'support_{next_hour}_mean-close'] = df[f'support_{next_hour}_mean'] - df['close'] + self.eps
            df[f'resist_{next_hour}-close'] = df[f'resist_{next_hour}'] - df['close'] + self.eps
            df[f'resist_{next_hour}_mean-close'] = df[f'resist_{next_hour}_mean'] - df['close'] + self.eps
        return df

    def trend(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_trens:
            mean_values = []
            std_values = []
            diff_values = []

            for i, _ in enumerate(df['close']):
                idx_start = i // hour * hour
                idx_end = (i // hour + 1) * hour
                idx_end = min(idx_end, len(df['close']))
                if idx_start >= idx_end:
                    mean_values.append(df['close'].iloc[idx_start])
                    std_values.append(std_values[-1])
                else:
                    mean_values.append(df['close'].iloc[idx_start: idx_end].mean())
                    std_values.append(df['close'].iloc[idx_start: idx_end].std())

                if i < hour:
                    diff_values.append(0)
                else:
                    diff_values.append(mean_values[-1] - mean_values[-hour - 1])

            merge_mean_values = []
            trend_mean_values = {df['time'].iloc[0]: df['close'].iloc[0]}
            trend_values = []

            diff_idx_start = 0
            diff_idx_end = 1

            while diff_idx_end < len(diff_values):
                if np.sign(diff_values[diff_idx_start]) == np.sign(diff_values[diff_idx_end]):
                    diff_idx_end += 1
                else:
                    merge_mean_values.extend(
                        [np.mean(mean_values[diff_idx_start: diff_idx_end])] * (diff_idx_end - diff_idx_start))

                    end_trend = mean_values[diff_idx_end - 1]
                    trend_mean_values[df['time'].iloc[diff_idx_end - 1]] = end_trend
                    step_trend = (mean_values[diff_idx_end - 1] -
                                  mean_values[diff_idx_start]) / \
                                 (diff_idx_end - diff_idx_start + self.eps)
                    trend_values.append(mean_values[diff_idx_start])
                    trend_values.extend(
                        [mean_values[diff_idx_start] + step_trend * i for i in range(1, diff_idx_end - diff_idx_start)])

                    diff_idx_start = diff_idx_end
                    diff_idx_end += 1

                if diff_idx_end == len(diff_values):
                    merge_mean_values.extend(
                        [np.mean(mean_values[diff_idx_start: diff_idx_end])] * (diff_idx_end - diff_idx_start))

                    end_trend = mean_values[diff_idx_end - 1]
                    trend_mean_values[df['time'].iloc[-1]] = end_trend
                    step_trend = (mean_values[diff_idx_end - 1] -
                                  mean_values[diff_idx_start]) / \
                                 (diff_idx_end - diff_idx_start + self.eps)
                    trend_values.append(mean_values[diff_idx_start])
                    trend_values.extend(
                        [mean_values[diff_idx_start] + step_trend * i for i in range(1, diff_idx_end - diff_idx_start)])

                    break

            df[f'mean_subtrend_{hour}'] = mean_values
            df[f'std_subtrend_{hour}'] = std_values
            df[f'std_subtrend_{hour}'].fillna(0, inplace=True)
            df[f'trend_values_{hour}'] = trend_values
            df[f'diff_trend_values_{hour}'] = df['close'] - trend_values
        return df

    def ao(self, df):
        """

        :param df:
        :return:
        """
        for hour_mn, hour_mj in self.hours_ao:
            median_price = (df['high'] + df['low']) / 2
            sma_mn = median_price.rolling(window=hour_mn, min_periods=1).mean()
            sma_mj = median_price.rolling(window=hour_mj, min_periods=1).mean()
            ao = sma_mn - sma_mj
            df[f'ao_{hour_mn}_{hour_mj}'] = ao
        return df

    def ac(self, df):
        """

        :param df:
        :return:
        """
        for hour_mn, hour_mj in self.hours_ao:
            median_price = (df['high'] + df['low']) / 2
            sma_mn = median_price.rolling(window=hour_mn, min_periods=1).mean()
            sma_mj = median_price.rolling(window=hour_mj, min_periods=1).mean()
            ao = sma_mn - sma_mj
            sma_ao = ao.rolling(window=hour_mn, min_periods=1).mean()
            df[f'ac_{hour_mn}'] = ao - sma_ao
        return df

    def alligator(self, df):
        """

        :param df:
        :return:
        """
        median_price = (df['high'] + df['low']) / 2
        for h1, h2, h3, h4 in self.hours_aligator:
            sum_lips = median_price.rolling(window=h2, min_periods=1).sum()
            diff_lips = (sum_lips / h2).shift(h1)
            diff_lips.fillna(0, inplace=True)
            smma_lips = (sum_lips - diff_lips + median_price) / h2

            sum_teeth = median_price.rolling(window=h3, min_periods=1).sum()
            diff_teeth = (sum_teeth / h3).shift(h2)
            diff_teeth.fillna(0, inplace=True)
            smma_teeth = (sum_teeth - diff_teeth + median_price) / h3

            sum_jaw = median_price.rolling(window=h4, min_periods=1).sum()
            diff_jaw = (sum_jaw / h4).shift(h3)
            diff_jaw.fillna(0, inplace=True)
            smma_jaw = (sum_jaw - diff_jaw + median_price) / h4

            df[f'lips_{h2}_{h1}'] = smma_lips
            df[f'teeth_{h3}_{h2}'] = smma_teeth
            df[f'jaw_{h4}_{h3}'] = smma_jaw
        return df

    def ema(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_ma:
            df[f'ewa_{hour}'] = df['close'].ewm(span=hour, min_periods=1, adjust=False, ignore_na=False).mean()
        return df

    def brp(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_ma:
            df[f'bears_{hour}'] = df['low'] - df['low'].ewm(span=hour, min_periods=1,
                                                            adjust=False, ignore_na=False).mean()
        return df

    def blp(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_ma:
            df[f'bulls_{hour}'] = df['high'] - df['high'].ewm(span=hour, min_periods=1,
                                                              adjust=False, ignore_na=False).mean()
        return df

    def cci(self, df):
        """

        :param df:
        :return:
        """
        tp = (df['high'] + df['low'] + df['close']) / 3

        for hour in self.hours_ma:
            sma_tp = tp.rolling(window=hour, min_periods=1).mean()
            d = tp - sma_tp
            m = np.abs(d).rolling(window=hour, min_periods=1).mean() * 0.015 + self.eps
            cci = d / m
            df[f'cci_{hour}'] = cci
        return df

    def dem(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_ma:
            demax = df['high'] - df['high'].shift(hour)
            demin = df['low'].shift(hour) - df['low']
            demax.fillna(0, inplace=True)
            demin.fillna(0, inplace=True)
            demax = demax.mask(demax < 0, 0)
            demin = demin.mask(demin < 0, 0)
            ma_demax = demax.rolling(window=hour, min_periods=1).mean() + self.eps
            ma_demin = demin.rolling(window=hour, min_periods=1).mean() + self.eps
            df[f'demark_{hour}'] = ma_demax / (ma_demax + ma_demin)
        return df

    def envlp(self, df):
        """

        :param df:
        :return:
        """
        k = 10
        for hour in self.hours_ma:
            upper_band = df['close'].rolling(window=hour, min_periods=1).mean() * (1 + k / 1000)
            lower_band = df['close'].rolling(window=hour, min_periods=1).mean() * (1 - k / 1000)
            df[f'envlp_up_{hour}'] = upper_band
            df[f'envlp_low_{hour}'] = lower_band
            df[f'diff_nearest_band_{hour}'] = np.min((abs(df['close'] - lower_band), abs(upper_band - df['close'])),
                                                     axis=0)
        return df

    def mf(self, df):
        """

        :param df:
        :return:
        """
        hl = df['high'] - df['low']
        df['mf_0'] = hl / df['volume']
        for hour in self.hours_ma:
            df[f'mf_{hour}'] = hl.rolling(window=hour, min_periods=1).mean() /\
                               df['volume'].rolling(window=hour, min_periods=1).mean()
        return df

    def mfi(self, df):
        """

        :param df:
        :return:
        """
        tp = (df['high'] + df['low'] + df['close']) / 3
        money_flow = tp * df['volume']

        pmf = money_flow.mask(money_flow.diff(1) < 0, 0)
        nmf = money_flow.mask(money_flow.diff(1) > 0, 0)

        for hour in self.hours_ma:
            money_ratio = pmf.rolling(window=hour, min_periods=1).sum() / (
                        nmf.rolling(window=hour, min_periods=1).sum() + self.eps)
            df[f'mfi_{hour}'] = (100 - (100 / (1 + money_ratio))) / 100
        return df

    def obv(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_ma:
            obv_0 = 0
            closes = df['close'].values
            volumes = df['volume'].values
            obv_i = [0] * min(hour, df.shape[0])
            for i in range(min(hour, df.shape[0]), df.shape[0]):
                if closes[i] > closes[i - min(hour, df.shape[0])]:
                    obv_i.append(obv_0 + volumes[i])
                    obv_0 += volumes[i]
                elif closes[i] < closes[i - min(hour, df.shape[0])]:
                    obv_i.append(obv_0 - volumes[i])
                    obv_0 -= volumes[i]
                else:
                    obv_i.append(obv_0)
            df[f'obv_{hour}'] = obv_i
        return df

    def sar(self, df):
        """

        :param df:
        :return:
        """
        acceleration = 0.02
        for hour in self.hours_ma:
            highs = df['high'].values
            sar_i_high = [0] * min(hour, df.shape[0])
            for i in range(min(hour, df.shape[0]), df.shape[0]):
                sar_0_high = (highs[i - min(hour, df.shape[0])] - sar_i_high[
                    i - min(hour, df.shape[0])]) * acceleration + sar_i_high[i - min(hour, df.shape[0])]
                sar_i_high.append(sar_0_high)

            df[f'sar_{hour}'] = sar_i_high
        return df

    def sd(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_ma:
            sd_hour = ((df['close'] - df['close'].rolling(window=hour, min_periods=1).mean()) ** 2)
            df[f'sd_{hour}'] = (sd_hour.rolling(window=hour, min_periods=1).mean()) ** 0.5
        return df

    def so(self, df):
        """

        :param df:
        :return:
        """
        for hour in self.hours_ma:
            k = (df['close'] - df['low'].rolling(window=hour, min_periods=1).min()) / \
                (df['high'].rolling(window=hour,
                                    min_periods=1).max() - df['low'].rolling(window=hour,
                                                                             min_periods=1).min() + self.eps)
            d = k.rolling(window=hour, min_periods=1).mean()
            df[f'so_{hour}'] = d
        return df

    def date_features(self, df):
        """

        :param df:
        :return:
        """
        df['hour_of_day'] = df['time'].apply(lambda x: pd.Timestamp(x).hour)
        df['day_of_week'] = df['time'].apply(lambda x: pd.Timestamp(x).day_of_week + 1)
        df['week_of_year'] = df['time'].apply(lambda x: pd.Timestamp(x).weekofyear)

        df['is_last_hour'] = 0
        df.loc[df['hour_of_day'] == max(df['hour_of_day'].unique()), 'is_last_hour'] = 1
        df['is_last_day'] = 0
        df.loc[df['day_of_week'] == max(df['day_of_week'].unique()), 'is_last_day'] = 1
        df['is_last_week'] = 0
        df.loc[df['week_of_year'] >= 50, 'is_last_week'] = 1

        df['till_holidays'] = df['time'].apply(lambda x: min(
            [(pd.Timestamp(h) - pd.Timestamp(x)).days for h in self.ru_holidays.keys() if
             pd.Timestamp(h) >= pd.Timestamp(x)])).values
        df['holidays_7'] = 0
        df.loc[df['till_holidays'] <= 7, 'holidays_7'] = 1

        return df

    @staticmethod
    def clear_df(df):
        """

        :param df:
        :return:
        """
        df['anomaly_tend'] = df[['open', 'close']].rolling(window=18,
                                                           min_periods=1,
                                                           method='table').apply(apply_rolling_table,
                                                                                raw=True,
                                                                                engine='numba').iloc[:, 0]

        extrimal_intervals = []
        extrimal_dates = df[df['anomaly_tend'] > 15]['time'].dt.date.sort_values().unique()

        i = 0
        while i < len(extrimal_dates):
            start_date = extrimal_dates[i]
            i += 1
            if i >= len(extrimal_dates):
                end_date = extrimal_dates[i - 1]
                extrimal_intervals.append(f'{start_date} - {end_date}')
                break
            while i < len(extrimal_dates) and (
                    pd.to_datetime(extrimal_dates[i]) - pd.to_datetime(extrimal_dates[i - 1])).days <= 3:
                i += 1
            end_date = extrimal_dates[i - 1]

            extrimal_intervals.append(f'{start_date} - {end_date}')

        df.drop(columns=['anomaly_tend'], inplace=True)
        dataframes = []
        d2_ = copy(df)

        for interval in extrimal_intervals:
            start_date, end_date = interval.split(' - ')
            d1_, d2_ = d2_.loc[d2_['time'] < start_date], d2_.loc[d2_['time'] > end_date]
            dataframes.append(d1_)

        dataframes.append(d2_)

        return dataframes

    def generate_feature_one_df(self, df):
        """

        :param df:
        :return:
        """
        df = df.loc[df['time'].dt.hour.between(10, 22)]

        df = self.diff_and_div(df)
        df = self.gap(df)
        df = self.ma(df)
        df = self.ma_subdata(df)
        df = self.rsi(df)
        df = self.macd(df)
        df = self.atr(df)
        df = self.general_points(df)
        df = self.trend(df)

        df = self.ao(df)
        df = self.ac(df)
        df = self.alligator(df)
        df = self.ema(df)
        df = self.brp(df)
        df = self.blp(df)
        df = self.cci(df)
        df = self.dem(df)
        df = self.envlp(df)
        df = self.mf(df)
        df = self.mfi(df)
        df = self.obv(df)
        df = self.sar(df)
        df = self.sd(df)
        df = self.so(df)
        df = self.date_features(df)

        return df

    def generate_feature(self, do_clear=False):
        """

        :param do_clear:
        :return:
        """
        self.ru_holidays = holidays.RU(years=range(self.df['time'].apply(lambda x: pd.Timestamp(x)).dt.year.min(),
                                                   self.df['time'].apply(lambda x: pd.Timestamp(x)).dt.year.max() + 2))

        if do_clear:
            self.df = self.clear_df(self.df)
        else:
            self.df = [self.df]

        self.df_feature = []

        for df in self.df:
            df_f = self.generate_feature_one_df(df)
            self.df_feature.append(df_f)

        self.df = pd.concat(self.df_feature)

        return self.df


# -*- coding: utf-8 -*-
"""The app module, containing the app factory function."""
import logging
import sys

from flask import Flask, render_template, globals, request

from flask_uber_clone import commands, public, rider, driver
from flask_uber_clone.extensions import (
    bcrypt,
    cache,
    csrf_protect,
    db,
    debug_toolbar,
    flask_static_digest,
    login_manager,
    migrate,
    flask_seeder,
    flask_breadcrumbs
)

@login_manager.user_loader
def load_user(user_id):
    """
    This will be used many times like on using current_user
    :param user_id: username
    :return: user or none
    """
    # http://librelist.com/browser/flask/2012/4/7/current-blueprint/#44814417e8289f5f5bb9683d416ee1ee
    if not request.endpoint:
        return None

    blueprint = globals.current_app.blueprints[request.blueprint]

    if hasattr(blueprint, "load_user"):
        return blueprint.load_user(user_id)

    # https://flask-login.readthedocs.org/en/latest/#how-it-works
    return None

def create_app(config_object="flask_uber_clone.settings"):
    """Create application factory, as explained here: http://flask.pocoo.org/docs/patterns/appfactories/.

    :param config_object: The configuration object to use.
    """
    app = Flask(__name__.split(".")[0])
    app.config.from_object(config_object)
    register_extensions(app)
    register_blueprints(app)
    register_errorhandlers(app)
    register_shellcontext(app)
    register_commands(app)
    configure_logger(app)
    return app

def register_extensions(app):
    """Register Flask extensions."""
    bcrypt.init_app(app)
    cache.init_app(app)
    db.init_app(app)
    csrf_protect.init_app(app)
    login_manager.init_app(app)
    debug_toolbar.init_app(app)
    migrate.init_app(app, db)
    flask_static_digest.init_app(app)
    flask_seeder.init_app(app, db)
    flask_breadcrumbs.init_app(app, init_menu=False)
    return None

def register_blueprints(app):
    """Register Flask blueprints."""
    app.register_blueprint(public.views.blueprint)
    app.register_blueprint(rider.views.blueprint, url_prefix="/ride")
    app.register_blueprint(driver.views.blueprint, url_prefix="/drive")
    return None

def register_errorhandlers(app):
    """Register error handlers."""

    def render_error(error):
        """Render error template."""
        # If a HTTPException, pull the `code` attribute; default to 500
        error_code = getattr(error, "code", 500)
        return render_template(f"{error_code}.html"), error_code

    for errcode in [401, 404, 500]:
        app.errorhandler(errcode)(render_error)
    return None

def register_shellcontext(app):
    """Register shell context objects."""

    def shell_context():
        """Shell context objects."""
        return {"db": db}

    app.shell_context_processor(shell_context)

def register_commands(app):
    """Register Click commands."""
    app.cli.add_command(commands.test)
    app.cli.add_command(commands.lint)

def configure_logger(app):
    """Configure loggers."""
    handler = logging.StreamHandler(sys.stdout)
    if not app.logger.handlers:
        app.logger.addHandler(handler)


import codecs
import locale
import re
import sys
from typing import List, Tuple

BOMS: List[Tuple[bytes, str]] = [
    (codecs.BOM_UTF8, "utf-8"),
    (codecs.BOM_UTF16, "utf-16"),
    (codecs.BOM_UTF16_BE, "utf-16-be"),
    (codecs.BOM_UTF16_LE, "utf-16-le"),
    (codecs.BOM_UTF32, "utf-32"),
    (codecs.BOM_UTF32_BE, "utf-32-be"),
    (codecs.BOM_UTF32_LE, "utf-32-le"),
]

ENCODING_RE = re.compile(rb"coding[:=]\s*([-\w.]+)")

def auto_decode(data: bytes) -> str:
    """Check a bytes string for a BOM to correctly detect the encoding

    Fallback to locale.getpreferredencoding(False) like open() on Python3"""
    for bom, encoding in BOMS:
        if data.startswith(bom):
            return data[len(bom) :].decode(encoding)
    # Lets check the first two lines as in PEP263
    for line in data.split(b"\n")[:2]:
        if line[0:1] == b"#" and ENCODING_RE.search(line):
            result = ENCODING_RE.search(line)
            assert result is not None
            encoding = result.groups()[0].decode("ascii")
            return data.decode(encoding)
    return data.decode(
        locale.getpreferredencoding(False) or sys.getdefaultencoding(),
    )


from dsa.errors import UserError, MappingError
from io import BytesIO
import zlib

# Most of these errors should never be seen by a user attempting to supply
# valid PNGs in good faith.
class BAD_PNG_HEADER(UserError):
    """Not a valid PNG file (PNG signature not detected)"""

class PNG_EOF(UserError):
    """Unexpected end of PNG file"""

class PNG_CHECKSUM(UserError):
    """PNG file contained invalid checksum for `{chunk}` chunk"""

class BAD_PLTE(UserError):
    """PNG file has invalid palette (not a multiple of 3 bytes)"""

class BAD_IEND(UserError):
    """PNG file corrupt at end (bad IEND chunk)"""

class UNSUPPORTED_PNG(UserError):
    """PNG file not supported (must be 8bpp palettized, non-interlaced)"""

class BAD_FILTER(MappingError):
    """PNG file had invalid image data (filter method = `{key}`)"""

class DUPLICATE_CHUNK(UserError):
    """PNG file contains disallowed duplicate `{chunk}` chunk"""

class MISSING_CHUNK(MappingError):
    """PNG file missing `{key}` chunk"""

class BAD_PALETTE_DATA(UserError):
    """Palette data must be a multiple of 32 bytes long"""

class TOO_MANY_PALETTES(UserError):
    """Too many palettes (max 16)"""

class BAD_BITMAP_DATA(UserError):
    """Bitmap data must be a multiple of 64 bytes long"""

def _num(quad):
    return int.from_bytes(quad, 'big')

def _quad(value):
    return value.to_bytes(4, 'big')

def _chunk(name, data):
    chunk = name + data
    return _quad(len(data)) + chunk + _quad(zlib.crc32(chunk))

def _compress(rows):
    # No filters are applied by the compressor.
    # We need to prepend a 0 filter-type byte to the beginning of each row.
    return zlib.compress(b'\x00' + b'\x00'.join(rows))

_PNG_SIGNATURE = b'\x89PNG\r\n\x1a\n'

def _to_png(rows, plte):
    # rows = 8bpp raw bitmap data, as a list of bytes objects
    # palette = 48-byte chunk specifying 16 palette indices
    result = bytearray(_PNG_SIGNATURE)
    result.extend(_chunk(
        b'IHDR',
        len(rows[0]).to_bytes(4, 'big') + # width
        len(rows).to_bytes(4, 'big') + # height
        bytes([8, 3, 0, 0, 0]) # bit depth, palettized, not interlaced
    ))
    result.extend(_chunk(b'PLTE', plte))
    result.extend(_chunk(b'IDAT', _compress(rows)))
    result.extend(_chunk(b'IEND', b''))
    return bytes(result)

def _png_chunk_gen(data):
    read = BytesIO(data).read
    header = read(8)
    BAD_PNG_HEADER.require(header == _PNG_SIGNATURE)
    while True:
        size = read(4)
        if not size:
            break
        size = int.from_bytes(size, 'big')
        name = read(4)
        data = read(size)
        checksum = read(4)
        PNG_EOF.require(len(checksum) == 4)
        PNG_CHECKSUM.require(
            _quad(zlib.crc32(name + data)) == checksum, chunk=name
        )
        yield (name, data)

def _validate(plte, iend):
    BAD_PLTE.require(len(plte) % 3 == 0)
    BAD_IEND.require(iend == b'')

def _parse_header(ihdr):
    width, height = _num(ihdr[:4]), _num(ihdr[4:8])
    UNSUPPORTED_PNG.require(ihdr[8:] == b'\x08\x03\x00\x00\x00')
    return width, height

def _process_scanline(method, prev_line, this_line):
    left = 0
    result = bytearray()
    for up, this, diag in zip(prev_line, this_line, bytes(1) + prev_line[:-1]):
        average = (left + up) // 2
        reference = left + up - diag
        paeth = min((left, up, diag), key = lambda x: abs(reference - x))
        prediction = BAD_FILTER.get({
            0: 0, 1: left, 2: up, 3: average, 4: paeth
        }, method)
        # PNG spec says we do unsigned math throughout here.
        left = (this + prediction) & 0xff
        result.append(left)
    return bytes(result)

def _scanlines(width, height, idat):
    data = zlib.decompress(idat)
    bytes_per_line = width + 1
    prev_line = b'\x00' * width
    for line_start in range(0, height*bytes_per_line, bytes_per_line):
        method = data[line_start]
        this_line = data[line_start+1:line_start+bytes_per_line]
        this_line = _process_scanline(method, prev_line, this_line)
        yield this_line
        prev_line = this_line

def _png_data(raw):
    result = {}
    for name, data in _png_chunk_gen(raw):
        if name not in result:
            result[name] = data
            continue
        # concatenate IDAT chunks, refuse duplicates of anything else.
        DUPLICATE_CHUNK.require(name == b'IDAT', chunk=name)
        result[name] += data
    plte = MISSING_CHUNK.get(result, b'PLTE')
    iend = MISSING_CHUNK.get(result, b'IEND')
    _validate(plte, iend)
    width, height = _parse_header(MISSING_CHUNK.get(result, b'IHDR'))
    idat = MISSING_CHUNK.get(result, b'IDAT')
    return plte, idat, width, height

def rgb_to_xy(r, g, b):
    value = ((r >> 3) | ((g >> 3) << 5) | ((b >> 3) << 10))
    return value.to_bytes(2, 'little')

def xy_to_rgb(x, y):
    return bytes([
        (x & 0x1f) << 3,
        ((x >> 5) << 3) | ((y & 3) << 6),
        (y << 1) & 0xf8
    ])

def _compact_palette(plte):
    return b''.join(rgb_to_xy(r, g, b) for r, g, b in zip(*([iter(plte)] * 3)))

def _expand_palette(tile):
    return b''.join(xy_to_rgb(x, y) for x, y in zip(*([iter(tile)] * 2)))

def _dummy_rows(count):
    return [
        b''.join(bytes([x]) * 16 for x in range(base, base+16))
        for base in range(0, 16*count, 16)
        for block_row in range(16)
    ]

_DUMMY_PLTE = b''.join(bytes([x]) * 3 for x in range(0, 0x11 * 16, 0x11))

def _unpack_palette(data):
    count, remainder = divmod(len(data), 32)
    BAD_PALETTE_DATA.require(remainder == 0)
    # Interpret as palette data.
    TOO_MANY_PALETTES.require(count <= 16)
    return _to_png(_dummy_rows(count), _expand_palette(data))

def _unpack_bitmap(data, width):
    size = len(data)
    # Data should already be expanded to 8bpp, with 8x8 tiles arranged.
    BAD_BITMAP_DATA.require(size % 64 == 0)
    stride = 8 * width
    rows = [data[i:i+stride] for i in range(0, size, stride)]
    return _to_png(rows, _DUMMY_PLTE)

# Filter interface.
pack_args = ('palette flag', {None: False, 'true': True, 'false': False})
# width is omitted when unpacking a palette.
unpack_args = ('width of image in tiles', 'integer?')

def pack(codec_lookup, data, use_palette):
    plte, idat, width, height = _png_data(data)
    return _compact_palette(plte) if use_palette else b''.join(
        _scanlines(width, height, idat)
    )

class View:
    def __init__(self, codec_lookup, data, width):
        self._raw_size, self._is_palette = len(data), width is None
        self._data = (
            _unpack_palette(data)
            if self._is_palette
            else _unpack_bitmap(data, width)
        )

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

    def pack_params(self, unpacked):
        return self._raw_size, [[str(self._is_palette).lower()]]


"""Module containing the MultiSeq class and support functionality.

This class along with yaml extensions implemented in yaml_parse
(!multiseq and !range) allow you to define an experiment once while also
using a range of values for some of the parameters.

The simplest usage would imply defining both the control class (MultiSeq)
and the experiment in the same file like so:

.. code-block:: none

    {
        multiseq: !multiseq: {
            # ...
        },
        train: !obj:pylearn2.train.Train {
            # ...
        }
    }

"""
__authors__ = "Nicu Tofan"
__copyright__ = "Copyright 2015, Universite de Montreal"
__credits__ = ["Nicu Tofan"]
__license__ = "3-clause BSD"
__maintainer__ = "LISA Lab"
__email__ = "pylearn-dev@googlegroups"

import re
import os
from numpy import linspace
from random import randint
from datetime import datetime
from uuid import uuid1

from pylearn2.utils import as_floatX, string_utils
from pylearn2.utils.string_utils import preprocess

# Regular expression used to extract data from !range
RANGE_REGEXP = r'^\$\{([0-9a-zA-Z_]+):([0-9\.]+),([0-9\.]+),([0-9]+)\}$'

class RangeVar(object):
    """
    One range variable.

    Parameters
    ----------

    name : str
        The name as it will be used in the custome environment.
    space : numpy.linspace instance
        The linear space from which we will extract values..
    index : int
        Current index inside space.

    """
    def __init__(self, name, space):
        super(RangeVar, self).__init__()

        self.name = name
        self.space = space
        self.index = -1

    def compute_next_id(self, method):
        """
        Selects the next index based on current value and provided settings.

        If ordered mode is selected the values will be incremented and,
        at the end, will wrap around.MultiSeq

        If random mode is selected a random index is picked up.

        Returns
        -------
        new_idx : int
            Newly assigned index.
        """
        new_idx = self.index
        if method == 'ordered':
            new_idx = new_idx + 1
            if new_idx >= len(self.space):
                new_idx = 0
        elif method == 'random':
            new_idx = randint(0, len(self.space)-1)
        else:
            raise NotImplementedError(method)

        self.index = new_idx
        return new_idx

    def value(self):
        """
        Returns the value based on internal range and index.

        We can't just `float(self.space[self.index])` because it may get
        promoted to float64 and that is only desirable if floatX is float64.

        We can't also just `as_floatX(self.space[self.index])` because there
        are assertions in code that explicitly expect a float and
        as_floatX will create a one dimensional array.

        So, there you go. The solution below appears to work. I'm sure you
        can see why.
        """
        return float(as_floatX(self.space[self.index]))

class MultiSeq(object):
    """
    Implements multiple experiments from a single source using value ranges.

    This is a helper class for yaml_parser module.

    Parameters
    ----------

    multiseq_tag: str
        Unique tag for this run generated using uuid or current time.

    iterations: int
        Number of iterations performed so far.

    dynamic_env: dict of strings
        Variables that are updated at the start of a new iteration

    range_vars: list
        the list of variables with an associated range

    """
    _instance = None
    def __init__(self):

        super(MultiSeq, self).__init__()
        MultiSeq._instance = self

        self.multiseq_tag = None
        self.grouping = 'once'
        self.iterations = 0
        self.dynamic_env = {}
        self.range_vars = {}
        self.__range_regexp__ = re.compile(RANGE_REGEXP)

    @staticmethod
    def get_instance():
        """
        Get existing instance or create a new one.
        """
        if MultiSeq._instance is None:
            MultiSeq()
        return MultiSeq._instance

    @staticmethod
    def get_value(rng_var):
        """
        Get existing instance or create a new one.
        """
        return rng_var.value()

    def reinit(self, grouping, gen_val, gen_tag, max_count, environ=None):
        """
        Prepare to iterate over examples.
        """
        self.iterations = 0
        self.grouping = grouping
        self.gen_tag = gen_tag
        self.gen_val = gen_val
        self.max_count = max_count
        if environ is None:
            self.dynamic_env = {}
        else:
            self.dynamic_env = environ
        self.__order_of_change__ = []
        self.__var_to_change__ = -1

    def unique_tag(self):
        """
        Generate a unique tag for current run.
        """
        if self.gen_tag == 'uuid':
            self.multiseq_tag = str(uuid1())
        else:
            date_tm = datetime.utcnow()
            self.multiseq_tag = '%4d_%02d_%02d_%02d_%02d_%02d_%06d' % \
                                (date_tm.year, date_tm.month, date_tm.day,
                                 date_tm.hour, date_tm.minute, date_tm.second,
                                 date_tm.microsecond)

    def first_iteration(self):
        """
        Prepares the object for first run.
        """
        # prepare the counter
        self.iterations = 0

        # generate a unique tag
        self.unique_tag()

        # update all variables to next value
        for dynk in self.range_vars:
            self.range_vars[dynk].compute_next_id(self.gen_val)

        # we will change variables in this order
        self.__var_to_change__ = -1
        self.__order_of_change__ = self.range_vars.keys()

        # update variables in custom dictionary
        self.set_iter_vars()

    def next_iteration(self):
        """
        Advances to next step.

        Returns False if we're at the end.
        """
        # update the counter
        self.iterations = self.iterations + 1
        if self.iterations >= self.max_count:
            return False

        # generate a unique tag
        self.unique_tag()

        if self.grouping == 'once':
            return False
        elif self.grouping == 'one':
            # select the variable that is going to be changed
            if len(self.__order_of_change__) == 0:
                return False
            self.__var_to_change__ = self.__var_to_change__ + 1
            if self.__var_to_change__ >= len(self.__order_of_change__):
                self.__var_to_change__ = 0

            # only update this variable to next value
            var_name = self.__order_of_change__[self.__var_to_change__]
            self.range_vars[var_name].compute_next_id(self.gen_val)
        else:
            # update all variables to next value
            for dynk in self.range_vars:
                self.range_vars[dynk].compute_next_id(self.gen_val)

        # update variables in custom dictionary
        self.set_iter_vars()

        # continue the iteration
        return True

    def set_iter_vars(self):
        """
        Set the variables that change with each iteration.
        """
        self.dynamic_env['MULTISEQ_ITER'] = str(self.iterations)
        self.dynamic_env['MULTISEQ_TAG'] = self.multiseq_tag

        # add our range variables to custom environment
        for dynk in self.range_vars:
            self.dynamic_env[dynk] = str(self.val_for_name(dynk))

    @staticmethod
    def _to_number(str_to_conv):
        """
        Attempts to convert the value to either an integer or a float.
        """
        try:
            return int(str_to_conv)
        except ValueError:
            return float(str_to_conv)

    def add_range_var(self, var_str):
        """
        Creates or updated a range variable from a string.

        Expected string has the form "${VARNAME:start,stop,num}"
        """
        match = self.__range_regexp__.match(var_str)
        if not match:
            raise ValueError('Value %s is not a valid !range value' % str(var_str))

        start_val = MultiSeq._to_number(match.group(2))
        end_val = MultiSeq._to_number(match.group(3))
        num_val = MultiSeq._to_number(match.group(4))
        lspc = linspace(start=start_val,
                        stop=end_val,
                        num=num_val,
                        endpoint=True)

        # create the variable
        rng_var = RangeVar(name=match.group(1), space=lspc)

        # insert it in our list
        self.range_vars[match.group(1)] = rng_var

        # and return initial value
        return rng_var

    def val_for_name(self, rng_var_name):
        """
        Gets current value of a range variable given range variable's name.
        """
        return self.range_vars[rng_var_name].value()

    def save_params(self):
        """
        Converts the input to a string and dumps it to provided file.

        The file is only generated if local environment contains
        a key called `PARAMETERS_FILE`. Otherwise we assume that the
        user does not want a separate file.

        TODO:generate a runnable yaml file.
        """
        if not self.dynamic_env.has_key('PARAMETERS_FILE'):
            return
        file_path = preprocess(self.dynamic_env['PARAMETERS_FILE'],
                               self.dynamic_env)
        with open(file_path, "w") as f_param:
            for k in self.dynamic_env.keys():
                f_param.write('%s: %s\n' % \
                              (k,
                               preprocess(self.dynamic_env[k],
                                          environ=self.dynamic_env)))
            f_param.write('#----------------------------------------\n')
            for k in os.environ.keys():
                f_param.write('%s: %s\n' % (k, os.environ[k]))


from django.shortcuts import render, get_object_or_404
from django.http import Http404, HttpResponseRedirect, HttpResponse
from pokeapp.models import Pokemon
from .forms import PostForm
import json, os
from pprint import pprint
from django.conf import settings
from django.contrib import messages

# Create your views here.
def index(request):
    all_pokemons = Pokemon.objects.all()
    context = {
        'all_pokemons' : all_pokemons
    }
    return render(request, 'pokemon/index.html', context)

def create(request):
    if request.method == 'POST':
        # create a form instance and populate it with data from the request:
        form = PostForm(request.POST)
        # check whether it's valid:
        if form.is_valid():
            new_pokemon = form.save(commit=False)
            # process the data in form.cleaned_data as required
            json_data = open(os.path.join(settings.BASE_DIR, 'pokemons.json'), encoding='latin1' )
            data = json.load(json_data)
            checkpokemon = Pokemon.objects.filter(name = request.POST['name']).count()
            if checkpokemon:
                 messages.warning(request, 'Este pokemon ya existe en la base de datos.')
                 return HttpResponseRedirect('http://127.0.0.1:8000/pokemons/create')
            for pokemon in data:
                if pokemon['name'] == request.POST['name']:
                    new_pokemon.picture = pokemon['image_url']
                    new_pokemon.pokemon_type = pokemon['types'][0]
                    new_pokemon.birthday = request.POST['birthday']
                    new_pokemon.name = request.POST['name']
                    new_pokemon.zodiac_sign = request.POST['zodiac_sign']
                    new_pokemon.blood_type = request.POST['blood_type']
                    new_pokemon.latitude = request.POST['latitude']
                    new_pokemon.longitude = request.POST['longitude']
                    new_pokemon.save()
                    messages.success(request, 'Pokemon ' + new_pokemon.name + ' agregado correctamente')
                    break
            else:
                messages.warning(request, 'Este no es un pokemon valido.')
    # if a GET (or any other method) we'll create a blank form
    else:
        form = PostForm()

    return render(request, 'pokemon/create.html', {'form': form})

def reports(request):
    if request.method == 'POST':
        month = request.POST['month']
        all_pokemons = Pokemon.objects.filter(birthday__month=month)
        context = {
            'all_pokemons' : all_pokemons
        }
        return render(request, 'pokemon/bymonth.html', context )
    else:
        return render(request, 'pokemon/reports.html', context=None)

def bytype(request):
    if request.method == 'POST':
        poke_type = request.POST['pokemon_type']
        all_pokemons = Pokemon.objects.filter(pokemon_type=poke_type)
        context = {
            'all_pokemons' : all_pokemons
        }
        return render(request, 'pokemon/bymonth.html', context )

def detail(request, pokemon_id):
    pokemon = get_object_or_404(Pokemon, pk=pokemon_id)
    return render(request, 'pokemon/detail.html', {'pokemon' : pokemon})


import numpy as np
import scipy.constants as SI
from abel import InteractionPoint, Event

class InteractionPointBasic(InteractionPoint):

    # perform a simple interaction (geometric lumosity)
    def interact(self, beam1, beam2):

        # create event
        event = Event(beam1, beam2)

        # geometric factor
        H_D = 1

        # find overlapping area
        sigx = np.sqrt(beam1.beam_size_x()*beam2.beam_size_x())
        sigy = np.sqrt(beam1.beam_size_y()*beam2.beam_size_y())

        # get charge
        N1 = abs(beam1.charge()/SI.e)
        N2 = abs(beam2.charge()/SI.e)

        # calculate the geometric luminosity (per crossing)
        lumi = H_D / (4*np.pi) * N1 * N2 / (sigx * sigy)

        # save to event
        event.luminosity_full = lumi
        event.luminosity_peak = lumi
        event.luminosity_geom = lumi

        return event


import cv2
import numpy as np
img = cv2.imread("D://python//opencv2//cat.jpg")
cv2.imshow('picture',img)
imgInfo = img.shape
height = imgInfo[0]
width = imgInfo[1]
matSrc = np.float32([[0,0],[0,height-1],[width-1,0]])
matDst = np.float32([[50,50],[300,height-200],[width-300,100]])
matAffine = cv2.getAffineTransform(matSrc,matDst)
dst = cv2.warpAffine(img,matAffine,(width,height))
cv2.imshow('Affine',dst)
cv2.imwrite("D://python//opencv2//Affine.jpg",dst)
cv2.waitKey(0)


import networkx as nx
import os
import time
import seaborn as sns
import matplotlib.pyplot as plt
import random

'''
Script to analyse graphs
'''

def count(graph_dir, plot = False):
    """
    Print stats about a dir of graphs
    :param graph_dir: path containing a list of nx graphs
    :param plot: if True, display the distributions
    """
    number_of_nodes = []
    number_of_cc = []
    number_of_node_per_cc = []
    for pdb in os.listdir(graph_dir):
        path = os.path.join(graph_dir, pdb)
        g = nx.read_gpickle(path)
        number_of_cc += [nx.number_connected_components(g)]
        ccs = []
        for cc in nx.connected_components(g):
            ccs.append(len(cc))
        number_of_node_per_cc += [sum(ccs) / len(ccs)]
        number_of_nodes += [nx.number_of_nodes(g)]
    avg_nn = (sum(number_of_nodes) / len(number_of_nodes))
    print('The average number of nodes is {}'.format(avg_nn))

    avg_cc = (sum(number_of_cc) / len(number_of_cc))
    print('The average number of cc is {}'.format(avg_cc))

    avg_ncc = (sum(number_of_node_per_cc) / len(number_of_node_per_cc))
    print('The average number of node per cc is {}'.format(avg_ncc))

    # Plot it if needed
    if plot:
        sns.distplot(number_of_nodes, kde=False, rug=True)
        plt.show()
        sns.distplot(number_of_cc, kde=False, rug=True)
        plt.show()
        sns.distplot(number_of_node_per_cc, kde=False, rug=True)
        plt.show()

    return number_of_nodes, number_of_cc, number_of_node_per_cc

def plot_random(graph_dir, number):
    """
    Plot a random subset
    :param graph_dir: path of the dir
    :param number: number of plots needed
    :return:
    """
    files = os.listdir(graph_dir)
    for i in range(number):
        index = random.randrange(0, len(files))
        path = os.path.join(graph_dir, files[index])
        g = nx.read_gpickle(path)
        printer = dict()
        for edge in g.edges.data():
            printer[(edge[0], edge[1])] = edge[2]['interaction']
        print(printer)
        pos = nx.spring_layout(g)
        nx.draw_networkx_edge_labels(g, pos=pos, edge_labels=printer)
        nx.draw_networkx(g, pos=pos)
        plt.show()

if __name__ == "__main__":
    start_time = time.time()

    # count('../data/output_graph/1radius_12minnodes')
    # count('../data/output_graph/3radius_12minnodes')

    # plot_random('../data/output_graph/v5_hbonds',7)
    # wait = input("press enter to continue")
    # plot_random('../data/output_graph/1radius_12minnodes',7)

    # print('no_hbonds')
    # # count('../data/output_graph/test/1a0i_ATP_0.p')
    # print('hbonds')
    plot_random('../data/output_graph/test',5)

    # print('hbond2')
    # count('../data/output_graph/hbond2')
    #
    # print('3radius_12minnodes')
    # count('../data/output_graph/3radius_12minnodes')
    # # print('1radius_12minnodes')
    # # count('../data/output_graph/1radius_12minnodes')
    # print('border')
    # count('../data/output_graph/border')

    # plot_count('../data/output_graph/border')
    # plot_random('../data/output_graph/border', 10)

    print("--- %s seconds ---" % (time.time() - start_time))


### 데이터프레임을 사용할 수 있는 라이브러리 불러들이기
import pandas as pd

import json

### 지도 시각화 라이브러리 불러들이기
import folium

class Employment_map :

    ### 클래스 생성시 아래 함수 모두 실행시키기
    # - 생성자에서 호출
    def __init__(self) :
        self.initDataFrame()
        self.initJson()

    ### 데이터 읽어들이기
    def initDataFrame(self) :
        file_path = "./yeargraphapp/employment/전력량_고용률_마커.csv"
        self.usage_data = pd.read_csv(file_path)
        self.usage_data.drop(columns=['Unnamed: 0'], inplace=True)

    # 특정 년, 월 받아와서 데이터 프레임 만들기
    def setDataFrame(self, year_data, month_data) :
        year = year_data
        if year is "ERROR":
            year = 2021
        year = int(year)
        month = month_data
        if month is "ERROR":
            month = 12
        month = int(month)

        self.df_ver = self.usage_data.query('연도 == @year & 월 == @month')
        self.df_data = self.df_ver[['시도', '고용률', '사용량', '위도', '경도']]
        self.df_data_test = self.df_data.to_dict(orient='records')

    def initJson(self) :
        jsonfile = open('./yeargraphapp/employment/Sido.json','r',encoding='utf8').read()
        self.jsondata = json.loads(jsonfile)
        self.jsondata_loc = {'type': 'FeatureCollection'}
        self.jsondata_pick = []

    def getMap(self) :
        for idx in self.jsondata['features']:
            idx['id'] = idx['properties']['CTP_KOR_NM']
            for i in self.df_data_test:
                if i['시도'] == idx['properties']['CTP_KOR_NM']:
                    idx['properties']['pop'] = i['고용률']
            self.jsondata_pick.append(idx)
        self.jsondata_loc['features'] = self.jsondata_pick

    ### 지도맵 그리기
    def map_base(self) :
        # 지도 생성, location : 시작위치, zoom_start : 시작 시 확대정도
        self.employment_map = folium.Map(location=[36.0068191, 127.6607805],zoom_start=7)
        # 색상단계도 생성, geo_data : 참고할 위치정보 데이터
        cho = folium.Choropleth(geo_data=self.jsondata_loc,
                        # data : geo_data와 비교해 맞출 데이터
                        data=self.df_data,
                        # columns : data에서 사용할 열
                        columns=['시도', '고용률'],
                        fill_color="YlGnBu",                      # fill_color : 색
                        key_on='feature.id',                      # key_on : data와 맞출 geo_data의 데이터
                        legend_name='맵').add_to(self.employment_map)    # legend_name : 범례 제목 입력
        cho.geojson.add_child(                                      # 속성 부여
            folium.features.GeoJsonPopup(['CTP_KOR_NM','pop'],labels=False,style='font-weight : bold')
        )
        for row in self.df_data_test:
            lat = row['위도']
            lon = row['경도']
            folium.CircleMarker(location=[lat, lon],
                                radius=row['사용량'] / 700000, # 전력량 값을 radius로 사용
                                color='red',
                                fill_color='red',
                                fill_opacity=0.3,
                                popup=(row['시도'], row['사용량']),
                                tooltip=row['시도']).add_to(cho)
        return self.employment_map._repr_html_()

    ### 데이터프레임 리턴하기
    def getDataFrame(self) :
        return self.employment_data


'''
Models tested and used for RODENet.
Brent Griffin, griffb@umich.edu, 190424
'''

import torch
import torch.nn as nn
import torch.nn.functional as F
import IPython

def count_parameters(model):
	return sum(p.numel() for p in model.parameters() if p.requires_grad)

# ResOpt
def pass_arg(**kwargs):
	model = ResOptNet(CustomBlock, **kwargs)
	return model

class ResOptNet(nn.Module):
	def __init__(self, block, layr=[1], kern=[1], fcs=[512], n1=5, n2=5,
			n2_in=True, n2_fc=True, avg_pool=7, par_net=False, par_n=1,
			batch_nrm=True, n2_mid=True, mid_rep=1, strd=[1], conv_kern=14,
			max_pool=True, fc_drop=0):
		super(ResOptNet, self).__init__()
		# Initial config.
		self.y_in=n2_in; self.y_fc=n2_fc; self.par_net=par_net
		self.n1=n1; self.par_n=par_n; self.inplanes=kern[0]
		self.bn = batch_nrm; self.n2_mid = n2_mid; self.mid_rep = mid_rep
		self.n_fc = len(fcs); self.n_blck = len(kern); self.max_pool = max_pool
		# Parallel process input masks (notes pg. 27 190923).
		if self.par_net:
			n_conv0 = n1 - par_n + 1 + n2 * n2_in / par_n
			template_idx = range(self.par_n, self.n1)
			idxs = []
			for i in range(par_n):
				if n2_in: idx = [i, i + n1] + template_idx
				else: idx = [i] + template_idx
				idxs.append(idx)
			self.net1_idxs = idxs
		else: n_conv0 = n1 + n2 * n2_in

		# Initial convolution and max pooling.
		self.conv1 = nn.Conv2d(n_conv0, self.inplanes, kernel_size = conv_kern,
				stride=2, padding=3, bias=False)
		if self.bn: self.bn1 = nn.BatchNorm2d(self.inplanes)
		self.relu = nn.ReLU(inplace=True)
		self.avgpool = nn.AvgPool2d(avg_pool, stride=1)
		if self.max_pool:
			self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

		# Residual layers.
		self.layer1 = self._make_layer(block,kern[1],layr[0],stride=strd[0])
		if self.par_net: self.inplanes *= par_n
		if self.n2_mid: self.inplanes += n2 * mid_rep
		self.layer2 = self._make_layer(block,kern[2],layr[1],stride=strd[1])
		self.layer3 = self._make_layer(block,kern[3],layr[2],stride=strd[2])
		if self.n_blck > 4:
			self.layer4 = self._make_layer(block,kern[4],layr[3],stride=strd[3])
		# Fully connected layers.
		self.fc0 = nn.Linear(fcs[0] + n2 * n2_fc, fcs[1]) # If adding distance.
		#self.fc0 = nn.Linear(fcs[0], fcs[1])
		if self.n_fc > 2: self.fc1 = nn.Linear(fcs[1] + n2 * n2_fc, fcs[2])
		if self.n_fc > 3: self.fc2 = nn.Linear(fcs[2] + n2 * n2_fc, fcs[3])
		self.fc = nn.Linear(fcs[-1] * block.expansion + n2 * n2_fc, 1)
		if fc_drop > 0:
			self.drop = True; self.dropout = nn.Dropout(p=fc_drop)
		else: self.drop = False

	def _make_layer(self, block, planes, blocks, stride=1):
		downsample = None
		if stride !=1 or self.inplanes != planes * block.expansion:
			if self.bn:
				downsample = nn.Sequential(
						nn.Conv2d(self.inplanes, planes * block.expansion,
								kernel_size=1, stride=stride, bias=False),
						nn.BatchNorm2d(planes * block.expansion),
				)
			else:
				downsample = nn.Conv2d(self.inplanes, planes * block.expansion,
								kernel_size=1, stride=stride, bias=False)
		layers = []
		layers.append(block(self.inplanes,planes,stride,downsample,bn=self.bn))
		self.inplanes = planes * block.expansion
		for i in range(1, blocks):
			layers.append(block(self.inplanes, planes,bn=self.bn))
		return nn.Sequential(*layers)

	def forward(self, x, y):
		# Concatenate with Residual Optics (see notes pg. 16 on 190710)
		if self.y_in:
			# y also includes camera move distance.
			x = torch.cat((x, y), 1) # Add camera move distance.
		if self.par_net:
			# Process each mask independently and then cat.
			x0 = self.net_1(x[:, self.net1_idxs[0]])
			for i in range(1, self.par_n):
				x0 = torch.cat((x0, self.net_1(x[:,self.net1_idxs[i]])), 1)
			x = x0
		else:
			x = self.conv1(x) # 480x640 (224x224) [56x56, 480x480]
			if self.bn: x = self.bn1(x)
			x = self.relu(x)
			if self.max_pool:
				x = self.maxpool(x) # 240x320 (112x112) [28x28, 240x240]
			x = self.layer1(x) # 120x160 (56x56) [14x14, 120x120]

		# Potential distance / residual concatenate point.
		if self.n2_mid:
			y0 = y
			for i in range(1, self.mid_rep):
				y0 = torch.cat((y0, y), 1)
			x = torch.cat((x, y0), 1)

		x = self.layer2(x) # 60x80 (28x28) [7x7, 60x60]
		#print('pre layer 3'); IPython.embed()
		x = self.layer3(x) # 30x40 (14x14) [4x4, 30x30]
		if not hasattr(self, 'n_blck'):
			self.n_blck = 5 # remove for later models
		#print('pre layer 4'); IPython.embed()
		if self.n_blck > 4:
			x = self.layer4(x) # 15x20 (7x7) [2x2, 15x15]
		#x = self.layer5(x) # 8x10 (4x4) [1x1, 8x8]
		#print('pre average pool'); IPython.embed()

		if not hasattr(self, 'drop'):
			self.drop = False # remove for later models

		x = self.avgpool(x) # 1x1
		#print('post average pool'); IPython.embed()
		x = x.view(x.size(0), -1)
		if self.y_fc:
			# Add camera move distance to fc layers.
			distances = y[:,:,0,0]
			x = torch.cat((x, distances), 1)
		x = self.fc0(x)
		x = self.relu(x)
		if self.drop: x = self.dropout(x)
		if self.n_fc >2:
			if self.y_fc: x = torch.cat((x, distances), 1)
			x = self.fc1(x)
			x = self.relu(x)
			if self.drop: x = self.dropout(x)
		if self.n_fc > 3:
			if self.y_fc: x = torch.cat((x, distances), 1)
			x = self.fc2(x)
			x = self.relu(x) # Uncommented and moved into if statement 191101.
			if self.drop: x = self.dropout(x)
		if self.y_fc: x = torch.cat((x, distances), 1)
		x = self.fc(x)

		return x

	def net_1(self, x):
		x = self.conv1(x) # 480x640 (224x224) [56x56, 480x480]
		if self.bn: x = self.bn1(x)
		x = self.relu(x)
		x = self.maxpool(x)
		x = self.layer1(x) # 240x320 (112x112) [28x28, 240x240]
		return x

class CustomBlock(nn.Module):
	expansion = 1
	def __init__(self, inplanes, planes, stride=1, downsample=None, bn=True):
		super(CustomBlock, self).__init__()
		self.conv1 = conv3x3(inplanes, planes, stride)
		if bn: self.bn1 = nn.BatchNorm2d(planes)
		self.relu = nn.ReLU(inplace=True)
		self.conv2 = conv3x3(planes, planes)
		if bn: self.bn2 = nn.BatchNorm2d(planes)
		self.downsample = downsample
		self.bn = bn
		#self.stride = stride

	def forward(self, x):
		residual = x

		out = self.conv1(x)
		if self.bn: out = self.bn1(out)
		out = self.relu(out)

		out = self.conv2(out)
		if self.bn: out = self.bn2(out)

		if self.downsample is not None:
			residual = self.downsample(x)

		out += residual
		out = self.relu(out)

		return out

class BasicBlock(nn.Module):
	expansion = 1
	def __init__(self, inplanes, planes, stride=1, downsample=None):
		super(BasicBlock, self).__init__()
		self.conv1 = conv3x3(inplanes, planes, stride)
		self.bn1 = nn.BatchNorm2d(planes)
		self.relu = nn.ReLU(inplace=True)
		self.conv2 = conv3x3(planes, planes)
		self.bn2 = nn.BatchNorm2d(planes)
		self.downsample = downsample
		#self.stride = stride

	def forward(self, x):
		residual = x

		out = self.conv1(x)
		out = self.bn1(out)
		out = self.relu(out)

		out = self.conv2(out)
		out = self.bn2(out)

		if self.downsample is not None:
			residual = self.downsample(x)

		out += residual
		out = self.relu(out)

		return out

# RoadResNet

def rode_res_224_nb(**kwargs):
	model = RODEResNet(NoBatchBlock, [2, 2, 2, 2, 2, 2],
			[1, 1, 2, 2, 4, 4], [16, 8], **kwargs)
	return model

def rode_res_224(**kwargs):
	model = RODEResNet(BasicBlock, [2, 2, 2, 2, 2, 2],
			[1, 1, 2, 2, 4, 4], [16, 8, 4, 2], **kwargs)
	return model

def rode_res_480_ro(**kwargs):
	# Include residual optics input.
	model = RODEResNet(BasicBlock, [2, 2, 2, 2, 2, 2, 2],
			[64, 64, 128, 128, 256, 256, 256], [1536, 768, 384, 192], **kwargs)
	return model

def rode_res_480(**kwargs):
	model = RODEResNet(BasicBlock, [2, 2, 2, 2, 2, 2, 2],
			[64, 64, 128, 128, 256, 256, 256], [1536, 768, 384, 192], **kwargs)
	#model = RODEResNet(BasicBlock, [2, 2, 2, 2, 2, 2, 2],
			#[1, 1, 2, 2, 4, 4, 4], [24, 12], **kwargs) # 4 x 6
	#model = RODEResNet(BasicBlock, [2, 2, 2, 2, 2, 2, 2],
			#[8, 8, 16, 16, 32, 32, 32], [192, 96], **kwargs)
	return model

def roderesnet34(**kwargs):
	#model = RODEResNet(BasicBlock, [3, 4, 6, 3], [16, 32, 64, 128], **kwargs)
	model = RODEResNet(BasicBlock, [3, 4, 6, 3], [32, 64, 128, 256], **kwargs)
	return model

def roderesnet18(**kwargs):
	#model = RODEResNet(BasicBlock, [2, 2, 2, 2], [32, 64, 128, 256], **kwargs)
	#model = RODEResNet(BasicBlock, [2, 2, 2, 2], [16, 32, 64, 128], **kwargs)
	model = RODEResNet(BasicBlock, [2, 2, 2, 2, 2], [32, 64, 128, 256, 512],
			**kwargs)
	return model

class RODEResNet(nn.Module):
	def __init__(self, block, layers, kernels, fcs=[512], norm=True, n=5):
		self.inplanes = kernels[0]
		super(RODEResNet, self).__init__()
		self.conv1 = nn.Conv2d(n, self.inplanes, kernel_size=7, stride=2,
				padding=3, bias=False)
		if norm:
			self.bn1 = nn.BatchNorm2d(self.inplanes)
		self.relu = nn.ReLU(inplace=True)
		self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

		self.layer1 = self._make_layer(block, kernels[0], layers[0])
		self.layer2 = self._make_layer(block, kernels[1], layers[1], stride=2)
		self.layer3 = self._make_layer(block, kernels[2], layers[2], stride=2)
		self.layer4 = self._make_layer(block, kernels[3], layers[3], stride=2)
		#self.avgpool = nn.AvgPool2d((15,20), stride=1)
		self.layer5 = self._make_layer(block, kernels[4], layers[4], stride=2)
		#self.avgpool = nn.AvgPool2d((8,10), stride=1)
		self.layer6 = self._make_layer(block, kernels[5], layers[5], stride=2)
		self.layer7 = self._make_layer(block, kernels[6], layers[6], stride=2)
		self.fc0 = nn.Linear(fcs[0] + (n-1), fcs[1])
		self.fc1 = nn.Linear(fcs[1], fcs[2])
		self.fc2 = nn.Linear(fcs[2], fcs[3])
		self.fc = nn.Linear(fcs[-1] * block.expansion, 1)

	def _make_layer(self, block, planes, blocks, stride=1, norm=True):
		downsample = None
		if stride !=1 or self.inplanes != planes * block.expansion:
			if norm:
				downsample = nn.Sequential(
						nn.Conv2d(self.inplanes, planes * block.expansion,
								kernel_size=1, stride=stride, bias=False),
						nn.BatchNorm2d(planes * block.expansion),
				)
			else:
				downsample = nn.Conv2d(self.inplanes, planes * block.expansion,
						kernel_size=1, stride=stride, bias=False)

		layers = []
		layers.append(block(self.inplanes, planes, stride, downsample))
		self.inplanes = planes * block.expansion
		for i in range(1, blocks):
			layers.append(block(self.inplanes, planes))

		return nn.Sequential(*layers)

	def forward(self, x, y):
		x = self.conv1(x) # 480x640 (224x224)
		x = self.bn1(x)
		x = self.relu(x)
		x = self.maxpool(x) # 240x320 (112x112)

		x = self.layer1(x) # 120x160 (56x56)
		# Concatenate here for Residual Optics (see notes pg. 16 on 190710)
		x = self.layer2(x) # 60x80 (28x28)
		x = self.layer3(x) # 30x40 (14x14)
		x = self.layer4(x) # 15x20 (7x7)
		x = self.layer5(x) # 8x10 (4x4)
		x = self.layer6(x) # 4x5 (2x2)
		x = self.layer7(x) # 2x3

		#x = self.avgpool(x) # 1x1
		x = x.view(x.size(0), -1)

		x = torch.cat((x, y), 1) # Add camera move distance.
		x = self.fc0(x)
		x = self.relu(x)
		x = self.fc1(x)
		x = self.relu(x)
		x = self.fc2(x)
		x = self.relu(x)
		x = self.fc(x)

		return x

class NoBatchBlock(nn.Module):
	expansion = 1
	def __init__(self, inplanes, planes, stride=1, downsample=None):
		super(NoBatchBlock, self).__init__()
		self.conv1 = conv3x3(inplanes, planes, stride)
		self.relu = nn.ReLU(inplace=True)
		self.conv2 = conv3x3(planes, planes)
		self.downsample = downsample

	def forward(self, x):
		residual = x
		out = self.conv1(x)
		out = self.relu(out)
		out = self.conv2(out)
		if self.downsample is not None:
			residual = self.downsample(x)
		out += residual
		out = self.relu(out)
		return out

# ResNet

def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
	# 3x3 convolution with padding
	return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
			padding=dilation, groups=groups, bias=False, dilation=dilation)

def resnet18(**kwargs):
	model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
	return model

def resnet34(**kwargs):
	model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
	return model

class ResNet(nn.Module):
	def __init__(self, block, layers, num_classes=1):
		self.inplanes = 64
		super(ResNet, self).__init__()
		self.conv1 = nn.Conv2d(2, self.inplanes, kernel_size=7, stride=2,
				padding=3, bias=False)
		#self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
		#		bias=False)
		self.bn1 = nn.BatchNorm2d(self.inplanes)
		self.relu = nn.ReLU(inplace=True)
		self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

		self.layer1 = self._make_layer(block, 64, layers[0])
		self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
		self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
		self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

		self.avgpool = nn.AvgPool2d((15,20), stride=1)
		#self.avgpool = nn.AvgPool2d(7, stride=1)
		self.fc = nn.Linear(512 * block.expansion, num_classes)

	def _make_layer(self, block, planes, blocks, stride=1):
		downsample = None
		if stride !=1 or self.inplanes != planes * block.expansion:
			downsample = nn.Sequential(
					nn.Conv2d(self.inplanes, planes * block.expansion,
							kernel_size=1, stride=stride, bias=False),
					nn.BatchNorm2d(planes * block.expansion),
			)

		layers = []
		layers.append(block(self.inplanes, planes, stride, downsample))
		self.inplanes = planes * block.expansion
		for i in range(1, blocks):
			layers.append(block(self.inplanes, planes))

		return nn.Sequential(*layers)

	def forward(self, x):
		x = self.conv1(x) # 224x224 (RODE: 480x640x2)
		x = self.bn1(x)
		x = self.relu(x)
		x = self.maxpool(x) # 112x112 (RODE: 240x320)

		x = self.layer1(x) # 56x56 (RODE: 120x160)
		x = self.layer2(x) # 28x28 (RODE: 60x80)
		x = self.layer3(x) # 14x14 (RODE: 30x40)
		x = self.layer4(x) # 7x7   (RODE: 15x20)

		x = self.avgpool(x) # 1x1
		x = x.view(x.size(0), -1)
		x = self.fc(x)

		return x

class LeNet(nn.Module):
	def __init__(self):
		super(LeNet, self).__init__()
		self.conv1 = nn.Conv2d(2, 6, 5) # 2 channel for two binary masks.
		#self.conv1 = nn.Conv2d(1, 6, 5)
		self.pool = nn.MaxPool2d(2, 2) # No pooling in future!
		self.conv2 = nn.Conv2d(6, 16, 5)
		self.fc1 = nn.Linear(16 * 117 * 157, 120)
		#self.fc1 = nn.Linear(16 * 5 * 5, 120)
		self.fc2 = nn.Linear(120, 84)
		self.fc3 = nn.Linear(84, 1) # Need output to be 1 (RO depth)
		#self.fc3 = nn.Linear(84, 10)

	def forward(self, x):
		x = self.pool(F.relu(self.conv1(x)))
		x = self.pool(F.relu(self.conv2(x)))
		x = x.view(-1, 16 * 117 * 157) # Account for larger input size.
		#x = x.view(-1, 16 * 5 * 5)
		x = F.relu(self.fc1(x))
		x = F.relu(self.fc2(x))
		x = self.fc3(x)
		return x


import LeducPoker.LeducPokerGame
from LeducPoker.LeducPokerGame import PlayerActions
import numpy as np
from abc import ABC, abstractmethod
from typing import List

class Policy(ABC):
    @abstractmethod
    def action_prob(self, infoset: LeducPoker.LeducPokerGame.LeducInfoset) -> List[float]:
        pass

    def get_action(self, infoset: LeducPoker.LeducPokerGame.LeducInfoset) -> int:
        act_prob = self.action_prob(infoset)
        retval = np.random.choice(
            [PlayerActions.FOLD, PlayerActions.CHECK_CALL, PlayerActions.BET_RAISE], 1, p=act_prob)
        return retval

class AlwaysCheckCall(Policy):
    def action_prob(self, infoset: LeducPoker.LeducPokerGame.LeducInfoset) -> List[float]:
        retval = [0, 0, 0]
        retval[PlayerActions.CHECK_CALL] = 1
        return retval

    def get_action(self, infoset: LeducPoker.LeducPokerGame.LeducInfoset) -> int:
        return PlayerActions.CHECK_CALL

class AlwaysBet(Policy):
    def action_prob(self, infoset: LeducPoker.LeducPokerGame.LeducInfoset) -> List[float]:
        retval = [0, 0, 0]
        if infoset.can_raise:
            retval[PlayerActions.BET_RAISE] = 1
        else:
            retval[PlayerActions.CHECK_CALL] = 1
        return retval

    def get_action(self, infoset: LeducPoker.LeducPokerGame.LeducInfoset) -> int:
        return PlayerActions.BET_RAISE

class NashPolicy(Policy):
    def __init__(self, p0_strat_filename, p1_strat_filename):
        self.player_policies = [
            self.load_policy(p0_strat_filename),
            self.load_policy(p1_strat_filename)
        ]

    def load_policy(self, policy_filename):
        char_to_card = {'J': 0, 'Q': 1, 'K': 2}
        char_to_action = {
            'r': PlayerActions.BET_RAISE,
            'c': PlayerActions.CHECK_CALL
        }

        retval = {}
        with open(policy_filename, "r") as policy_filename:
            for line in policy_filename:
                if line[0] == "#":
                    continue

                cards, actions, probs = line.split(":")

                player_card = char_to_card[cards[0]]
                if len(cards) > 1:
                    board_card = char_to_card[cards[1]]
                else:
                    board_card = None

                round_actions = actions.split("/")[1:]
                round1_actions = tuple(char_to_action[c] for c in round_actions[0])
                if len(round_actions) > 1:
                    round2_actions = tuple(char_to_action[c] for c in round_actions[1])
                else:
                    round2_actions = ()

                infoset_template = (
                    player_card, board_card, round1_actions, round2_actions
                )

                raise_prob, call_prob, fold_prob = list([float(p) for p in probs.strip(" \n").split(" ")])
                infoset_probs = [0, 0, 0]
                infoset_probs[PlayerActions.BET_RAISE] = raise_prob
                infoset_probs[PlayerActions.CHECK_CALL] = call_prob
                infoset_probs[PlayerActions.FOLD] = fold_prob

                assert infoset_template not in retval
                retval[infoset_template] = infoset_probs

        return retval

    def action_prob(self, infoset: LeducPoker.LeducPokerGame.LeducInfoset) -> List[float]:
        assert not infoset.is_terminal

        if infoset.board_card is not None:
            board_card = infoset.board_card % 3
        else:
            board_card = None

        infoset_template = (
            infoset.card % 3, board_card, infoset.bet_sequences[0], infoset.bet_sequences[1]
        )

        return self.player_policies[infoset.player_to_act][infoset_template]


import datetime
import math

def get_julian_datetime(date):
    """
    Convert a datetime object into julian float.
    Args:
        date: datetime-object of date in question

    Returns: float - Julian calculated datetime.
    Raises:
        TypeError : Incorrect parameter type
        ValueError: Date out of range of equation
    """

    # Ensure correct format
    if not isinstance(date, datetime.datetime):
        raise TypeError('Invalid type for parameter "date" - need datetime')
    if date.year < 1801 or date.year > 2099:
        raise ValueError('Datetime must be between year 1801 and 2099')

    # Perform the calculation
    julian_datetime = (367 * date.year
                       - int((7*(date.year+int((date.month+9)/12.0)))/4.0)
                       + int((275*date.month)/9.0)
                       + date.day
                       + 1721013.5
                       + (date.hour
                          + date.minute/60.0
                          + date.second/math.pow(60, 2))/24.0
                       - 0.5*math.copysign(1, (100*date.year
                                               + date.month-190002.5))
                       + 0.5)

    return julian_datetime

# GATHER OBSERVATIONS
lcdb_fdir = "/home/xubu/NEFE/LCDB/LC_DAT_PUB.TXT"
observations = []
dates = []
with open(lcdb_fdir, "r") as lcdb_file:
    lines = lcdb_file.readlines()[5:]
    counter = 0
    numslst = []
    object_counter = 0
    quality_dict = {"": -1.0, "0": 0.0,
                    "1-": 1.0, "1": 2.0, "1+": 3.0,
                    "2-": 4.0, "2": 5.0, "2+": 6.0,
                    "3-": 7.0, "3": 8.0}
    for lnidx, line in enumerate(lines):
        if len(line.strip()) == 0:
            continue
        varstrs = []
        split_pos = [7, 9, 40, 61, 70, 72, 78, 80, 82, 91, 93, 100, 102, 105,
                     107, 114, 116, 130, 146, 148, 153, 158, 161, 167, 171,
                     175, 179, 185, 189]
        for idx, pos in enumerate(split_pos):
            if idx == 0:
                previous_pos = 0
            else:
                previous_pos = split_pos[idx-1]
            varstrs.append(line[previous_pos:pos])
        for idx, item in enumerate(varstrs):
            varstrs[idx] = item.strip()

        MPCnum, entryflag, namestr, col_3, family, csource, clss = varstrs[0:7]
        dsource, dflag, diam, hsource, h, hband, asource = varstrs[7:14]
        albedoflag, alb, pflag, period, pdescrip, ampflag = varstrs[14:20]
        ampmin, ampmax, u, notes, binary, private, pole = varstrs[20:27]
        survey, extnotes = varstrs[27:]
        if hband == "":
            hband = "V"
        try:
            MPCnum = int(MPCnum)
        except ValueError:
            if MPCnum != "":
                print(lnidx)
                print("MPCnum: ", MPCnum)
            MPCnum = None
        qual_obs = quality_dict[u]
        float_list = [diam, h, alb, period, ampmin, ampmax]
        for idx, item in enumerate(float_list):
            try:
                float_list[idx] = float(item)
            except ValueError:
                if item != "":
                    print(lnidx)
                    print(idx, "float item: ", item)
                float_list[idx] = None
        diam, h, alb, period, ampmin, ampmax = float_list
        if MPCnum is not None:
            object_counter = object_counter + 1
            mpc_number = MPCnum
            name = namestr
            desig = col_3
            diameter = diam
            obj_class = clss
            fam = family
            abs_mag = h
            albedo = alb
            qual_gen = qual_obs
            minamp = ampmin
            maxamp = ampmax
            bin_flag = binary
            pole_availability = pole
            srvy_nm = survey
        try:
            date = datetime.datetime.strptime(col_3, '%Y-%m-%d')
        except ValueError:
            date = None
        if (isinstance(date, datetime.datetime) and isinstance(period, float)
            and isinstance(diameter, float)):
            date = str(get_julian_datetime(date))
            observations.append([object_counter, mpc_number, name, desig,
                                 date, period, diameter, abs_mag,
                                 obj_class, fam, albedo, qual_obs, qual_gen,
                                 minamp, maxamp,
                                 bin_flag, pole_availability, srvy_nm])
            dates.append(date)
            counter = counter + 1


from collections import defaultdict
from typing import Any

def reformat_emissions(emissions: list[dict], subsector: str) -> list[dict]:
    """Select `subsector` emissions and simplify format.

    Expected format of `emissions` arg is:
    `emissions = [country_a, country_b ...]`, where
    `country_a =  {"name": country_code, "data": subsector_data}`, where
    `subsector_data =  {"name": subsector, "emissions": subs_emissions}`, where
    `subs_emissions =  ["year": y, ghg: value_in_tonnes ...]`

    Format returned by this function is:
    `emissions = [yearly_emissions_a ...]`, where
    `yearly_emissions_a = {"year": y, country_code: co2e_in_tonnes ...}`
    for the requested subsector.

    Warning! Only the "co2100" attribute (CO2 equivalent after 100 years)
    is selected.

    Args:
        emissions: Data received from Climate TRACE. See format above.
        subsector: Name of subsector to select the data for.

    Returns:
    """
    # Keys are years, value is dict like `yearly_emissions_a` from docstring
    country_emissions = defaultdict(dict)

    for country in emissions:
        country_code = country["name"]

        for subsector_data in country["data"]:
            if subsector_data["name"] != subsector:
                continue

            # Add each year of emissions for this country
            for yearly_e in subsector_data["emissions"]:
                year = yearly_e["year"]
                # Convert to million tonnes
                million_co2 = yearly_e["co2100"] / 1000000
                country_emissions[year][country_code] = million_co2

    # Add year attribute to each dict
    country_emissions = {
        y: dict(data, **{"year": y})
        for y, data in country_emissions.items()
    }
    # Drop the keys and return as list of yearly emissions
    parsed = list(country_emissions.values())
    return parsed


from contextlib import contextmanager
import http.client as http_client
import logging
from datetime import datetime
from time import sleep
import json

import tempfile
import os

from io import BytesIO
from matplotlib import pyplot as plt
from PIL import Image
from h5py import File as H5File
import xarray as xa
import numpy as np
import geopandas as gpd
import contextily as ctx

from satstac import Catalog

import requests
from cachecontrol import CacheController, CacheControlAdapter

def _build_session():
  """Builds a requests session that caches responses where possible, making redirects faster.

  Returns:
      requests.Session -- A shared session to use for the notebook
  """
  result = requests.session()

  # Set up caching.  Particularly obey and cache 307 redirects to avoid duplicate expensive calls when we already
  # have a result
  cache_adapter = CacheControlAdapter()
  cache_adapter.controller = CacheController(cache=cache_adapter.cache, status_codes=(200, 203, 300, 301, 307))

  result.mount('http://', cache_adapter)
  result.mount('https://', cache_adapter)
  return result

# Session accessible by callers
session = _build_session()

def debug_http():
  """Adds debugging output to HTTP requests to show redirects, headers, etc
  """
  http_client.HTTPConnection.debuglevel = 1
  logging.basicConfig()
  logging.getLogger().setLevel(logging.DEBUG)
  requests_log = logging.getLogger("requests.packages.urllib3")
  requests_log.setLevel(logging.DEBUG)
  requests_log.propagate = True

def request(*args, **kwargs):
  """Thin wrapper around requests.Request, logging URL sent and Content-Type received

  See https://requests.readthedocs.io/en/master/api/#requests.Request for args

  Returns:
      requests.Response -- The response to the request
  """
  req = requests.Request(*args, **kwargs)
  prepped = session.prepare_request(req)

  print('%s %s' % (prepped.method, prepped.path_url))
  response = session.send(prepped)
  #print('Received %s' % (response.headers.get('Content-Type', 'unknown content',)))
  return response

def get(*args, **kwargs):
  """Performs a GET request using the request wrapper

  See https://requests.readthedocs.io/en/master/api/#requests.Request for args

  Returns:
      requests.Response -- The response to the request
  """
  return request('GET', *args, **kwargs)

def post(*args, **kwargs):
  """Performs a POST request using the request wrapper

  See https://requests.readthedocs.io/en/master/api/#requests.Request for args

  Returns:
      requests.Response -- The response to the request
  """
  return request('POST', *args, **kwargs)

def show_shape(filename, basemap=True):
  """Plots the shapefile in the given filename with optional basemap (ESRI or GeoJSON)

  Arguments:
      filename {string} -- The filename of the shapefile to display

  Keyword Arguments:
      basemap {bool} -- Whether to display a basemap under the shapefile (default: {True})
  """
  shape = gpd.read_file(filename).to_crs(epsg=3857)
  plot = shape.plot(alpha=0.5, edgecolor='k', figsize=(8, 8))
  if basemap:
    ctx.add_basemap(plot)

def show(response, varList=[], color_index=None, immediate=True, flip=True):
  """Shows a variety of responses possible from Harmony for its example data

  Handles NetCDF files with red_var, green_var, blue_var, and alpha_var bands, compositing output
  into a single colored image, ESRI Shapefiles with basemaps, and any type of image that can be
  read by PIL, including GeoTIFF

  Arguments:
      response {requests.Response} -- The response containing the data to display
      varList {array} -- If set, only plot the variables listed in varList.  Otherwise, plot all.

  Keyword Arguments:
      color_index {number} -- Set for monochromatic images to put the output in a color band (0=red, 1=green, 2=blue) (default: {None})
      immediate {bool} -- True if the data should be shown immediately in the notebook (default: {True})
      flip {bool} -- Whether to flip the numpy arrays for stacked plots

  """

  # show_netcdf (look at dimensions, decide how to display); show_image
  plt.rcParams['figure.figsize'] = [16, 8]
  arrays = []

  check_status(response)
  content_type = response.headers['Content-Type']
  print('Content-type: ', content_type)

  if content_type == 'binary/octet-stream' or content_type == 'application/octet-stream':
    print('WARNING: Let service developer know to set their content_type correctly!')

  if content_type == 'application/x-netcdf' or content_type == 'application/x-netcdf4' or content_type == 'application/netcdf' or content_type == 'binary/octet-stream' or content_type == 'application/octet-stream':
    # Show NetCDF4
    data = H5File(BytesIO(response.content), 'r')

    # If user didn't provide any specific vars to plot check if the nc4 file has rgba vars to plot,
    # otherwise pull all vars into varList
    is_default_rgba = False
    if (len(varList) == 0):
      rgba_var_list = []
      # don't look for an exact match (e.g. == 'red_var')
      # so that we can find subsetted rgba vars as well, which have longer var names
      for rgba_key in ['red', 'green', 'blue', 'alpha']:
        var_iterator = (data_key for data_key in data.keys() if rgba_key in data_key)
        rgba_var_list.append(next(var_iterator, None))
      if None in varList:
        varList = data.keys()
      else:
        is_default_rgba = True
        varList = rgba_var_list

    #Plot the variables requested
    for var in varList:
      if var in data and len(data[var].shape) > 0:
          if (len(data[var].shape) < 3) and not is_default_rgba:
            #Simple plot for 1D or 2D
            plt.plot(data[var])
            plt.show()
          else:
            #Setup for 3D display
            ds = data[var]
            sliced_ds = None
            if len(ds.shape) > 2:
              sliced_ds = ds[0,:]
            else:
              sliced_ds = ds[:,:]
            values = sliced_ds
            if flip:
              values = np.flip(sliced_ds, 0)
            where = (values != ds.attrs.get('_FillValue', None))
            scale = ds.attrs.get('scale_factor', [1])[0]
            offset = ds.attrs.get('add_offset', [0])[0]
            array = np.where(where, values * scale + offset, values)
            arrays.append(array)
      else:
        print('Error: ', var, 'not found in dataset')
    if (len(arrays) != 0):
      #plot the 3D data
      plt.imshow(np.dstack(arrays))
  elif content_type in ['application/zip', 'application/shapefile+zip']:
    # Show ESRI Shapefiles
    tmp = tempfile.NamedTemporaryFile(suffix='.shp.zip', delete=False)
    try:
      tmp.write(response.content)
      show_shape('zip://' + tmp.name, immediate)
    finally:
      os.unlink(tmp.name)
  elif 'application/json' in content_type:
    #Most likely an error
    print(response.json())
    assert(False)
  else:
    # Show Images
    if color_index == None:
      plt.imshow(Image.open(BytesIO(response.content)))
    else:
      gray_image = Image.open(BytesIO(response.content))
      # Move 1-channel green_var TIFF to second channel of RGB
      image = Image.new('RGB', gray_image.size)
      # There's probably a better way to do this with numpy
      if color_index == 0:
        image.putdata([(g, 0, 0) for g in gray_image.getdata()])
      if color_index == 1:
        image.putdata([(0, g, 0) for g in gray_image.getdata()])
      if color_index == 2:
        image.putdata([(0, 0, g) for g in gray_image.getdata()])
      plt.imshow(image)
  if immediate:
    plt.show()

def get_data_urls(response):
  """Returns the data URLs in an async response

  Arguments:
      response {response.Response} -- The async job response

  Returns:
      string[] -- An array of URLs for data links
  """
  return [link['href'] for link in response.json()['links'] if link.get('rel', 'data') == 'data']

def show_async(response, varList = []):
  """Shows an asynchronous Harmony response.

  Polls the output, displaying it as it changes, displaying any http data
  links in the response as they arrive, and ultimately ending once the request
  is successful or failed

  Arguments:
      response {response.Response} -- the response to display
      varList {array} -- If set, only plot the variables listed in varList.  Otherwise, plot all.

  Returns:
      response.Response -- the response from the final successful or failed poll
  """
  def show_response(response, link_count):
    print('Async response at', datetime.now().strftime("%H:%M:%S"))
    print(json.dumps(response.json(), indent=2))
    links = get_data_urls(response)
    new_links = links[slice(link_count, None)]
    for link in new_links:
      if link.startswith('http'):
        show(get(link), varList)
    return len(links)

  check_status(response)
  displayed_link_count = 0
  body = response.json()
  displayed_link_count = show_response(response, displayed_link_count)
  waiting_message_printed = False
  while body['status'] not in ['successful', 'failed', 'canceled']:
    if not waiting_message_printed:
      print('Waiting for updates...')
      waiting_message_printed = True
    sleep(1)
    progress = body['progress']
    status = body['status']
    response = session.get(response.url)
    body = response.json()
    if progress != body['progress'] or status != body['status']:
      displayed_link_count = show_response(response, displayed_link_count)
      waiting_message_printed = False

  assert(body['status'] not in ['failed'])
  check_stac(response)
  print('Async request is complete')
  return response

def print_async_status(body):
  """Prints the status, progress and any messages for the async job

  Arguments:
      body {json} -- the response body to display

  """
  print('JobID:',body['jobID'],'Status:',body['status'],'(',body['progress'],'%) Messages:', body['message'])

def show_async_condensed(response, varList = [], show_results=True):
  """Shows a condensed version of the asynchronous Harmony response.  Useful for getting status if you don't care about the results.

  Polls the output, displaying status as it changes, and ultimately ending once the request
  is successful or failed

  Arguments:
      response {response.Response} -- the response to display
      varList {array} -- If set, only plot the variables listed in varList.  Otherwise, plot all.
      show_results {bool} -- True will display the results as they arrive.  (default: {True})
  """
  def show_response_condensed(response, varList, link_count):
    links = get_data_urls(response)
    new_links = links[slice(link_count, None)]
    for link in new_links:
      if link.startswith('http'):
        show(get(link), varList)
    return len(links)

  check_status(response)

  displayed_link_count = 0
  body = response.json()
  print ('Getting results for request')
  print_async_status(body)
  if show_results:
    displayed_link_count = show_response_condensed(response, varList, displayed_link_count)
  while body['status'] not in ['successful', 'failed', 'canceled']:
    sleep(0.5)
    progress = body['progress']
    status = body['status']
    response = session.get(response.url)
    body = response.json()
    if progress != body['progress'] or status != body['status']:
      if show_results:
        displayed_link_count = show_response_condensed(response, varList, displayed_link_count)
      print_async_status(body)

  assert(body['status'] not in ['failed'])
  check_stac(response)
  print('Async request is complete')

def check_bbox_subset(response, req_lat_min, req_lat_max, req_lon_min, req_lon_max):
  """Asserts if the spatial extents of the data in the response are within the requested bbox of a spatial subset

  #####  CHECK_BBOX_SUBSET currently is not in use; placeholder for the next round of regression test work

  Arguments:
      response {response.Response} -- the response to display
      req_lat_min -- The minimimum latitude from the request bbox for a spatial subset
      req_lat_max -- The maximum latitude from the request bbox for a spatial subset
      req_lon_min -- The minimimum longitude from the request bbox for a spatial subset
      req_lon_max -- The maximum longitude from the request bbox for a spatial subset
  """

  data = H5File(BytesIO(response.content), 'r')

  attr_data = data['lat'][:]

  print('Orig min and max: ', attr_data.min(), attr_data.max() )
  lat_min = (attr_data.min() + 180) % 360 - 180
  lat_max = (attr_data.max() + 180) % 360 - 180
  print(lat_min)
  print(lat_max)

  assert lat_max <= req_lat_max
  assert lat_min >= req_lat_min

  attr_data = data['lon'][:]
  lon_min = (attr_data.min() + 180) % 360 - 180
  lon_max = (attr_data.max() + 180) % 360 - 180
  print(lon_min)
  print(lon_max)

  assert lon_max <= req_lon_max
  assert lon_min >= req_lon_min

def check_status(response):
  """Asserts if the response is a 200, if not, print out the response code

  Arguments:
      response {response.Response} -- the response to display
  """
  if (response.status_code != 200):
    errStr = 'Request failed with status code ' + str(response.status_code)
    assert False, errStr

def check_stac(response):
  """Asserts if the response contains a valid STAC catalog and prints it out.  More robust assertions could
    be done here in the future to confirm that the STAC metadata is valid per the request parameters

  Arguments:
      response {response.Response} -- the response to display
  """
  for i in range(len(response.json()['links'])):
    if response.json()['links'][i]['title'] == 'STAC catalog':
        stac_url = response.json()['links'][i]['href']

  assert(stac_url)
  cat = Catalog.open(stac_url)

  for i in cat.items():
    assert(i.id)
    assert(i.datetime)
    assert(i.bbox)
    assert(i.assets.keys())
    print('STAC Item')
    print('\t', 'ID:', i.id)
    print('\t', 'Date:', i.datetime)
    print('\t', 'Bounding Box:', i.bbox)
    print('\t', 'File:', list(i.assets.keys()))