import pandas as pd
import numpy as np
import math
import matplotlib.pyplot as plt
import seaborn as sns
import yahoo_finance as yf
import pandas_datareader as pdr
from datetime import date
from sklearn.cluster import KMeans
from statsmodels.tsa.api import adfuller
from sklearn.preprocessing import StandardScaler
from dateutil.relativedelta import relativedelta
from statsmodels.regression.rolling import RollingOLS
import plotly.express as px
import plotly.graph_objects as go

# Modify some settings
plt.rcParams['figure.figsize'] = (15, 7)
plt.rcParams['font.size'] = 16
plt.rcParams['figure.dpi'] = 150

pd.options.display.max_rows = 20
pd.options.display.max_columns = 15

class StatArbBuilder:
    def __init__(self,
                 s1,
                 s2,
                 start,
                 end,):
        self._s1 = s1 # name of stock one
        self._s2 = s2 # name of stock two
        self._start = start
        self._end = end
        self._valid_pair = True
        try:
            self._df1 = pdr.get_data_yahoo(self._s1, self._start, self._end) # dataframe of stock one
        except:
            print("Can't get data for {}".format(self._s1))
            self._valid_pair = False
        try:
            self._df2 = pdr.get_data_yahoo(self._s2, self._start, self._end) # dataframe of stock two
        except:
            print("Can't get data for {}".format(self._s2))
            self._valid_pair = False

    def build(self):

        return self._valid_pair

class StatArb(object):
    '''
    Attributes
    ----------
        s1: str
            name of stock 1
        s2: str
            name of stock 2
        start: datetime
            start date of backtest period
        end: datetime
            end date of backtest period
        ma: int
            moving average window used in calc of rolling z-score (default = 28)
        floor / ceiling: float
            trade entry criteria (default = +/-2.0)
        stop_loss_long / short: float
            trade stop loss criteria (default = +/-3.0)
        beta_lookback: int
            lookback window used when calculating hedge ratio (default = 14)
        allocation: int
            dollar value of funds to use (default = 10,000)
        exit_zscore: float
            Z-score at which we close position (default = 0)
        show_plot: boolean
            define if you want to see charts
    ----------

    Methods
    -------
        show_pair()
            displays daily price charts of chosen stocks through chosen dates
        test_coint()
            returns if pair is cointergrated to quickly check if its worth back testing
        backtest_pair()
            backtests pair using all methods below
        generate_spread()
            takes daily closing prices, caluclating the hedge ratios and spreads of the assets
        generate_cointergration()
            Checks if the pair is currently cointergrated and to what significance
        generate_signal()
            takes spread and generates signals indicating trade entry / exit / stop out for long and short positions
        generate_trades()
            defines the postions of both stocks in the pair. i.e. long trade, stock 1 = positive, stock 2 = negative
        generate_portfolio()
            shows our portfolio returns and equity curve
        generate_order_book()
            shows the dates on which we would of placed orders (buy or sell) for each stock and quantity of shares for each
        generate_metrics()
            Shows various portfolio and return metrics e.g. sharp ratio, avg wins / lossesetc
    -------
    '''
    def __init__(self, s1, s2, start, end, ma=28, floor=-2.0, ceiling=2.0, stop_loss_long=-3.0, stop_loss_short=3.0, beta_lookback=28, allocation=10000, exit_zscore=0, show_plot = True):
        self.s1 = s1 # name of stock one
        self.s2 = s2 # name of stock two
        self.df1 = pdr.get_data_yahoo(s1, start, end) # dataframe of stock one
        self.df2 = pdr.get_data_yahoo(s2, start, end) # dataframe of stock two
        self.df = pd.DataFrame(index = self.df1.index) # new df for data_cleaning method
        self.signals = pd.DataFrame(index = self.df1.index)
        self.trades = pd.DataFrame(index = self.df1.index)
        self.portfolio = pd.DataFrame(index = self.df1.index)
        self.metrics = pd.DataFrame(index = self.df1.index)
        self.book = pd.DataFrame(index = range(len(self.df1)))
        self.ma = ma # moving average period
        self.floor = floor # buy threshold for z-score
        self.ceiling = ceiling # sell threshold for z-score
        self.stop_loss_long = stop_loss_long # z-score continues to drop through our floor
        self.stop_loss_short = stop_loss_short # z-score continues to rise through our ceiling
        self.Close = 'Close Long'
        self.Cover = 'Cover Short'
        self.beta_lookback = beta_lookback # lookback of beta for hedge ratio
        self.start = start # begining of test period
        self.end = end # end of test period
        self.exit_zscore = exit_zscore # z-score at which trade is closed
        self.allocation = allocation # dollar value of funds
        self.pair = self.s1 + " and " + self.s2 + " Backtest"
        self.show_plot = show_plot

    def show_pair(self):
        '''
		parameters
		----------
			self: object
                statarb object
		returns
		----------
		'''
        if self.show_plot == False:
            print("show_plot = False")
        else:
            with plt.style.context(['seaborn-paper']):
                ma1 = self.df1['Close'].rolling(window=self.ma).mean()
                std1 = self.df1['Close'].rolling(window=self.ma).std()
                upper1 = ma1 + (std1 * 2)
                lower1 = ma1 - (std1 * 2)

                ma2 = self.df2['Close'].rolling(window=self.ma).mean()
                std2 = self.df2['Close'].rolling(window=self.ma).std()
                upper2 = ma2 + (std2 * 2)
                lower2 = ma2 - (std2 * 2)

                plt.plot(self.df1['Close'],label=self.s1)
                plt.plot(upper1, 'r', alpha=0.5)
                plt.plot(lower1, 'r', alpha=0.5)
                plt.plot(ma1, 'r', alpha=0.5)
                plt.fill_between(self.df1.index, upper1, lower1, alpha=0.1)
                plt.plot(self.df2['Close'],label=self.s2)
                plt.plot(upper2, 'g', alpha=0.5)
                plt.plot(lower2, 'g', alpha=0.5)
                plt.plot(ma2, 'g', alpha=0.5)
                plt.fill_between(self.df1.index, upper2, lower2, alpha=0.1)
                plt.title(self.s1 + ' and ' + self.s2 + ' ' + str(self.start) + ' to ' + str(self.end))
                plt.legend(loc=0)
                plt.show()

    def test_coint(self):
        '''
		parameters
		----------
			self: object
                statarb object
		returns
		----------
            nothing
                a check of cointergration. Useful for running through combinations of stocks in a cluster
		'''

        self.generate_spread()
        self.generate_cointergration()

    def backtest_pair(self):
        '''
		parameters
		----------
			self: object
                statarb object
		returns
		----------
            nothing
                fully backtests a given pair of stocks
		'''

        self.get_spread = self.generate_spread()
        self.get_coint = self.generate_cointergration()
        self.get_signals = self.generate_signal()
        self.get_trades = self.generate_trades()
        self.get_portfolio = self.generate_portfolio()
        self.get_book = self.generate_order_book()
        self.get_metrics = self.generate_metrics()

    def generate_spread(self):
        '''
		parameters
		----------
			self: object
                statarb object
		returns
		----------
			self: df
                dataframe with chosen stock prices cov, var, beta, hedge ratio and spreads
		'''

        # take closing price of chose stocks and add to new dataframe
        self.df[self.s1] = self.df1['Close']
        self.df[self.s2] = self.df2['Close']

        # find beta values for pair of stocks. This beta is not the CAPM beta,
        # its the general beta representing the partial slope coeffiecient in a multivariate
        # (in this case univariate) regression. Given this, it also represents the min variance hedge ratio.
        # this is a rolling regression.
        ############# need to experiment with different values for lookback window

        rolling_model = RollingOLS(self.df[self.s1], self.df[self.s2], window=self.beta_lookback)
        fitted = rolling_model.fit()

        self.df['Hedge_Ratio'] = fitted.params

        # the spread. For each stock_1 purchased we sell n * stock_2 where n is our hedge ratio
        # If the stocks are cointegrated, it implies the spread equation is stationary, I.E. mean and var are same over time
        # if we choose a hedge ratio such that the spread = 0, if there is cointegration the expected value of
        # the spread will stay = 0. Therefore, any deviation from this will present an opportunity of Stat Arb
        # We check for cointegration in the next method.
        self.df['Spread'] = np.log(self.df[self.s1]) - (self.df['Hedge_Ratio']*np.log(self.df[self.s2]))
        self.df['Spread_2'] = np.log(self.df[self.s2]) - (self.df['Hedge_Ratio']*np.log(self.df[self.s1]))

        return self.df

    def generate_cointergration(self):
        '''
		parameters
		----------
			self: object
                statarb object
		returns
		----------
		'''
        ####### can imporve this method with matrix vectorization
        # find coint
        # tells us on a given confidence level weather the par is cointegrated and thus stationary
        adf = adfuller(self.df['Spread'].dropna())
        # print appropriate response
        if adf[0] < adf[4]['1%']:
            print('Spread is Cointegrated at 1% Significance Level')
        elif adf[0] < adf[4]['5%']:
            print('Spread is Cointegrated at 5% Significance Level')
        elif adf[0] < adf[4]['10%']:
            print('Spread is Cointegrated at 10% Significance Level')
        else:
            print('Spread is not Cointegrated')
        return

    def generate_signal(self):
        '''
		parameters
		----------
			self: object
                statarb object
		returns
		----------
			signas: df
                dataframe with 1 or 0 values for long, short, and exit signals, and in position markers
		'''
        # use z scores to generate buy, sell, exit signals
        # floor and ceiling threshold should be between 1.5 and 2 sigma (change depending on backtest results)
        # LONG SIGNAL = LONG THE SPREAD: BUY STOCK 1, SELL STOCK 2
        # SHORT SIGNAL = SHORT THE SPREAD: SELL STOCK 1, BUY STOCK 2

        # with an assumed distribution of spread ~N(0, 1), its is easy to form threshold levels
        # these thresholds will act as signal levels
        # Z = (X - mean) / SD
        # given time series mean and SD will be rolling, using a moving average window
        # create stock z score of the pair spread
        self.signals['Z_Score'] = ((self.df['Spread'] - self.df['Spread'].rolling(window = self.ma).mean())/
                                   (self.df['Spread'].rolling(window = self.ma).std()))

        # create prior stock z score
        self.signals['Prior_Z_Score'] = self.signals['Z_Score'].shift(1)

        self.Short_Signal = False
        self.Long_Signal = False
        self.In_Short = False
        self.In_Long = False
        self.Stopped_Short = False
        self.Stopped_Long = False

        self.signals['Short_Signal'] = 0.0
        self.signals['Long_Signal'] = 0.0
        self.signals['In_Short'] = 0.0
        self.signals['In_Long'] = 0.0
        self.signals['Cover_Short'] = 0.0
        self.signals['Close_Long'] = 0.0

        for i, j in enumerate(self.signals.iterrows()):
            current_z = j[1]['Z_Score']
            # are we already in a short trade?
            if self.In_Short == True:
                # heave we been stopped out already?
                # define stop loss criteria
                if current_z >= self.stop_loss_short:
                    # exit trade if stop loss hit
                    # indicate we have been stopped out
                    self.In_Short = False
                    self.Stopped_Short = True
                    self.signals.iloc[i]['In_Short'] = 0.0
                    self.Short_Signal = False
                    self.signals.iloc[i]['Cover_Short'] = 1.0
                # if not stopped, have we hit close criteria?
                elif current_z <= self.exit_zscore:
                    self.In_Short = False
                    self.signals.iloc[i]['In_Short'] = 0.0
                    self.Short_Signal = False
                    self.signals.iloc[i]['Cover_Short'] = 1.0
                # if not stopped and not closed, still in trade
                else:
                    self.signals.iloc[i]['In_Short'] = 1.0
            else:
                # why are we not in a short
                # have we been stopped out or did we close position?
                if self.Stopped_Short == True:
                    self.In_Short = False
                    self.Short_Signal = False
                    self.signals.iloc[i]['In_Short'] = 0.0
                    # reset cover
                    self.signals.iloc[i]['Cover_Short'] = 0.0
                    # if stopped, wait untill we reach exit critera to re-enter teade
                    if current_z <= self.exit_zscore:
                        self.Stopped_Short = False
                # define trade entry criteria
                elif current_z >= self.ceiling:
                    self.In_Short = True
                    self.Stopped_Short = False
                    self.signals.iloc[i]['In_Short'] = 1.0
                    self.signals.iloc[i]['Cover_Short'] = 0.0
                    if self.Short_Signal == False:
                        self.signals.iloc[i]['Short_Signal'] = 1.0
                        self.Short_Signal = True

            # are we already in a long trade?
            if self.In_Long == True:
                # define stop loss criteria
                if current_z <= self.stop_loss_long:
                    # exit trade if stop loss hit
                    # indicate we have been stopped out
                    self.In_Long = False
                    self.Stopped_Long = True
                    self.signals.iloc[i]['In_Long'] = 0.0
                    self.Long_Signal = False
                    self.signals.iloc[i]['Close_Long'] = 1.0
                # if not stopped, have we hit close criteria?
                elif current_z >= self.exit_zscore:
                    self.In_Long = False
                    self.signals.iloc[i]['In_Long'] = 0.0
                    self.Long_Signal = False
                    self.signals.iloc[i]['Close_Long'] = 1.0
                # if not stopped and not closed, still in trade
                else:
                    self.signals.iloc[i]['In_Long'] = 1.0
            else:
                # why are we not in a long
                # have we been stopped out or did we close position?
                if self.Stopped_Long == True:
                    self.In_Long = False
                    self.Long_Signal = False
                    self.signals.iloc[i]['In_Long'] = 0.0
                    # reset close
                    self.signals.iloc[i]['Close_Long'] = 0.0
                    # if stopped, wait untill we reach exit critera to re-enter teade
                    if current_z >= self.exit_zscore:
                        self.Stopped_Long = False
                # define trade entry criteria
                elif current_z <= self.floor:
                    self.In_Long = True
                    self.Stopped_Long = False
                    self.signals.iloc[i]['In_Long'] = 1.0
                    self.signals.iloc[i]['Close_Long'] = 0.0
                    if self.Long_Signal == False:
                        self.signals.iloc[i]['Long_Signal'] = 1.0
                        self.Long_Signal = True

        self.df['Floor'] = self.floor
        self.df['Ceiling'] = self.ceiling
        self.df['Long_Stop_Loss'] = self.stop_loss_long
        self.df['Short_Stop_Loss'] = self.stop_loss_short
        self.signals['exit_zscore'] = self.exit_zscore

        if self.show_plot == False:
            print("show_plot = False")
        else:

            self.chart = ChartBuilder(self.s1, self.s2, self.start, self.end)
            self.chart.plot(self.signals, pd.DataFrame(), 'Z_Score')

        return self.signals

    def generate_trades(self):
        '''
		parameters
		----------
			self: object
                statarb object
		returns
		----------
			trades: df
                dataframe with negative or positive price values for each stock depending on long or short trade
                total indicates net value of trade
		'''
        self.trades['Positions'] = self.signals['In_Long'] - self.signals['In_Short']
        # Long stock shows negative value = price to represent cash outflow of bought share
        # short stock shows positive value = price to represent cash inflow of borrowed shares
        self.trades[self.s1] = -1 * (self.df[self.s1] * self.trades['Positions'])
        self.trades[self.s2] = (self.df[self.s2] * self.trades['Positions'])
        # Total shows current cumulative value of positions
        self.trades['Total'] = self.trades[self.s1] + self.trades[self.s2]

        return self.trades

    def generate_portfolio(self):
        '''
		parameters
		----------
			self: object
                statarb object
		returns
		----------
			portfolio: df
                dataframe with return values for trades, and overall portfolio vlaue
		'''

        # create percentage returns stream
        # find daily change of total market value of positions
        self.portfolio['Returns'] = self.trades['Total'].pct_change()
        self.portfolio['Returns'].fillna(0.0, inplace = True)
        # account for % changes where start = 0 as this would be infinite
        self.portfolio['Returns'].replace([np.inf, -np.inf], 0.0, inplace = True)
        self.portfolio['Returns'].replace(-1.0, 0.0, inplace = True)

        # getting equity curve
        self.portfolio['Cumulative_Returns'] = (self.portfolio['Returns'] + 1.0).cumprod()
        self.portfolio['Portfolio_Value'] = (self.allocation * self.portfolio['Cumulative_Returns'])
        self.portfolio['Portfolio_Returns'] = self.portfolio['Portfolio_Value'].pct_change()
        self.portfolio['Allocation'] = self.allocation

        self.max_returns = np.fmax.accumulate(self.portfolio['Returns'])
        self.portfolio['Max_Drawdown'] = (self.portfolio['Returns'] / self.max_returns) - 1

        self.portfolio['In_Long'] = self.signals['In_Long']
        self.portfolio['In_Short'] = self.signals['In_Short']

        # plot portfolio valuation
        if self.show_plot == False:
            print("show_plot = False")
        else:

            self.chart = ChartBuilder(self.s1, self.s2, self.start, self.end)
            self.chart.plot(self.signals, self.portfolio, 'Portfolio_Value')

        return self.portfolio

    def generate_order_book(self):
        '''
		parameters
		----------
			self: object
                statarb object
		returns
		----------
			book: df
                dataframe showing dates that a buy or sell order is made and for which stock at which quantity
		'''

        self.book = pd.DataFrame(index = range(len(self.signals)))
        #self.book = self.bk
        self.book['Date'] = ""
        self.book['Position_Type'] = ""
        self.book['Stock'] = ""
        self.book['Order_Type'] = ""
        self.book['Qty'] = 0
        self.book['Price'] = 0
        self.book['Profit/Loss'] = 0
        count = 0
        for i, j in enumerate(self.signals.iterrows()):

            if self.signals.iloc[i]['Long_Signal'] == 1.0:
                # set type of order
                self.book.loc[count, 'Date'] = self.signals.index[i]
                self.book.loc[count, 'Position_Type'] = "Long"
                self.book.loc[count, 'Stock'] = self.s1
                self.book.loc[count, 'Order_Type'] = "Buy"
                self.long_quant = np.floor((self.allocation) / self.df.iloc[i][self.s1])
                self.book.loc[count, 'Qty'] = self.long_quant
                self.enter_long_stock_1_price = self.df.iloc[i][self.s1]
                self.book.loc[count, 'Price'] = self.enter_long_stock_1_price

                self.book.loc[count+1, 'Date'] = self.signals.index[i]
                self.book.loc[count+1, 'Position_Type'] = "Long"
                self.book.loc[count+1, 'Stock'] = self.s2
                self.book.loc[count+1, 'Order_Type'] = "Sell"
                self.short_quant = np.floor(self.long_quant * self.df.iloc[i]['Hedge_Ratio'])
                self.book.loc[count+1, 'Qty'] = self.short_quant
                self.enter_long_stock_2_price = self.df.iloc[i][self.s2]
                self.book.loc[count+1, 'Price'] = self.enter_long_stock_2_price
                count += 2
            elif self.signals.iloc[i]['Close_Long'] == 1.0:
                self.book.loc[count, 'Date'] = self.signals.index[i]
                self.book.loc[count, 'Position_Type'] = "Close Long"
                self.book.loc[count, 'Stock'] = self.s1
                self.book.loc[count, 'Order_Type'] = 'Sell'
                self.book.loc[count, 'Qty'] = self.long_quant
                self.exit_long_stock_1_price = self.df.iloc[i][self.s1]
                self.book.loc[count, 'Price'] = self.exit_long_stock_1_price
                self.book.loc[count, 'Profit/Loss'] = ((self.exit_long_stock_1_price - self.enter_long_stock_1_price)
                                                       *self.long_quant)

                self.book.loc[count+1, 'Date'] = self.signals.index[i]
                self.book.loc[count+1, 'Position_Type'] = "Close Long"
                self.book.loc[count+1, 'Stock'] = self.s2
                self.book.loc[count+1, 'Order_Type'] = 'Buy'
                self.book.loc[count+1, 'Qty'] = self.short_quant
                self.exit_long_stock_2_price = self.df.iloc[i][self.s2]
                self.book.loc[count+1, 'Price'] = self.exit_long_stock_2_price
                self.book.loc[count+1, 'Profit/Loss'] = -1*((self.exit_long_stock_2_price - self.enter_long_stock_2_price)
                                                            *self.short_quant)
                count += 2

            if self.signals.iloc[i]['Short_Signal'] == 1.0:
                # set type of order
                self.book.loc[count, 'Date'] = self.signals.index[i]
                self.book.loc[count, 'Position_Type'] = "Short"
                self.book.loc[count, 'Stock'] = self.s2
                self.book.loc[count, 'Order_Type'] = "Buy"
                self.long_quant = np.floor((self.allocation) / self.df.iloc[i][self.s2])
                self.book.loc[count, 'Qty'] = self.long_quant
                self.enter_long_stock_2_price = self.df.iloc[i][self.s2]
                self.book.loc[count, 'Price'] = self.enter_long_stock_2_price

                self.book.loc[count+1, 'Date'] = self.signals.index[i]
                self.book.loc[count+1, 'Position_Type'] = "Short"
                self.book.loc[count+1, 'Stock'] = self.s1
                self.book.loc[count+1, 'Order_Type'] = "Sell"
                self.short_quant = np.floor(self.long_quant * self.df.iloc[i]['Hedge_Ratio'])
                self.book.loc[count+1, 'Qty'] = self.short_quant
                self.enter_long_stock_1_price = self.df.iloc[i][self.s1]
                self.book.loc[count+1, 'Price'] = self.enter_long_stock_1_price
                count += 2

            elif self.signals.iloc[i]['Cover_Short'] == 1.0:
                self.book.loc[count, 'Date'] = self.signals.index[i]
                self.book.loc[count, 'Position_Type'] = "Close Short"
                self.book.loc[count, 'Stock'] = self.s2
                self.book.loc[count, 'Order_Type'] = 'Sell'
                self.book.loc[count, 'Qty'] = self.long_quant
                self.exit_long_stock_2_price = self.df.iloc[i][self.s2]
                self.book.loc[count, 'Price'] = self.exit_long_stock_2_price
                self.book.loc[count, 'Profit/Loss'] = ((self.exit_long_stock_2_price - self.enter_long_stock_2_price)
                                                       *self.long_quant)

                self.book.loc[count+1, 'Date'] = self.signals.index[i]
                self.book.loc[count+1, 'Position_Type'] = "Close Short"
                self.book.loc[count+1, 'Stock'] = self.s1
                self.book.loc[count+1, 'Order_Type'] = 'Buy'
                self.book.loc[count+1, 'Qty'] = self.short_quant
                self.exit_long_stock_1_price = self.df.iloc[i][self.s1]
                self.book.loc[count+1, 'Price'] = self.exit_long_stock_1_price
                self.book.loc[count+1, 'Profit/Loss'] = -1*((self.exit_long_stock_1_price - self.enter_long_stock_1_price)
                                                            *self.short_quant)
                count += 2

        self.book['Cum_Sum_P_L'] = self.book['Profit/Loss'].cumsum()
        self.book = self.book.drop(self.book.index[count:])
        return self.book

    def generate_metrics(self):
        '''
		parameters
		----------
			self: object
                statarb object
		returns
		----------
			metrics: df
                dataframe of performance metrics
		'''
        # calculate summary statistics
        self.mu = (self.portfolio['Returns'].mean())
        self.sigma = (self.portfolio['Returns'].std())
        self.sharpe = (self.mu - 0.005) / self.sigma
        # where True, yield x, otherwise yield y
        self.wins = (np.where(self.portfolio['Cumulative_Returns'] > 0.0, 1.0, 0.0)).sum()
        self.losses = (np.where(self.portfolio['Cumulative_Returns'] < 0.0, 1.0, 0.0)).sum()
        if self.losses == 0:
            self.total_trades = self.wins
        else:
            self.total_trades = self.wins + self.losses
        # win loss ratio
        if self.losses == 0:
            self.wl_ratio = 1
        else:
            self.wl_ratio = (self.wins / self.losses)
        # probability of win and loss
        self.p_win = (self.wins / self.total_trades)
        if self.losses == 0:
            self.p_loss = 0
        else:
            self.p_loss = (self.losses / self.total_trades)
        # avg win / loss return
        self.avg_win_return = (self.portfolio['Cumulative_Returns'] > 0.0).mean()
        self.avg_loss_return = (self.portfolio['Cumulative_Returns'] < 0.0).mean()
        # payout ratio
        if self.avg_loss_return == 0:
            self.payout_ratio = 1
        else:
            self.payout_ratio = (self.avg_win_return / self.avg_loss_return)

        self.difference_in_years = relativedelta(self.end, self.start).years

        self.metrics['CAGR'] = (((self.portfolio.iloc[-1]['Portfolio_Value']/self.portfolio.iloc[0]['Portfolio_Value'])
                                 **(1/self.difference_in_years)) - 1)
        self.metrics['Sharpe Ratio'] = self.sharpe
        self.metrics['Wins'] = self.wins
        self.metrics['P(Wins)'] = self.p_win
        self.metrics['Avg_Win_Return'] = self.avg_win_return
        self.metrics['Losses'] = self.losses
        self.metrics['P(Loss)'] = self.p_loss
        self.metrics['Avg_Loss_Return'] = self.avg_loss_return
        self.metrics['WL_Ratio'] = self.wl_ratio

        return self.metrics

class ChartBuilder(object):

    def __init__(self, s1, s2, start, end, exit=0, long_stop = -3, short_stop = 3, long_enter = -2, short_enter = 2):

        self.s1 = s1
        self.s2 = s2
        self.start = start
        self.end = end
        self.exit = exit
        self.long_stop = long_stop
        self.short_stop = short_stop
        self.long_enter = long_enter
        self.short_enter = short_enter

        self.pair = str("{} and {} ".format(self.s1, self.s2))
        self.scale = str(" ({} - {})".format(self.start, self.end))

# Customized Area chart

    def plot(self, signal_data, metric_data, metric, add_signal_levels = True, add_trades = True):

        self.signal_data = signal_data
        self.metric_data = metric_data
        self.data = pd.concat([self.signal_data, self.metric_data], axis = 1)
        self.metric = str(metric)
        self.add_signal_levels = add_signal_levels
        self.add_trades = add_trades
        self.title = self.pair + self.metric + self.scale
        self.y0 = self.data[self.metric].max()
        self.y1 = self.data[self.metric].min()
        c_line = px.line(x = self.data.index, y = self.data[self.metric], title = self.title)

        c_line.update_xaxes(
            title_text = 'Date',
            rangeslider_visible = False,
            rangeselector = dict(
                buttons = list([
                    dict(count = 1, label = '1M', step = 'month', stepmode = 'backward'),
                    dict(count = 6, label = '6M', step = 'month', stepmode = 'backward'),
                    dict(count = 1, label = 'YTD', step = 'year', stepmode = 'todate'),
                    dict(count = 1, label = '1Y', step = 'year', stepmode = 'backward'),
                    dict(step = 'all')])))

        if self.add_signal_levels == True:
            if self.metric == 'Z_Score':
                c_line.add_shape(type = "line", x0 = self.start, x1 = self.end, y0 = self.exit, y1 = self.exit, line=dict(color = 'orange', width = 2, dash = 'dash'))
                c_line.add_shape(type = "line", x0 = self.start, x1 = self.end, y0 = self.short_stop, y1 = self.short_stop, line=dict(color = 'red', width = 2, dash = 'dash'))
                c_line.add_shape(type = "line", x0 = self.start, x1 = self.end, y0 = self.long_stop, y1 = self.long_stop, line=dict(color = 'red', width = 2, dash = 'dash'))
                c_line.add_shape(type = "line", x0 = self.start, x1 = self.end, y0 = self.long_enter, y1 = self.long_enter, line=dict(color = 'green', width = 2, dash = 'dash'))
                c_line.add_shape(type = "line", x0 = self.start, x1 = self.end, y0 = self.short_enter, y1 = self.short_enter, line=dict(color = 'green', width = 2, dash = 'dash'))

        if self.add_trades == True:
            enter = self.data['Short_Signal'].gt(0.5)
            exit = self.data['Cover_Short'].gt(0.5)
            shorts = pd.concat([enter, exit], axis=1)
            shorts = shorts[~((~shorts).all(axis=1))]
            shorts.reset_index(level=0, inplace=True)
            ff = []
            ll = []
            f = shorts.loc[shorts['Short_Signal'] == True, 'Date']
            l = shorts.loc[shorts['Cover_Short'] == True, 'Date']
            run_f = True
            run_l = True
            if shorts.iloc[0]['Short_Signal'] == False:
                ff.append(self.start)
                for i in f:
                    ff.append(i)
                run_f = False
            elif shorts.iloc[-1]['Short_Signal'] == True:
                for i in l:
                    ll.append(i)
                ll.append(self.end)
                run_l = False

            if run_f == True:
                for i in f:
                    ff.append(i)
            if run_l == True:
                for i in l:
                    ll.append(i)

            ts = pd.DataFrame(columns = ['first', 'last'])
            ts['first'] = ff
            ts['last'] = ll

            for index, row in ts.iterrows():
                c_line.add_shape(type = 'rect',
                              x0 = row['first'],
                              y0 = self.y0,
                              x1 = row['last'],
                              y1 = self.y1,
                              line=dict(color="rgba(0,0,0,0)", width=3,), fillcolor='red', opacity = 0.2, layer='below')

            enter = self.data['Long_Signal'].gt(0.5)
            exit = self.data['Close_Long'].gt(0.5)
            longs = pd.concat([enter, exit], axis=1)
            longs = longs[~((~longs).all(axis=1))]
            longs.reset_index(level=0, inplace=True)
            ff = []
            ll = []
            f = longs.loc[longs['Long_Signal'] == True, 'Date']
            l = longs.loc[longs['Close_Long'] == True, 'Date']
            run_f = True
            run_l = True
            if longs.iloc[0]['Long_Signal'] == False:
                ff.append(self.start)
                for i in f:
                    ff.append(i)
                run_f = False
            elif longs.iloc[-1]['Long_Signal'] == True:
                for i in l:
                    ll.append(i)
                ll.append(self.end)
                run_l = False

            if run_f == True:
                for i in f:
                    ff.append(i)
            if run_l == True:
                for i in l:
                    ll.append(i)

            ts = pd.DataFrame(columns = ['first', 'last'])
            ts['first'] = ff
            ts['last'] = ll

            for index, row in ts.iterrows():
                c_line.add_shape(type = 'rect',
                              x0 = row['first'],
                              y0 = self.y0,
                              x1 = row['last'],
                              y1 = self.y1,
                              line=dict(color="rgba(0,0,0,0)", width=3,), fillcolor='green', opacity = 0.2, layer='below')

        c_line.update_yaxes(title_text = self.metric)
        c_line.update_layout(showlegend = False,
            title = {
                'text': self.title,
                'y':0.9,
                'x':0.5,
                'xanchor': 'center',
                'yanchor': 'top'})

        c_line.show()


import math
from itertools import chain

import torch
from ppo_pytorch.common.make_grid import make_grid
from ppo_pytorch.actors.actors import ModularActor, ActorOutput
from ppo_pytorch.actors.utils import image_to_float, make_conv_heatmap
from torch import nn as nn, autograd as autograd
from torch.autograd import Variable

class GanCNNActor(ModularActor):
    """
    Convolution network.
    """

    def __init__(self, obs_space, action_space, head_factory, cnn_kind='large',
                 activation=nn.ReLU, actor_includes_conv=True, **kwargs):
        """
        Args:
            obs_space: Env's observation space
            action_space: Env's action space
            head_factory: Function which accept (hidden vector size, `ProbabilityDistribution`) and return `HeadBase`
            cnn_kind: Type of cnn.
                'small' - small CNN from arxiv DQN paper (Mnih et al. 2013)
                'large' - bigger CNN from Nature DQN paper (Mnih et al. 2015)
                'custom' - largest CNN of custom structure
            activation: Activation function
        """
        super().__init__(obs_space, action_space, **kwargs)
        self.activation = activation
        self.actor_includes_conv = actor_includes_conv
        assert cnn_kind in ('small', 'large', 'custom')

        ngf, ndf = 64, 64
        nz = 64
        nc = 4
        self.disc_start_net = nn.Sequential(
            nn.Sequential(
                nn.Conv2d(nc, ndf, 4, 2, 1, bias=True),
                nn.LeakyReLU(0.2, inplace=True),
            ),
            nn.Sequential(
                nn.Conv2d(ndf, ndf * 2, 4, 2, 0, bias=True),
                nn.LeakyReLU(0.2, inplace=True),
            ),
            nn.Sequential(
                nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=True),
                nn.LeakyReLU(0.2, inplace=True),
            ),
            nn.Sequential(
                nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 0, bias=True),
            ),
        )
        self.disc_end_net = nn.Sequential(
            nn.LeakyReLU(0.2),
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=True),
        )
        self.gen_start_net = nn.Sequential(
            nn.Conv2d(ndf * 8, nz, 4, 1, 0, bias=True),
        )
        self.gen_end_net = nn.Sequential(
            nn.Sequential(
                nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),
                nn.BatchNorm2d(ngf * 8),
                nn.LeakyReLU(0.2, inplace=True),
            ),
            nn.Sequential(
                nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 0, bias=False),
                nn.BatchNorm2d(ngf * 4),
                nn.LeakyReLU(0.2, inplace=True),
            ),
            nn.Sequential(
                nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
                nn.BatchNorm2d(ngf * 2),
                nn.LeakyReLU(0.2, inplace=True),
            ),
            nn.Sequential(
                nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 0, bias=False),
                nn.BatchNorm2d(ngf),
                nn.LeakyReLU(0.2, inplace=True),
            ),
            nn.Sequential(
                nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=True),
            ),
        )

        # create head
        self.actor_linear = nn.Sequential(
            nn.Linear(nz, nz*2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(nz*2, nz*2),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.head = head_factory(nz*2, self.pd)

        self.actor_linear_head = nn.Sequential(
            self.actor_linear,
            self.head,
        )

        self.reset_weights()

    def params_disc(self):
        return chain(self.disc_start_net.parameters(), self.disc_end_net.parameters())

    def params_gen(self):
        return chain(self.gen_start_net.parameters(), self.gen_end_net.parameters())

    def params_head(self):
        return chain(self.head.parameters(), self.actor_linear.parameters())

    def params_actor(self):
        if self.actor_includes_conv:
            return chain(self.disc_start_net.parameters(), self.params_head())
        else:
            return self.params_head()

    def reset_weights(self):
        super().reset_weights()
        self.head.reset_weights()

    def forward_on_hidden_code(self, hidden):
        return self.actor_linear_head(hidden)

    def generate(self, hidden_code, state):
        x = hidden_code.contiguous().view(*hidden_code.shape, 1, 1)
        for i, conv in enumerate(self.gen_end_net):
            # run conv layer
            x = conv(x)

        # log
        if self.do_log:
            state = image_to_float(state.data)
            img = torch.cat([x.data.unsqueeze(2), state.unsqueeze(2)], 2)
            img = img[:min(img.shape[0], 2)]
            img = img.view(-1, 1, *img.shape[3:])
            img = img.clamp(-1, 1).div_(2).add_(0.5)
            img = make_grid(img, nrow=round(math.sqrt(img.shape[0])), normalize=False, fill_value=0.1)
            self.logger.add_image('reconstructed state', img, self._step)

        return x

    def discriminate(self, state, detach_hidden_from_conv=False):
        x = state
        for i, layer in enumerate(self.disc_start_net):
            # run conv layer
            x = layer(x)
            # log
            if self.do_log:
                self.log_conv_activations(i, layer[0], x)
                self.log_conv_filters(i, layer[0])

        # if self.do_log:
        #     self.logger.add_histogram('conv linear', x, self._step)

        conv_out = x
        disc = self.disc_end_net(x).view(-1)
        hidden = self.gen_start_net(x.detach() if detach_hidden_from_conv else x).view(x.shape[0], -1)
        return disc, hidden, conv_out

    def forward(self, input) -> ActorOutput:
        log_policy_attention = self.do_log and input.is_leaf
        input = image_to_float(input)
        if log_policy_attention:
            input = Variable(input.data, requires_grad=True)

        disc, hidden, conv_out = self.discriminate(input)

        # hidden = self.conv_actor_linear(conv_out)

        ac_out = self.actor_linear_head(hidden)
        ac_out.conv_out = conv_out
        ac_out.hidden_code = hidden
        ac_out.gan_d = disc

        if log_policy_attention:
            self.log_policy_attention(input, ac_out)

        return ac_out

    def log_conv_activations(self, index: int, conv: nn.Conv2d, x: Variable):
        img = x[0].data.unsqueeze(1).clone()
        img = make_conv_heatmap(img)
        img = make_grid(img, nrow=round(math.sqrt(conv.out_channels)), normalize=False, fill_value=0.1)
        self.logger.add_image('conv activations {} img'.format(index), img, self._step)
        self.logger.add_histogram('conv activations {} hist'.format(index), x[0], self._step)

    def log_conv_filters(self, index: int, conv: nn.Conv2d):
        channels = conv.in_channels * conv.out_channels
        shape = conv.weight.data.shape
        kernel_h, kernel_w = shape[2], shape[3]
        img = conv.weight.data.view(channels, 1, kernel_h, kernel_w).clone()
        max_img_size = 100 * 5
        img_size = channels * math.sqrt(kernel_h * kernel_w)
        if img_size > max_img_size:
            channels = channels * (max_img_size / img_size)
            channels = math.ceil(math.sqrt(channels)) ** 2
            img = img[:channels]
        img = make_conv_heatmap(img, scale=2 * img.std())
        img = make_grid(img, nrow=round(math.sqrt(channels)), normalize=False, fill_value=0.1)
        self.logger.add_image('conv featrues {} img'.format(index), img, self._step)
        self.logger.add_histogram('conv features {} hist'.format(index), conv.weight.data, self._step)

    def log_policy_attention(self, states, head_out):
        states_grad = autograd.grad(
            head_out.probs.abs().mean() + head_out.state_values.abs().mean(), states,
            only_inputs=True, retain_graph=True)[0]
        img = states_grad.data[:4]
        img.abs_()
        img /= img.view(4, -1).pow(2).mean(1).sqrt_().add_(1e-5).view(4, 1, 1, 1)
        img = img.view(-1, 1, *img.shape[2:]).abs()
        # img = make_conv_heatmap(img, scale=2*img.std())
        img = make_grid(img, 4, normalize=True, fill_value=0.1)
        self.logger.add_image('state attention', img, self._step)


from config import Config
from surface_descriptors.mesh import Mesh

if __name__ == '__main__':
    for mesh_num in range(*Config.categories_indices[Config.proc_category]):
        if mesh_num in Config.excluded_set:
            continue

        mesh = Mesh(obj_path=Config.data_path / f'obj/{mesh_num}.obj',
                    save_path=Config.save_path / 'data')
        all_features = mesh.calculate_surface_descriptors(save=True)
        print(f'Shape: {all_features.shape}\nMin: {all_features.min()}\nMax: {all_features.max()}')


from pyspark import SparkConf, SparkContext
conf=SparkConf().setMaster('local').setAppName('FriendsByAge')
sc=SparkContext(conf=conf)

def parseLine(line):
    fields=line.split(',')
    age=int(fields[2])
    numFriends=int(fields[3])
    return (age, numFriends)

lines=sc.textFile("file:///SparkCourse/fakefriends.csv")
rdd=lines.map(parseLine)
totalByAge=rdd.mapValues(lambda x:(x,1)).reduceByKey(lambda x,y:(x[0]+y[0],x[1]+y[1]))
averagesByAge=totalByAge.mapValues(lambda x:x[0]/x[1])
results=averagesByAge.collect()
a=results.sort()

for result in results:
    print ('People in age of: ', result[0], ' have in average: {:.0f}'.format(result[1]), ' friends', sep='')


import  requests
from bs4 import BeautifulSoup
import operator

def start(url):
    word_list = []
    source_code = requests.get(url).text
    soup = BeautifulSoup(source_code, "html.parser")

    for post_text in soup.findAll('a', {'class': 'title text-semibold'}):
        content = post_text.string       # Get only the text
        words = content.lower().split()  # covert all words to lowercase. split() returns a list of the woords seperated by whitespaces

        for each_word in words:
            word_list.append(each_word)
    clean_up_list(word_list)

def clean_up_list(mylist):
    clean_list = []
    for word in mylist:
        symbols = "0123456789!\"£$%^&*(){}_-:@~<>?[];'#,./'"
        for i in range(0, len(symbols)):
            word = word.replace(symbols[i], "")
        if len(word) > 0:
            #print(word)
            clean_list.append(word)
    create_dictionary(clean_list)

def create_dictionary(clean_list):
    word_count = {}

    # The word is the key each time you find a word that is not in the dictinary.
    for word in clean_list:
        if word in word_count:
            word_count[word] += 1   # using the word as the key, increments the value stored
        else:
            word_count[word] = 1    # stores the value one with the wword as its key
    for key, value in sorted(word_count.items(), key=operator.itemgetter(1)):       # 1 is sort by value. O is sort by key
        print(key, value)

start('https://thenewboston.com/forum/category.php?id=15')


from typing import BinaryIO, Union, Optional
import subprocess
import logging

log = logging.getLogger(__name__)

def _run_cmd(args):
    log.info("Running %s", " ".join(args))
    subprocess.run(args, check=True)

# pylint: disable=invalid-name
def dd(file_in: Union[str, BinaryIO], file_out: BinaryIO, out_offset: int, out_size: Optional[int] = None):
    "Run dd with the given arguments"
    # Try to guess block size. We would like to use as big block as
    # possible. But we need take into account that "seek" parameter
    # uses block size as the unit.
    blocksize: int = 65536
    while out_offset % blocksize != 0:
        blocksize //= 2

    if isinstance(file_in, str):
        file_in_path = file_in
    else:
        file_in_path = file_in.name
    args = [
        "dd",
        f"if={file_in_path}",
        f"of={file_out.name}",
        f"bs={blocksize}",
        f"seek={out_offset // blocksize}",
        "status=progress",
        "conv=sparse",
        "conv=notrunc",
    ]  # yapf: disable
    if out_size:
        args.append(f"count={out_size // blocksize}")
    _run_cmd(args)

def simg2img(file_in: Union[str, BinaryIO], file_out: BinaryIO):
    "Run simg2img with the given arguments"
    if isinstance(file_in, str):
        file_in_path = file_in
    else:
        file_in_path = file_in.name
    args = [
        "simg2img",
        file_in_path,
        file_out.name,
    ]  # yapf: disable
    _run_cmd(args)

def mkext4fs(file_out: BinaryIO, contents_dir=None):
    "Create ext4 fs in given file"
    args = ["mkfs.ext4", file_out.name]
    if contents_dir:
        args.append("-d")
        args.append(contents_dir)

    _run_cmd(args)

def mkvfatfs(file_out: BinaryIO, sector_size=None):
    "Create FAT fs in given file"
    args = ["mkfs.vfat"]
    if sector_size:
        args.append("-S")
        args.append(str(sector_size))
    args.append(file_out.name)

    _run_cmd(args)

def mcopy(img: BinaryIO, file: str, name: str):
    "Copy a file to a vfat image with a given name"
    args = ["mcopy", "-i", img.name, file, "::" + name]

    _run_cmd(args)

def mmd(img: BinaryIO, folders: list):
    "Create directories inside a vfat image"
    args = ["mmd", "-i", img.name]
    args.extend(folders)

    _run_cmd(args)


from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure

from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *

import sys
from PyQt5.QtWidgets import *
from PyQt5 import*
from PyQt5 import uic
import numpy as np
from sympy import*
import matplotlib.pyplot as plt
from sympy.core.numbers import Infinity

form_class6 = uic.loadUiType("FinR.ui")[0]

class FinR(QMainWindow, form_class6):

    val4else = ['T(x)','q','q4C','x','L','m','M','t']

    def __init__(self):

        self.graphwindow = Graphwindow()
        self.opengl = Opengl()

        super().__init__()
        self.setupUi(self)

        self.Radiobutton = [self.radioButton_1,self.radioButton_2,self.radioButton_3,
            self.radioButton_4]

        def BC():
            for a,b in zip(self.Radiobutton,range(4)):
                if a.isChecked():
                    self.Radiosel[0] = b
                    print(b)

        self.Radiosel = np.ones(1,dtype= np.int0)*5

        for a in self.Radiobutton:
            a.toggled.connect(BC)

        self.pushButton_1.clicked.connect(self.Calfin)
        self.pushButton_2.clicked.connect(self.show_graph)
        self.pushButton_3.clicked.connect(self.show_opengl)

    def Calfin(self):

        Bc = self.Radiosel

        #공통 변수 설정..
        try:
            w = float(self.lineEdit_w.text())

        except ValueError:
            w = 1

        x = symbols('x')
        L = symbols('L')

        k = float(self.lineEdit_k.text())
        h = float(self.lineEdit_h.text())
        t = float(self.lineEdit_t.text())
        self.val4else[7] = t

        T_b = float(self.lineEdit_tb.text())
        T_a =float(self.lineEdit_ta.text())

        theta_b = T_b - T_a
        P = 2*(w+t)
        A = w*t

    #입력받은 값을 이용해서 m,M,theha_b 를 얻음

        M = self.fun_M(h,P,k,A,theta_b)
        m = self.fun_m(h,P,k,A)
        theta  = self.Case_A_temperature(x,m,L,k,h)

        if Bc == 0:             #BC1

            L = float(self.lineEdit_1_L.text())
            theta  = self.Case_A_temperature(x,m,L,k,h)

            T = theta*(theta_b) + T_a
            q = self.Case_A_heatrate(M,m,L,k,h)

        #따로 저장해 놓을것 --> 다른곳에서 쓸 수 있음
            self.val4else[0] = T
            self.val4else[1] = q
            self.val4else[4] = L

            try:
                x1 = float(self.lineEdit_1_x.text())

                self.textBrowser_1.append(" T({0}) = {1} ".format(str(x1),str(T.subs(x,x1))))
                self.textBrowser_1.append(" T(x) = {0} ".format(str(N(T,3))))
                self.textBrowser_2.append(" q = {0}".format(str(q)))

            except ValueError:

                self.textBrowser_1.append(" T({0}) ={1} ".format((str(x),str(T))))
                self.textBrowser_1.append(" T(x) ={0} ".format(str(N(T,3))))
                self.textBrowser_2.append(" q = {0} ".format(str(q)))

        elif Bc == 1:               #BC2

            L = float(self.lineEdit_2_L.text())
            theta  = self.Case_B_temperature(x,m,L)

            T = theta*(theta_b) + T_a
            q = self.Case_B_heatrate(M,m,L)

        #따로 저장해 놓을것 --> 다른곳에서 쓸 수 있음
            self.val4else[0] = T
            self.val4else[1] = q
            self.val4else[4] = L

            print(self.val4else[0])

            try:
                x1 = float(self.lineEdit_2_x.text())

                self.textBrowser_1.append(" T({0}) = {1} ".format(str(x1), str(T.subs(x,x1))))
                self.textBrowser_1.append(" T(x) = {0} ".format(str(T)))
                self.textBrowser_2.append(" q = {0}".format(str(q)))

            except ValueError:

                self.textBrowser_1.append(" T({0}) ={1} ".format(str(x),str(T.subs(x,x))))
                self.textBrowser_2.append(" q = {0} ".format(str(q)))

        elif Bc == 2:               #BC3

            L   = float(self.lineEdit_3_L.text())
            T_L = float(self.lineEdit_3_T.text())

            theta_L = T_L - T_a
            theta_b = T_b - T_a

            theta  = self.Case_C_temperature(x,m,L,theta_L,theta_b)

            T = theta*(theta_b) + T_a

            q = self.Case_C_heatrate(M,m,L,theta_L,theta_b)

        #따로 저장해 놓을것 --> 다른곳에서 쓸 수 있음

            self.val4else[0] = T
            self.val4else[1] = q[0]
            self.val4else[2] = q[1]
            self.val4else[4] = L

            try:
                x1 = float(self.lineEdit_3_x.text())

                self.textBrowser_1.append(" T({0}) ={1} ".format(str(x1),str(T.subs(x,x1))))
                self.textBrowser_1.append(" T(x) ={0} ".format(str(T)))

                self.textBrowser_2.append(" q0 = {0}".format(str(q[0])))
                self.textBrowser_2.append(" qL = {0}".format(str(q[1])))

            except ValueError:

                self.textBrowser_1.append(" T({0}) ={1} ".format(str(x),str(T.subs(x,x))))
                self.textBrowser_1.append(" T(x) ={0} ".format(str(T)))

                self.textBrowser_2.append(" q0 = {0}".format(str(q[0])))
                self.textBrowser_2.append(" qL = {0}".format(str(q[1])))

        elif Bc == 3:               #BC3

            theta  = self.Case_D_temperature(x,m)

            T = theta*(theta_b) + T_a
            q = self.Case_D_heatrate(M)

            L = 10

        #따로 저장해 놓을것 --> 다른곳에서 쓸 수 있음
            self.val4else[0] = T
            self.val4else[1] = q
            self.val4else[4] = L

            try:
                x1 = float(self.lineEdit_4_x.text())

                self.textBrowser_1.append(" T({0}) ={1} ".format(str(x1),str(T.subs(x,x1))))
                self.textBrowser_1.append(" T(x) ={0} ".format(str(T)))

                self.textBrowser_2.append(" q = {0}".format(str(q)))

            except ValueError:

                self.textBrowser_1.append(" T({0}) ={1} ".format(str(x),str(T.subs(x,x))))
                self.textBrowser_1.append(" T(x) ={0} ".format(str(T)))

                self.textBrowser_2.append(" q = {0} ".format(str(q)))

        else:
            pass

    def fun_m(self,h,P,k,A):
        return sqrt((h*P)/(k*A))

    def fun_M(self,h,P,k,A,theta_b):

        return sqrt(h*P*k*A)*theta_b

    def Case_A_temperature(self,x, m, L, k, h):
        # convection
        d1 = N(cosh(m*(L-x)) + (h/(m*k))*sinh(m*(L-x)),2)
        d2 = N(cosh(m*L) + (h/(m*k))*sinh(m*L),2)

        return N(d1/d2,2)

    def Case_A_heatrate(self,M, m, L, k, h):
        # convection
        d1 = sinh(m*L) + (h/(m*k))*cosh(m*L)
        d2 = cosh(m*L) + (h/(m*k))*sinh(m*L)
        return M*d1/d2

    def Case_B_temperature(self,x, m, L):
        # adiabatic
        return cosh(m*(L-x))/cosh(m*L)

    def Case_B_heatrate(self,M, m, L):
        # adiabatic
        return M*tanh(m*L)

    def Case_C_temperature(self,x, m, L, theta_L, theta_b):
        # prescribed temperature
        d1 = (theta_L/theta_b)*sinh(m*x) + sinh(m*(L-x))
        return d1/sinh(m*L)

    def Case_C_heatrate(self,M, m, L, theta_L, theta_b):
        # prescribed temperature
        q0 = M/sinh(m*L) * (cosh(m*L) - (theta_L/theta_b))
        qL = M/sinh(m*L) * (1 - cosh(m*L)*(theta_L/theta_b))
        return [q0, qL]

    def Case_D_temperature(self,x,m):
        # infinite fin
        return exp(-m*x)

    def Case_D_heatrate(self,M):
        # infinite fin
        return M

    def show_graph(self):
        self.graphwindow.show()

    def show_opengl(self):
        self.opengl.show()

class Graphwindow(QMainWindow):

    def __init__(self):
        super().__init__()

        self.main_widget = QWidget()                            #Qwiget 을 만듬
        self.setCentralWidget(self.main_widget)                 #QMainWindow에 이걸해야지 뭐든지 추가할 수있는거임..

        button = QPushButton("Show Graph")                      #버튼을 생성함
        self.canvas = FigureCanvas(Figure(figsize=(4, 3)))      #canvas도 생성함
        vbox = QVBoxLayout(self.main_widget)                    #layout도 생성함
        self.ax = self.canvas.figure.subplots()

        vbox.addWidget(self.canvas)                             #layout에 canvas를 추가
        vbox.addWidget(button)                                  #이렇게 추가해야지 실제로 적용되는것.
        self.addToolBar(NavigationToolbar(self.canvas, self))   #toolbor 추가

        self.setWindowTitle('Graph')                            #이름,크기지정
        self.setGeometry(300, 100, 600, 400)
        button.clicked.connect(self.update_chart)               #버튼 기능

    def update_chart(self):                                     #그래프 추가하기
        x = symbols("x")
        f = lambdify(x,FinR.val4else[0],"numpy")

        x1 = 0
        x2 = FinR.val4else[4]
        ex = (x2-x1)/20

        x = np.linspace(x1,x2,100)
                                                                #grid 추가하는 방법 찾아보기.....

        self.ax.grid()
        self.ax.set_xlim([x1-ex, x2+ex])
        self.plot = self.ax.plot(x,f(x),'r')                    #이렇게 새로운 값을 대입할 수 있는것

        self.canvas.draw()                                      #show 말고 draw 그림

class Opengl(QOpenGLWidget):                                    #그림그리기.....

    x = symbols('x')

    def  __init__(self):
        super().__init__()

    def initializeGL(self):
        glPolygonMode(GL_FRONT, GL_FILL)
        glPolygonMode(GL_BACK, GL_FILL)

        glShadeModel(GL_SMOOTH)
        glEnable(GL_NORMALIZE)

        glClearDepth(1.0)
        glPolygonMode(GL_FRONT_AND_BACK, GL_FILL)

        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
        glEnable(GL_BLEND)

        glEnable(GL_DEPTH_TEST)
        glEnable(GL_TEXTURE_2D)

    def resizeGL(self, width, height):
        glGetError()

        aspect = width if (height == 0) else width / height

        glViewport(0, 0, width, height)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(45, aspect, 0.1, 1000.0)
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()

    def paintGL(self):

        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glLoadIdentity()

        glClearColor(0.2, 0.2, 0 , 0)                                #바탕화면 색깔
        glPushMatrix()
        glTranslatef(0.0, 0.0, -3.0)                                 #시점조정하기...
        self.Draw()
        glPopMatrix()
        glFlush()

    def Draw(self):

        def Vertices():                                              #좌표......

            x = np.linspace(-1.5,1.5,200)
            ver= []

            t = FinR.val4else[7]
            L = FinR.val4else[4]

            for a in x:
                ver1 = [a,-0.2]
                ver2 = [a, 0.2]
                ver.append(ver1)
                ver.append(ver2)

            return ver

        def Color():                                                   #색배합 조정하기

            expr = lambdify(self.x,FinR.val4else[0],'numpy')
            x = np.linspace(0,FinR.val4else[4],200)
            T = np.array(expr(x))
            T_avg = np.average(T)                   #평균값

            dT = T-T_avg                            # dT= 온도 - 평균
            dT_max = np.max(dT)
            dT_min = np.min(dT)

            color= []

            for T in dT:
                if T > 0:
                    if    dT_max >= T >=  0.93*dT_max:
                        color1 = [253/253,0,0]
                        color2 = [253/253,0,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.93*dT_max> T >= 0.86*dT_max:
                        color1 = [253/253,34/253,0]
                        color2 = [253/253,34/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.86*dT_max> T >= 0.79*dT_max:
                        color1 = [253/253,68/253,0]
                        color2 = [253/253,68/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.79*dT_max> T >= 0.72*dT_max:
                        color1 = [253/253,102/253,0]
                        color2 = [253/253,102/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.72*dT_max> T >= 0.65*dT_max:
                        color1 = [253/253,136/253,0]
                        color2 = [253/253,136/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.65*dT_max> T >= 0.58*dT_max:
                        color1 = [253/253,170/253,0]
                        color2 = [253/253,170/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.58*dT_max> T >= 0.51*dT_max:
                        color1 = [253/253,204/253,0]
                        color2 = [253/253,204/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.51*dT_max> T >= 0.44*dT_max:
                        color1 = [253/253,238/253,0]
                        color2 = [253/253,238/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.44*dT_max> T >= 0.37*dT_max:
                        color1 = [233/253,253/253,0]
                        color2 = [233/253,253/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.37*dT_max> T >= 0.30*dT_max:
                        color1 = [199/253,253/253,0]
                        color2 = [199/253,253/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.30*dT_max> T >= 0.23*dT_max:
                        color1 = [165/253,253/253,0]
                        color2 = [165/253,253/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.23*dT_max> T >= 0.16*dT_max:
                        color1 = [131/253,253/253,0]
                        color2 = [131/253,253/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.16*dT_max> T >= 0.09*dT_max:
                        color1 = [97/253,253/253,0]
                        color2 = [97/253,253/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.09*dT_max> T >= 0.02*dT_max:
                        color1 = [63/253,253/253,0]
                        color2 = [63/253,253/253,0]
                        color.append(color1)
                        color.append(color2)

                    elif 0.02*dT_max> T >= 0.00*dT_max:
                        color1 = [29/253,253/253,0]
                        color2 = [29/253,253/253,0]
                        color.append(color1)
                        color.append(color2)

                else:
                    if dT_min <= T < 0.93*dT_min:               #가장 차가운
                        color1 = [0,25/253,253/253]
                        color2 = [0,25/253,253/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.93*dT_min <= T < 0.86*dT_min:
                        color1 = [0,59/253,253/253]
                        color2 = [0,59/253,253/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.86*dT_min <= T < 0.79*dT_min:
                        color1 = [0,93/253,253/253]
                        color2 = [0,93/253,253/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.79*dT_min <= T < 0.72*dT_min:
                        color1 = [0,127/253,253/253]
                        color2 = [0,127/253,253/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.72*dT_min <= T < 0.65*dT_min:
                        color1 = [0,161/253,253/253]
                        color2 = [0,161/253,253/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.65*dT_min <= T < 0.58*dT_min:
                        color1 = [0,195/253,253/253]
                        color2 = [0,195/253,253/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.58*dT_min <= T < 0.51*dT_min:
                        color1 = [0,229/253,253/253]
                        color2 = [0,229/253,253/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.51*dT_min <= T < 0.44*dT_min:
                        color1 = [0,253/253,243/253]
                        color2 = [0,253/253,243/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.44*dT_min <= T < 0.37*dT_min:
                        color1 = [0,253/253,209/253]
                        color2 = [0,253/253,209/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.37*dT_min <= T < 0.30*dT_min:
                        color1 = [0,253/253,175/253]
                        color2 = [0,253/253,175/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.30*dT_min <= T < 0.23*dT_min:
                        color1 = [0,253/253,141/253]
                        color2 = [0,253/253,141/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.23*dT_min <= T < 0.16*dT_min:
                        color1 = [0,253/253,107/253]
                        color2 = [0,253/253,107/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.16*dT_min <= T < 0.09*dT_min:
                        color1 = [0,253/253,73/253]
                        color2 = [0,253/253,73/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.09*dT_min <= T < 0.02*dT_min:
                        color1 = [0,253/253,39/253]
                        color2 = [0,253/253,39/253]
                        color.append(color1)
                        color.append(color2)

                    elif 0.02*dT_min <= T < 0.00*dT_min:
                        color1 = [0,253/253,5/253]
                        color2 = [0,253/253,5/253]
                        color.append(color1)
                        color.append(color2)

            return color

        glBegin(GL_TRIANGLE_STRIP)                              #(GL_TRIANGLE_STRIP으로 그릴것

        for a,b in zip(Vertices(),Color()):

            glColor3fv(b)                                       #색깔이 먼저 언급되어야한다.
            glVertex2fv(a)                                      #그다음이 좌표임

        glEnd()

        glBegin(GL_TRIANGLE_STRIP)                              #핀그림

        glColor3fv((253/253,0,0))
        glVertex2fv((-2 ,- 2))
        glColor3fv((253/253,0,0))
        glVertex2fv((-1.5,-2))
        glColor3fv((253/253,0,0))
        glVertex2fv((-2 ,  2))
        glColor3fv((253/253,0,0))
        glVertex2fv((-1.5, 2))

        glEnd()

def Fin_show():
    app = QApplication(sys.argv)
    myWindow = FinR()
    myWindow.show()
    sys.exit(app.exec_())


import numpy as np
from IPython.display import display, HTML

# Association of the values used in the data storage matrix
YES = 1
NO = 2
MAYBE = 0

class ZebraPuzzle:

    def __init__(self, categories_attributes: dict):
        self.categories_attributes: dict = categories_attributes
        # Prepare some convenience variables
        self.categories: list = list(self.categories_attributes.keys())
        self.number_categories: int = len(self.categories)
        self.number_attributes: int = len(self.categories_attributes[self.categories[0]])

        # Prepare some more convenience variables
        self.attributes = []
        for i, k in enumerate(self.categories):
            for j, n in enumerate(self.categories_attributes[k]):
                self.attributes.append(n.lower())
        self._tot: int = int(self.number_categories * self.number_attributes)

        # Check if each category has same number of attributes
        for k in self.categories_attributes:
            assert len(self.categories_attributes[k]) == self.number_attributes

        # Prepare the data-container
        self._data = np.zeros([self._tot, self._tot], dtype=int)
        # Initialize the trivial data
        assert YES == 1 and NO == 2, "If you want to change the values of YES and NO, please also change the code below accordingly."
        su = 2 * np.ones([self.number_attributes, ] * 2) - np.identity(self.number_attributes)
        for i in range(self.number_categories):
            self._data[i * self.number_attributes:(i + 1) * self.number_attributes,
            i * self.number_attributes:(i + 1) * self.number_attributes] = su

    def get_category(self, x: int) -> int:
        return int(x / self.number_attributes)

    def get_category_from_name(self, name: str) -> int:
        return self.get_category(self.get_id(name))

    def get_id(self, name: str) -> int:
        return self.attributes.index(name.lower())

    def get_assoc_name(self, a: str, b: str) -> int:
        return self.get_assoc(self.get_id(a), self.get_id(b))

    def get_assoc(self, a: int, b: int) -> int:
        assert a != b, "Requesting the association of a attribute with itself does not make sense."
        # assert that a and b are not from the same category
        assert self._data[a, b] == self._data[b, a]
        return self._data[min(a, b), max(a, b)]

    def add_assoc_name(self, name_a: str, name_b: str, assoc: bool):
        k_a = self.get_id(name_a)
        k_b = self.get_id(name_b)
        self.add_assoc(k_a, k_b, assoc)

    def add_assoc(self, a: int, b: int, assoc: bool):
        v = YES if assoc else NO
        assert self.get_assoc(a, b) == MAYBE or self.get_assoc(a, b) == v, \
            f"Try to overwrite field {a, b} with value {self.get_assoc(a, b)}."

        # In case of a positive assoc, we can already remove quite a few possibilities.
        if assoc:
            for i in range(self.number_attributes):
                self._data[self.get_category(a) * self.number_attributes + i, b] = NO
                self._data[a, self.get_category(b) * self.number_attributes + i] = NO
                self._data[b, self.get_category(a) * self.number_attributes + i] = NO
                self._data[self.get_category(b) * self.number_attributes + i, a] = NO

        self._data[a, b] = v
        self._data[b, a] = v

    def generate_matrix_html(self) -> str:
        from IPython.display import display, HTML
        # Some css to make the table look nice
        table = """

    """
        # Now generate the required HTML
        table += '
'
        # Print categories
        table += '
		'
        for ki in range(1, self.number_categories):
            k = self.categories[ki]
            table += f"{k}		"
        table += "

		"
        # Print attributes
        for ki in range(1, self.number_categories):
            k = self.categories[ki]
            for att in self.categories_attributes[k]:
                table += f"{att}	"
            table += "	"
        table += "

"
        # Print main body
        for ki in range(0, self.number_categories - 1):
            # Print categories on left
            k = self.categories[ki]
            table += f"{k}	"
            for i, att in enumerate(self.categories_attributes[k]):
                # Print attributes on left
                it = ki * self.number_attributes + i
                table += f"{att}	"
                for kj in range(1, self.number_categories):
                    for j, att2 in enumerate(self.categories_attributes[k]):
                        # Print values
                        jt = kj * self.number_attributes + j
                        if kj <= ki:
                            c = "-"
                        elif self.get_assoc(jt, it) == YES:
                            c = 'o'
                        elif self.get_assoc(jt, it) == NO:
                            c = 'x'
                        else:
                            c = '.'
                        table += f"{c}	"
                    table += f"|	"
                table += "

"
            table += "-	" * (2 + (self.number_categories - 1) * (self.number_attributes + 1))
            table += "

"
        table += '

'
        table += f"Filled about {int(np.round(self.get_percentage_solved())):d}% of the puzzle."
        return table

    def display_matrix(self):
        # Display output
        html_table = HTML(self.generate_matrix_html())
        display(html_table)

    def _exclude(self, x: int, y: int):
        assert self.get_assoc(x, y) == YES
        kx = self.get_category(x)
        ky = self.get_category(y)
        for xi in range(0, self._tot):
            if xi != y and self.get_assoc(xi, y) == NO:
                if self.get_category(xi) != kx and self.get_category(xi) != ky:
                    if self.get_assoc(x, xi) == MAYBE:
                        print(f"By simple logic '{self.attributes[x]}' and '{self.attributes[xi]}'" +
                              f" are negatively associated.")
                        self.add_assoc(x, xi, False)

    def exclude_all(self):
        for i in range(self._tot):
            for j in range(i + 1, self._tot):
                if self.get_assoc(i, j) == YES:
                    self._exclude(i, j)
                    self._exclude(j, i)

    def find_solved(self):
        for ki in range(self.number_categories):
            for kj in range(ki + 1, self.number_categories):
                sr, er = ki * self.number_attributes, (ki + 1) * self.number_attributes
                for x in range(sr, er):
                    sc, ec = kj * self.number_attributes, (kj + 1) * self.number_attributes
                    d = [r == NO for r in self._data[x, sc:ec]]
                    # print(x, f"{sc}:{ec}", d)
                    if sum(d) == self.number_attributes - 1:
                        for y in range(sc, ec):
                            if self.get_assoc(x, y) == 0:
                                print(f"By simple logic '{self.attributes[x]}' and '{self.attributes[y]}'" +
                                      f" are positively associated.")
                                # print(self._data[x, sc:ec])
                                self.add_assoc(x, y, True)
                # Other way around.
                # TODO: this is ugly code-duplication. Re-implement this at your leisure.
                sc, ec = kj * self.number_attributes, (kj + 1) * self.number_attributes
                for y in range(sc, ec):
                    d = [r == NO for r in self._data[sr:er, y]]
                    # print(x, f"{sc}:{ec}", d)
                    if sum(d) == self.number_attributes - 1:
                        for x in range(sr, er):
                            if self._data[x, y] == 0:
                                print(f"By simple logic '{self.attributes[x]}' and '{self.attributes[y]}'" +
                                      f" are positively associated.")
                                # print(self._data[sc:ec, y])
                                self.add_assoc(x, y, True)

    def iterate_logic(self):
        p = 0
        while self.get_percentage_solved() > p:
            p = self.get_percentage_solved()
            self.exclude_all()
            self.find_solved()

    def get_percentage_solved(self):
        # Estimates the percentage to which this puzzle is solved.
        n = self._tot
        m = self.number_attributes
        n_yes = self.number_categories ** 2 * m
        init = self.number_categories * (m ** 2 - m) * NO + self.number_categories * m * YES

        return 100 * (np.sum(self._data) - init) / (YES * n_yes + NO * (n ** 2 - n_yes) - init)


from PyPDF2 import PdfFileWriter, PdfFileReader
from PyPDF2.generic import NameObject, createStringObject

def read(input_pdf):

    pdf_reader = PdfFileReader(input_pdf)
    print(pdf_reader.getDocumentInfo())

if __name__ == '__main__':
    # read(input_pdf='ROSCON2019_Workshop_Aliasroboticscopy.pdf')
    read(input_pdf='AliasRobotics - RT_workshop_ROSCon_2019.pdf')


from typing import Optional, Union
from transformers import PreTrainedModel
from data.rl_data import RL_Dataset
from models.base import InputType
from models.iql_model import PerTokenIQL
import torch
import wandb
from utils.torch_utils import get_transformer_logs, map_pytree

class CQLModel(PerTokenIQL):
    def __init__(self,
                 model: PreTrainedModel,
                 dataset: RL_Dataset,
                 device: Union[torch.device, str] = "cuda",
                 alpha: float = 0.005,
                 gamma=1.0,
                 beta=1.0,
                 transition_weight=0.0,
                 clip_weight: Optional[float] = None,
                 value_max: Optional[float] = None,
                 value_min: Optional[float] = None,
                 detach_v: bool = False,
                 detach_pi: bool = False,
                 detach_q: bool = False,
                 double_q: bool = False,
                 seperate_policy: bool = False,
                 seperate_target: bool = False,
                 exp_weights: bool = False,
                 advanced_mlp: bool = False,
                 cql_temp: float = 1.0,
                ):
        super().__init__(model, dataset, device, alpha, gamma, beta, transition_weight,
                         clip_weight, value_max, value_min, detach_v=detach_v, detach_pi=detach_pi,
                         detach_q=detach_q, double_q=double_q, tau=0.5, seperate_policy=seperate_policy,
                         seperate_target=seperate_target, exp_weights=exp_weights,
                         dm_margin=0.0, advanced_mlp=advanced_mlp, cql_temp=cql_temp)

    def get_qvs(self, items: InputType,
                prefix_embs: Optional[torch.Tensor]=None,
                prefix_attn_mask: Optional[torch.Tensor]=None,
                remove_prefix_position_embs: bool=False,
                qv_kwargs=None, policy_kwargs=None, target_kwargs=None,
                **kwargs):
        prepared_inputs = self.prepare_inputs(items)
        tokens, attn_mask = prepared_inputs['tokens'], prepared_inputs['attn_mask']
        s_idx, a_idx = prepared_inputs['state_idxs'], prepared_inputs['action_idxs']
        rs, terminals = prepared_inputs['rewards'], prepared_inputs['terminals']
        self_outputs = self(tokens, attn_mask, s_idx, a_idx,
                            prefix_embs, prefix_attn_mask,
                            remove_prefix_position_embs,
                            qv_kwargs, policy_kwargs, target_kwargs,
                            **kwargs)
        model_outputs, vs, qs = self_outputs['model_outputs'], self_outputs['vs'], self_outputs['qs']
        target_qs, logits = self_outputs['target_qs'], self_outputs['logits']
        vt = vs[:, :-1]
        vtp1 = vs[:, 1:]
        select_tokens = torch.gather(tokens[:, 1:], dim=1, index=a_idx)
        cql_term = self.get_cql_loss(qs, select_tokens, terminals)
        full_qs = qs
        if self.double_q:
            q1, q2 = qs
            q1 = torch.gather(q1, dim=2, index=select_tokens.unsqueeze(2)).squeeze(2)
            q2 = torch.gather(q2, dim=2, index=select_tokens.unsqueeze(2)).squeeze(2)
            qs = (q1, q2,)
        else:
            qs = torch.gather(qs, dim=2, index=select_tokens.unsqueeze(2)).squeeze(2)
        dm_term = self.get_dm_loss(full_qs, qs, terminals, self.dm_margin)
        all_target_qns = torch.cat((target_qs[:, 1:, :], torch.zeros((target_qs.shape[0], 1, target_qs.shape[2],)).to(self.device)), dim=1)
        target_qs = torch.gather(target_qs, dim=2, index=select_tokens.unsqueeze(2)).squeeze(2)
        with torch.no_grad():
            weights = self.get_weights(tokens, vt, target_qs, s_idx, a_idx, terminals)
        return {
                    'tokens': tokens,
                    'attn_mask': attn_mask,
                    'model_outputs': model_outputs,
                    'vs': vt,
                    'qs': qs,
                    'vns': vtp1,
                    'target_vs': vt,
                    'target_qs': target_qs,
                    'target_vns': vtp1,
                    'rs': rs,
                    'terminals': terminals,
                    'logits': logits,
                    'weights': weights,
                    'cql_term': cql_term,
                    'dm_term': dm_term,
                    'all_target_qns': all_target_qns,
                }

    def get_q_loss(self, all_target_qns, qs, rs, gamma, terminals):
        all_target_qns = all_target_qns.detach()
        max_qns = torch.max(all_target_qns, dim=-1).values
        if self.double_q:
            q1, q2 = qs
            l1 = ((((1 - terminals[:, 1:]) * max_qns * gamma + rs - q1) ** 2) * (1 - terminals[:, :-1])).sum() / max((1 - terminals[:, :-1]).sum().item(), 1.0)
            l2 = ((((1 - terminals[:, 1:]) * max_qns * gamma + rs - q2) ** 2) * (1 - terminals[:, :-1])).sum() / max((1 - terminals[:, :-1]).sum().item(), 1.0)
            return l1 + l2
        return ((((1 - terminals[:, 1:]) * max_qns * gamma + rs - qs) ** 2) * (1 - terminals[:, :-1])).sum() / max((1 - terminals[:, :-1]).sum().item(), 1.0)

    def get_loss(self,
                 items: InputType,
                 awac_weight=0.0,
                 v_loss_weight=0.0,
                 q_loss_weight=0.0,
                 cql_loss_weight=0.0,
                 dm_loss_weight=0.0,
                 mc_returns=False):
        prepared_inputs = self.prepare_inputs(items)
        a_idx = prepared_inputs['action_idxs']
        get_qvs_outputs = self.get_qvs(items,
                                       qv_kwargs={'output_attentions': True},
                                       policy_kwargs={'output_attentions': True},
                                       target_kwargs={'output_attentions': True},
                                       skip_policy_on_train=(awac_weight == 0.0),
                                      )
        tokens, attn_mask, model_outputs = get_qvs_outputs['tokens'], get_qvs_outputs['attn_mask'], get_qvs_outputs['model_outputs']
        vs, qs = get_qvs_outputs['vs'], get_qvs_outputs['qs']
        vns, target_qs, rs = get_qvs_outputs['vns'], get_qvs_outputs['target_qs'], get_qvs_outputs['rs']
        terminals, logits, weights = get_qvs_outputs['terminals'], get_qvs_outputs['logits'], get_qvs_outputs['weights']
        all_target_qns = get_qvs_outputs['all_target_qns']

        logs = {}
        transformer_logs = {}
        transformer_logs['qv_transformer_logs'] = get_transformer_logs(model_outputs['qv_model_outputs'].attentions, self.model, attn_mask)
        if self.lm_policy is not None and (not (self.training and awac_weight == 0.0)):
            transformer_logs['policy_transformer_logs'] = get_transformer_logs(model_outputs['policy_model_outputs'].attentions, self.lm_policy, attn_mask)
        if self.lm_target is not None:
            transformer_logs['target_transformer_logs'] = get_transformer_logs(model_outputs['target_model_outputs'].attentions, self.lm_target, attn_mask)
        n = max((1 - terminals[:, :-1]).sum().item(), 1)
        rs_downstream = self.get_downstream_rs(rs, self.gamma)
        if mc_returns:
            v_loss = self.get_v_loss(vs, rs_downstream, terminals)
        else:
            v_loss = self.get_v_loss(vs, target_qs, terminals)
        q_loss = self.get_q_loss(all_target_qns, qs, rs, self.gamma, terminals)
        cql_loss = get_qvs_outputs['cql_term']
        dm_loss = get_qvs_outputs['dm_term']
        token_loss = self.awac_loss(tokens, attn_mask, logits, weights)
        logs['token_loss'] = (token_loss.item(), n)
        loss = awac_weight * token_loss + v_loss_weight * v_loss + q_loss_weight * q_loss + cql_loss_weight * cql_loss + dm_loss_weight * dm_loss
        logs['v_loss'] = (v_loss.item(), n)
        logs['q_loss'] = (q_loss.item(), n)
        logs['cql_loss'] = (cql_loss.item(), n)
        logs['dm_loss'] = (dm_loss.item(), n)
        advantages = sum([((target_qs[i] - vs[i])[:(1 - terminals[i, :-1]).sum().long().item()]).detach().cpu().tolist() for i in range(tokens.shape[0])], [])
        if self.double_q:
            q1, q2 = qs
            logs['q1_avg'] = ((q1 * (1 - terminals[:, :-1])).sum().item() / max(n, 1), n)
            logs['q1_var'] = (((((q1 - logs['q1_avg'][0]) ** 2)*(1 - terminals[:, :-1])).sum() / max(n, 1)).item(), 1)
            logs['q2_avg'] = ((q2 * (1 - terminals[:, :-1])).sum().item() / max(n, 1), n)
            logs['q2_var'] = (((((q2 - logs['q2_avg'][0]) ** 2)*(1 - terminals[:, :-1])).sum() / max(n, 1)).item(), 1)
        else:
            logs['q_avg'] = ((qs * (1 - terminals[:, :-1])).sum().item() / max(n, 1), n)
            logs['q_var'] = (((((qs - logs['q_avg'][0]) ** 2)*(1 - terminals[:, :-1])).sum() / max(n, 1)).item(), 1)
        logs['v_avg'] = ((vs * (1 - terminals[:, :-1])).sum().item() / max(n, 1), n)
        logs['v_var'] = (((((vs - logs['v_avg'][0]) ** 2)*(1 - terminals[:, :-1])).sum() / max(n, 1)).item(), 1)
        act_weights = torch.gather(weights, dim=1, index=a_idx)
        logs['act_weight_avg'] = (((act_weights * (1 - terminals[:, :-1])).sum() / max(n, 1)).item(), n)
        logs['transformer'] = transformer_logs
        postproc_f = lambda l: l.update({'loss': awac_weight * l['token_loss'] + q_loss_weight * l['q_loss'] + v_loss_weight * l['v_loss'] + cql_loss_weight * l['cql_loss'] + dm_loss_weight * l['dm_loss']})
        hist_f = lambda l: l.update({'advantage_hist': wandb.Histogram(advantages)})
        return loss, logs, [postproc_f, hist_f]


#-*- coding:utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F

from torch.utils import model_zoo
from torchvision import models

# 最复杂
class FCN8(nn.Module):
    def __init__(self, in_channels=6, n_classes=2,pretrained=True):
        super(FCN8, self).__init__()
        feats = list(models.vgg16(pretrained=pretrained).features.children())

        self.first_layer = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1) #改造VGG 第一层channel为6
        self.feats = nn.Sequential(*feats[1:9])
        self.feat3 = nn.Sequential(*feats[10:16])
        self.feat4 = nn.Sequential(*feats[17:23])
        self.feat5 = nn.Sequential(*feats[24:30])

        # first_conv2d_flag = 0
        # for m in self.modules():
        #     if isinstance(m, nn.Conv2d):
        #         if first_conv2d_flag==0:
        #             first_conv2d_flag = 1
        #             continue
        #         m.requires_grad = False

        self.fconn = nn.Sequential(
            nn.Conv2d(512, 4096, 7),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Conv2d(4096, 4096, 1),
            nn.ReLU(inplace=True),
            nn.Dropout(),
        )
        self.score_feat3 = nn.Conv2d(256, n_classes, 1)
        self.score_feat4 = nn.Conv2d(512, n_classes, 1)
        self.score_fconn = nn.Conv2d(4096, n_classes, 1)

    def forward(self, x):
        feats = self.first_layer(x)
        print(feats.size())
        feats = self.feats(feats)
        feat3 = self.feat3(feats) # size=(N, 512, x.H/8, x.W/8)
        feat4 = self.feat4(feat3) # size=(N, 512, x.H/16, x.W/16)
        feat5 = self.feat5(feat4) # size=(N, 512, x.H/32, x.W/32)
        print(feat5.size())
        fconn = self.fconn(feat5)

        score_feat3 = self.score_feat3(feat3)
        score_feat4 = self.score_feat4(feat4)
        score_fconn = self.score_fconn(fconn)

        print(score_fconn.size(),score_feat4.size())
        score = F.upsample(score_fconn, score_feat4.size()[2:])

        score += score_feat4
        score = F.upsample(score, score_feat3.size()[2:])
        score += score_feat3
        out = F.upsample(score, x.size()[2:])
        # print( out.size() )
        return out

class FCN16(nn.Module):
    def __init__(self, in_channels=6, n_classes=2, pretrained=True):
        super(FCN16, self).__init__()
        feats = list(models.vgg16(pretrained=pretrained).features.children())

        self.feats = nn.Sequential(*feats[0:16])
        self.feat4 = nn.Sequential(*feats[17:23])
        self.feat5 = nn.Sequential(*feats[24:30])
        self.fconn = nn.Sequential(
            nn.Conv2d(512, 4096, 7),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Conv2d(4096, 4096, 1),
            nn.ReLU(inplace=True),
            nn.Dropout(),
        )
        self.score_fconn = nn.Conv2d(4096, n_classes, 1)
        self.score_feat4 = nn.Conv2d(512, n_classes, 1)

    def forward(self, x):
        feats = self.feats(x)
        feat4 = self.feat4(feats)
        feat5 = self.feat5(feat4)
        fconn = self.fconn(feat5)

        score_feat4 = self.score_feat4(feat4)
        score_fconn = self.score_fconn(fconn)

        score = F.upsample(score_fconn, score_feat4.size()[2:])
        score += score_feat4

        return F.upsample(score, x.size()[2:])

class FCN32(nn.Module):
    def __init__(self, in_channels=6, n_classes=2, pretrained=True):
        super(FCN32, self).__init__()
        self.feats = models.vgg16(pretrained=pretrained).features
        self.fconn = nn.Sequential(
            nn.Conv2d(512, 4096, 7),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Conv2d(4096, 4096, 1),
            nn.ReLU(inplace=True),
            nn.Dropout(),
        )
        self.score = nn.Conv2d(4096, n_classes, 1)

    def forward(self, x):
        feats = self.feats(x)
        fconn = self.fconn(feats)
        score = self.score(fconn)

        return F.upsample(score, x.size()[2:])

if __name__ == '__main__':
    model = models.vgg16(pretrained=False)
    print(model)


from functools import partial

import torch.optim as optim
from ppo_pytorch.common import DecayLR
from ppo_pytorch.actors import MLPActor, ActionValuesHead
from torch import nn as nn

def create_small_mlp_kwargs(dueling=True, activation=nn.Tanh, lr_decay_frames=1e5, **kwargs):
    defaults = dict(
        optim=partial(optim.Adam, lr=3e-4),
        eval_batch_size=1,
        replay_batch_size=32,
        trainable_batch_size=64,
        replay_size=50_000,
        target_net_update_freq=1000,
        eps=1,
        train_interval=4,
        reward_discount=0.995,
        model_factory=partial(MLPActor, activation=activation, head_factory=partial(ActionValuesHead, dueling=dueling)),
        cuda_eval=False,
        cuda_replay=False,
        double=True,
        learning_starts=2048,
        reward_scale=0.2,
        lr_scheduler=partial(DecayLR, start_value=1, end_value=0.05, end_epoch=lr_decay_frames),
    )
    defaults.update(kwargs)
    return defaults